Repository navigation
Conversation
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unsupported custom symmetry subclasses are silently ignored instead of being rejected or enforced.
1 open finding
What changed in this PR
Adds exact symmetry-aware Gaussian process kernels while preserving recommender-level data augmentation.
Changes:
- Implements permutation-, mirror-, and dependency-invariant kernel wrappers.
- Adds GP symmetry validation, parameter tying, RGPE propagation, and tests.
- Updates documentation, examples, constraints, and search-space lookup behavior.
| File | Description |
|---|---|
baybe/surrogates/gaussian_process/_symmetry.py |
Constructs and validates invariant kernels. |
baybe/surrogates/gaussian_process/components/_gpytorch.py |
Adds GPyTorch symmetry wrappers and parameter tying. |
baybe/surrogates/gaussian_process/core.py |
Exposes GP symmetries and integrates kernel construction. |
baybe/surrogates/gaussian_process/_override/kernel.py |
Adds kernel-tree traversal. |
baybe/surrogates/gaussian_process/_override/__init__.py |
Exports kernel-tree traversal. |
baybe/symmetries/base.py |
Documents invariant-kernel usage. |
baybe/symmetries/dependency.py |
Validates dependency roles and clarifies discretization. |
baybe/constraints/discrete.py |
Rejects self-dependent parameters. |
baybe/searchspace/core.py |
Generalizes name collections. |
baybe/searchspace/discrete.py |
Rejects bare-string name queries. |
baybe/searchspace/continuous.py |
Rejects bare-string name queries. |
baybe/recommenders/pure/bayesian/base.py |
Distinguishes augmentation from kernel symmetry. |
tests/test_invariant_kernels.py |
Tests invariance, validity, and posterior behavior. |
tests/validation/test_symmetry_validation.py |
Tests invalid GP symmetry configurations. |
tests/validation/test_constraint_validation.py |
Tests self-dependency rejection. |
tests/test_searchspace.py |
Tests exact parameter-name lookup. |
tests/test_rgpe_surrogate.py |
Tests symmetry propagation into RGPE models. |
tests/test_parameter_kernel_overrides.py |
Tests symmetry composition with overrides. |
tests/test_iterations.py |
Adds end-to-end symmetric GP coverage. |
tests/test_gp.py |
Tests preset symmetry configuration. |
tests/hypothesis_strategies/surrogates.py |
Generates symmetry-enabled GP configurations. |
examples/Symmetries/permutation.py |
Compares constraints, augmentation, and invariant kernels. |
docs/components/symmetries.md |
Documents invariant-kernel mathematics and limitations. |
docs/components/surrogates.md |
Links surrogate users to symmetry documentation. |
docs/components/recommenders.md |
Documents the kernel-based alternative. |
docs/components/parameters.md |
Documents override compatibility. |
docs/components/constraints.md |
Documents self-dependency restrictions. |
CHANGELOG.md |
Records API and feature changes. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+68
to
+70
| permutations = [s for s in symmetries if isinstance(s, PermutationSymmetry)] | ||
| mirrors = [s for s in symmetries if isinstance(s, MirrorSymmetry)] | ||
| dependencies = [s for s in symmetries if isinstance(s, DependencySymmetry)] |
A single string passed as names was matched via substring membership, so looking up "n1_not_discrete" also returned a parameter named "n1". DependencySymmetry passed its causing parameter name this way and could thus validate the wrong parameter. Strings are now rejected with a ValueError and the affected call site passes a tuple.
Adds iter_gpytorch_kernel_tree, which yields every kernel of a GPyTorch kernel tree with the input columns it acts on, and rebuilds get_active_dimensions on top of it. The results are unchanged for all kernels built by BayBE, but nested active_dims and arbitrary composite kernels are now handled correctly.
Replaces the blanket ban on symmetries sharing parameters with rules based on the roles of the parameters. In particular, dependencies are now allowed on permuted parameters as long as they are closed under the permutations, which enables slot-based mixtures.
Scienfitz
force-pushed
the
feature/symmetric_kernels
branch
from
October 7, 2026 18:37
a117378 to
545f3e7
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

Closes #622 and #620
This PR adds a
symmetriesattribute to GPs. In contrast toRecommender.symmetries(which triggers data augmentation) this triggers the construction of symmetric kernels in theGP if possible.This is a complex topic and I don't consider this PR wisdoms last word, but tis a start and somewhat complete in terms of coverage.
TODO, Pending Until Architecture Agreement
Proof of Concept
Here is the result from the expanded Symmetry example. It shows
-------------------- LLM based below --------------------
Mathematical Treatment of the Kernel Symmetries
Let$k$ denote the fully resolved kernel of the GP, i.e. the surrogate kernel including all parameter kernel overrides and the task kernel. All symmetries act on the normalized model inputs $x, x'$ , i.e. after the GP's $k$ .
Normalizeinput transform. None of the constructions below add hyperparameters, all learnable parameters remain those ofA symmetrized kernel$\tilde k$ makes the model exactly invariant under a transformation $T$ if $\tilde k(Tx, x') = \tilde k(x, x')$ for all $x, x'$ . The posterior mean $\mu(x) = m(x) + \tilde k(x, X) \alpha$ and the posterior variance only depend on the inputs via $\tilde k$ and the prior mean $m$ . The posterior is therefore invariant for any hyperparameters, provided that $m$ is constant (enforced).
Permutation Symmetry
Equation (tied single sum). Let$\pi$ range over all $n!$ permutations of the $n$ positions of a $\pi x'$ permutes the positions of $x'$ in all groups in lockstep; parameters with several computational columns (e.g. one-hot encodings) are moved as blocks:
PermutationSymmetry.Validity. This is a valid kernel and invariant under all permutations if the base kernel is invariant under jointly permuting both inputs:
In that case,$k_{\text{perm}}(x, x') = \tfrac{1}{n!} \sum_{\pi, \sigma} k(\pi x, \sigma x')$ . BayBE guarantees the condition by sharing the hyperparameters of permuted positions:
raw_lengthscale,raw_period_length) are tied entry-wise via atorch.nn.utils.parametrizeparametrization (TiedEntries). The positions thus share a single learned lengthscale, e.g.override_kernels, orMatern(["x1"]) * Matern(["x2"])): the kernel of each position reuses the parameter objects of the corresponding first-position kernel.Before fitting, a numerical check evaluates a copy of the tied kernel with random hyperparameters on random inputs and verifies$k(\pi x, \pi x') = k(x, x')$ .
Limitations and whether they occur in BayBE.
Matern(["x1"]) * RBF(["x2"])) or with a factoryparameter_selectorexcluding only some permuted parameters. Both raise aValueErrorbefore fitting. Different kernel overrides on permuted parameters are already rejected, sincePermutationSymmetryrequires equivalent parameters._constraintsnaming convention.Cost.$n!$ evaluations of the inner kernel per permutation symmetry with $n$
positions.
Mirror Symmetry
Equation (double sum). Let$m$ reflect the mirrored column at the normalized mirror point $c' = (c - \ell) / (u - \ell)$ , i.e. replace its value $v$ by $2c' - v$ , where $[\ell, u]$ are the scaling bounds of the parameter:
Validity. This is a valid kernel and invariant under the reflection for any base kernel, since it equals the inner product of the feature maps$\phi(x) + \phi(m x)$ . No hyperparameter tying is needed. For reflection-invariant base kernels (all kernels depending only on per-dimension distances, e.g. Matérn, RBF, RQ, periodic), the double sum equals twice the single sum $k(x, x') + k(x, m x')$ . The learned output scale absorbs that factor.
Limitations and whether they occur in BayBE.
MirrorSymmetryitself.Cost. 4 evaluations of the inner kernel per mirror symmetry.
Dependency Symmetry
Equation (indicator gating). For$D$ dependency symmetries, let $a_d(x) \in {0, 1}$ indicate whether dependency $d$ is active at $x$ , and let $P(x) = (a_1(x), \dots, a_D(x))$ be the activity pattern of $x$ . For a pattern $P$ , let $x_P$ denote $x$ with the affected columns of all dependencies inactive in $P$ set to a constant (the lower bound $0$ of the normalized range):
Validity. Each term is the product of the rank-one mask$\mathbf{1}[P(x) = P],\mathbf{1}[P(x') = P]$ and a kernel evaluated on transformed inputs, so the sum is positive semi-definite.
off1,off2)Limitations and whether they occur in BayBE.
DependencySymmetryitself.SubstanceParameterandCustomDiscreteParameterwith decorrelation, which can map an active and an inactive value to identical encodings. RaisesIncompatibleSearchSpaceErrorbefore fitting.ThresholdConditionon a categorical parameter. Raises an error before fitting viaCondition.evaluate.Cost.$2^D$ evaluations of the inner kernel for $D$ dependency symmetries (one per activity pattern, each masked to the matching pairs).
Overall Cost
With$G$ permutation symmetries with $n_1, \dots, n_G$ positions, $M$ mirror symmetries and $D$ dependency symmetries, each evaluation of the symmetric kernel requires
evaluations of the resolved kernel$k$ . This factor applies to building the kernel matrices, in every step of the hyperparameter optimization and at prediction time. The cost of the subsequent linear algebra is unchanged, since the number of training points stays the same. Data augmentation, in contrast, multiplies the number of training points.
Combining Symmetries
Stack and Order
The symmetric kernel is built at the very end of$k$ is then wrapped, from the inside out:
_resolve_kernel, i.e. after all parameter kernel overrides and the task kernel have been applied. The resolved kernelBefore wrapping, the hyperparameters of all permuted positions in$k$ are tied, and the joint permutation invariance of $k$ is verified numerically.
columns to determine the activity patterns and to fix inactive affected columns.
permuting both inputs. The gated kernel satisfies this if
dependencies are closed under the permutation (see below). The permutation then only
renumbers the dependencies, i.e.
would determine the activity patterns on unpermuted inputs and break the permutation
invariance.
are neither permuted by another symmetry nor causing parameters. The reflections
therefore commute with the permutations and never change an activity pattern, so all
inner invariances are preserved.
Because the posterior also depends on the prior mean, only constant means (or task-wise constant means in transfer learning) are accepted. In RGPE mode, each inner GP of the ensemble builds the same stack on its own search space.
Allowed Combinations (LLM based)
Each wrapper always yields a valid kernel, but the invariances only combine if the roles of the involved parameters are compatible. A parameter can be permuted, mirrored, causing (decides whether a dependency is active) or affected (only matters if a dependency is active).
validate_symmetriesenforces the following rules at construction time:Closure means that renaming the causing and affected parameters of any dependency according to any permutation yields another existing dependency with an identical condition (compared via
==). Dependencies that touch no permuted parameter satisfy this automatically. Permutation groups are additionally limited to 5 positions.Slot-based mixtures. The closure rule enables the main use case of combining symmetries. A mixture modeled with slots permutes the substance labels and amounts of all slots in lockstep, and each label only matters if the amount of its slot is positive:
PermutationSymmetry([["Solvent_1", "Solvent_2", "Solvent_3"], ["Fraction_1", "Fraction_2", "Fraction_3"]])DependencySymmetry("Fraction_i", ThresholdCondition(0, ">"), ["Solvent_i"])forPermuting the slots maps the dependency of slot$i$ onto the dependency of slot $\pi(i)$ , so the set of dependencies is unchanged and all invariances combine. This is exactly what $2^3 \cdot 3! = 48$ kernel evaluations.
DiscretePermutationInvarianceConstraint.to_symmetry()andDiscreteDependenciesConstraint.to_symmetries()produce for such a setup, provided all slots use identical conditions. Mixtures whose slots use equivalent but differently specified conditions are rejected. A 3-slot mixture costs