To Augment or Not to Augment? Diagnosing Distributional Symmetry Breaking
Authors: Hannah Lawrence
Elyssa Hofgard
Vasco Portilheiro
Yuxuan Chen
Tess Smidt
Robin Walters
Venue: ICLR 2026 Poster (2026)
Keywords: equivariance
symmetry
symmetry breaking
canonicalization
data augmentation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present an interpretable metric for quantifying distributional symmetry-breaking in point cloud datasets, and show that benchmark point cloud datasets are highly anisotropic, with implications for when equivariant methods are advisable.

Abstract: Symmetry-aware methods for machine learning, such as data augmentation and equivariant architectures, encourage correct model behavior on all transformations (e.g. rotations or permutations) of the original dataset. These methods can improve generalization and sample efficiency, under the assumption that the transformed datapoints are highly probable, or "important", under the test distribution. In this work, we develop a method for critically evaluating this assumption. In particular, we propose a metric to quantify the amount of symmetry breaking in a dataset, via a two-sample classifier test that distinguishes between the original dataset and its randomly augmented equivalent. We validate our metric on synthetic datasets, and then use it to uncover surprisingly high degrees of symmetry-breaking in several benchmark point cloud datasets, constituting a severe form of dataset bias. We show theoretically that distributional symmetry-breaking can prevent invariant methods from performing optimally even when the underlying labels are truly invariant, for invariant ridge regression in the infinite feature limit. Empirically, the implication for symmetry-aware methods is dataset-dependent: equivariant methods still impart benefits on some symmetry-biased datasets, but not others, particularly when the symmetry bias is predictive of the labels. Overall, these findings suggest that understanding equivariance — both when it works, and why — may require rethinking symmetry biases in the data.

Published as a conference paper at ICLR 2026
To AUGMENT OR NOT TO AUGMENT? DIAGNOSING DISTRIBUTIONAL SYMMETRY BREAKING
Hannah Lawrence"! Elyssa Hofgard" Vasco Portilheiro?, Yuxuan Chen', Tess Smidt',
Robin Walters²
MIT, University College London, Northeastern University
Equua contribution.
ABSTRACT
Symmetry-aware methods for machine learning, such as data augmentation and equivariant architectures, encourage correct model behavior on all transforma- tions (e.g. rotations or permutations) of the original dataset. These methods can improve generalization and sample efficiency, under the assumption that the trans- formed datapoints are highly probable, or "important", under the test distribution. In this work, we develop a method for critically evaluating this assumption. particular, we propose metric to quantify the amount of symmetry breaking in a dataset, via a two-sample classifier test that distinguishes between the original dataset and its randomly augmented equivalent. We validate metric syn- thetic datasets, and then use it to uncover surprisingly high degrees of symmetry- breaking in several benchmark point cloud datasets, constituting severe form of dataset bias. We show theoretically that distributional symmetry-breaking can prevent invariant methods from performing optimally even when the underlying labels are truly invariant, for invariant ridge regression in the infinite feature limit. Empirically, the implication for symmetry-aware methods dataset-dependent: equivariant methods still impart benefits on some symmetry-bee biased datasets, but not others, particularly when the symmetry bias is predictive of the labels. Over- all, these findings suggest that understanding equivariance both when it works, and why - may require rethinking symmetry biases in the data.
INTRODUCTION
By integrating physical symmetries into model architectures as group invariances, equivariant neural networks often achieve superior performance across materials science (Liao al., (023). robotics (Wang et al. 2024a), drug discovery (Igashov et al.) 2024), fluid dynamics (Wang et al. (021), computer vision (Esteves et al. 2019), and beyond. Given a group or permutations, a function neural network f is equivariant if f(gx) (x) Vgg E C. Their success is typically explained in terms of improved sample efficiency and generalizability. resulting from the ability relate data sample and transformed data sample gx (Cohen & Welling 20166). Alternatively, data augmentation may be used to enforce equivariance by applying a random g to each input x in the training set and its corresponding label. For all of these equivariance methods, is an explicit assumption that the ground truth function ff is equivariant (Thomas et al.) [0018] Cohen & WWelling 2016a Cohen al., 2018). However, there is also often an implicit assumption that transformed samples gx occur relatively uniformly in distribution, i.e. p(x) ~ p(gx). Theoretical results on the benefits of equivariance and data augmentation almost always assume that x and gx
symmetries, like rotations
are equally likely under the data distribution (Flesedy & Zaidi 2021 Chen et al.
2020
Mei et al.,
2021 Lyle et al. 2020 Tahmasebi & Jegelka 2023).
In this paper, we study distributional symmetry breaking (Wang et al., 2024b d), or equivalently symmetry bias-when a datapoint x and its transform gx are not equally likely under the data distri- bution. Formally, this means that px (x) & px (gx) for some x E X and g G, where px : - R is the data distribution, or equivalently that all points the orbit {gx]geG of x are not equally likely. This paper takes step towards the goal of understanding how distributional symmetry breaking affects the performance of equivariant methods, including the ubiquitous practice of data augmentation. In- tuirively, although equivariance can help performance by providing the correct inductive bias all

---

Published as a conference paper at ICLR 2026
arbitrary orientations
"canonical"
96
orieniation
Equally Likely
Likely
Unlikely
(a)
(b)
(c)
Figure 1: (a) Distributional symmetry breaking: Baseballs are likely to occur in any orientation in photos, and are therefore uniform across orbits. In contrast, coffee mugs are more likely to appear with the handle on the side. The latter is an example of distributional symmetry breaking. (b) Canonicalization: Canonicalization is when object only ever appears in one, "canonical", orientation. This is the strongest form of distributional symmetry breaking. (c) Inherent vs. user- defined canonicalization:
Datapoints can be canonicalized for reasons that are inherent, such the orientation of digit determining whether it is a 6 or 9. However, it can also be user-defined, such as the orientation of crystal lattice, without any deeper connection the data-generating process.
transformations of the input data, it may also discard useful information. For example, consider the oft-discussed example of classifying "6"s and "9"s in the MNIST dataset. The two digits look very similar when rotationally aligned, but are easily distinguishable under their naturally occurring ori- entations. Thus under rotational augmentation, this task becomes much more difficult. In general, this discarded orientation information may be inherent, such as the previous MNIST example, or user-defined, such as the conventions used to orient crystal structures (Figure D In practice, Cohen itt al. (2018) demonstrated that rotational equivariance only improves performance on MNIST when the dataset is artificially rotated. Thus high level, distributional symmetry breaking can impact how non-equivariant methods perform relative to equivariant methods in-distribution.
Yet, quantifying the amount of symmetry breaking in distribution remains challenging, particularly in the absence of domain knowledge (Wang et al.) 2024b 2023 (024d). We thus propose metric to measure the degree of distributional symmetry breaking, which can place distribution on the spectrum between fully symmetrized (or "isotropic")
all points in each orbit {gr) are equally likely-on one side, and fully canonicalized-wher where only a single sample x in each orbit (gxJeeG is in-distribution-a the other (Figure Dp. We aim for this metric to prove useful both as a practical tool for data exploration, and as a lens for rethinking the more fundamental questions of why, and when, equivariant methods succeed.
Concretely, we propose a two-sample classifier test (Lopez-Paz & Oquab 2017). in which a model is trained to distinguish between samples from px (the original data distribution) and Px (the aug- mented data distribution) (Figure 2}. The accuracy of this classifier on a held-out test set is a natural, interpretable measure, between 0 and of distance between px and px. This (1) allows for inter- pretability methods (applied to the classifier itself), and (2) sidesteps the kernel selection required by Chiu & Bloem-Reddy (2023) in their tests for distributional symmetry, offloading impactful choice of architecture, Applying this metric to a variety of datasets, including QM9 (Wu al.) 2017), revised MD17 (Christensen von Lilienfeld 2020), OC20 (Chanussot* et al. 2021). and ModelNet40 (Wu et al., 2015), we find that all are highly non-uniform under 3D rotations.
the less
Complementing these empirical findings, we provide nuanced theory on the trade-offs between dif- ferent equivariant methods under distributional symmetry-breaking and show equivariant methods can be harmful depending on properties of the data distribution. We use ridge(less) regression as a model, which captures some of the behavior of neural networks when applied in the neural tangent kernel space (D' (D'Ascoli et al. 2020 Atanasov et al., [0233 Jacot et al. 2018). We show that even when the ground-truth function is invariant, data augmentation can be harmful when invariant and non-invariant features are strongly correlated. As our main contributions, we:
Define a flexible metric for measuring distributional symmetry breaking in a dataset (Section 2 This is tool for probing datasets' symmetry biases without a priori knowledge of their creation.
I 'Distributional symmetry breaking differs from functional symmetry breaking WWag al. 2024d), where the mapping between inputs and outputs is not fully equivariant (e.g. during a phase transition in a material). For the purposes of this paper. we treat "6" and "9" as distinct digits that simply look similar distributional, not functional, symmetry breaking. This is supported by the observation in e.g. Wang et al. (2024b) that they be correctly classified with high accuracy.
2

---

Published as a conference paper at ICLR 2026
Sumplic
Ortpral
Sec: Ponble Aparacy
SO%
IMD17 Aspile
Clast Lbbe
Canllrt Alrehass
100%
Clucsly
Figure 2: (left) Visualizations of unrotated samples from several materials datasets, with their canon- icalization visible. (right) A classifier test for determining if a sample from the original dataset, or rotated. With distributional symmetry breaking, no classifier can achieve better than 50% test accuracy. However, the original dataset was fully canonicalized, the classifier can theoretically achieve perfect accuracy (for an infinite group; otherwise, 2
Provide novel theoretical analysis invariant ridge regression under distributional asymmetry, showing that data augmentation sometimes hurts (Section 4
Use metric discover that point cloud benchmarks, including QM9 and ModelNet40, highly canonicalized (Section 5) We correspondingly evaluate the impact of equivariant meth- ods (augmentation, constrained architecture, and stochastic averaging), using auxiliary task- dependent metric for finer-grained analysis. We find surprising results on the relation to distribu- tional symmetry breaking, as summarized in 6
PROPOSED METRICS
Consider data points E 2. drawn from distribution px. acted on by compact group G. We as- sume that there a ground truth labeling function f: 2 - S that is equivariant, i.c. f(gx) = (x). We do not assume that (x(x)= px(xx;; instead, we wish to quantify the degree to which px breaks distributional symmetry by failing to satisfy this equality, i.e., to define a metric m(px) which measures how close Px is to symmetric. To this end, define the symmetrized density (x(x) == JxecPX(gx)ee invariant measure μ 2" it minimizes [(((() px (x)))) over all invariant densities i.
The density px is the closest invariant distribution to px: for any G-
We assume a dataset of samples from px. and obtain samples from Px by applying random G- augmentations. As metric of distributional symmetry breaking, some notion of distance between px and Px based on a finite number of samples - but this not straightforward choose compute.
wish to approximate
Chiu & BlodmyReddy (2023) set d to be the maximum mean discrepancy (MMD) with respect to non-parametric two sample statistical test. However, there is not always clear choice of kernel. For example, for materials datasets of geometric graphs, Chiu & "Bloem-Reddy (2023) do not provide an applicable kernel that includes chemical information.
some choice of kernel, corresponding
Rectifying this requires choosing kernel suitable for 2", which may be non- non-trivial, and as noted in Lopez-Paz & Oquab) (017), may not return values in units that are directly interpretable.
We propose instead applying two sample classifier test, common tool for detecting and quanti- fying distribution shift in machine learning (Lopez-Paz & Oquab) (017). We train small neural network NN to distinguish between distributions as a binary classification task, and define the dis- tance d between distributions as the test accuracy:
detass(po.P1 P1) = Ee- Ex-p. [1(NN(x) =c)]
Our metric is then m(px) detaso (px. p.). Concretely, construct binary classification dataset from an original dataset as shown in Figure and Algorithm with half of the dataset transformed by random group elements (label 1), and the rest unchanged (label 0).
3

---

Published as a conference paper at ICLR 2026
Algorithm Metric for Distributional Symmetry Breaking, m(px)
1: Inputs: Unlabeled train/test sets Dirain and Dien, group G, binary classifier network NN
2: For split E {trrain test):
3: Randomly divide sppli into equally sized Dspti and Dxplit
4: For each x E Dxpiit uniformly sample g ~ G and apply g to
5: Define classification dataset Splla := {(x, 0) : x E Dspii) U ((xx1,) Depia)
6: Train binary classifier NN on the dataset D Dinii with the standard BCE loss
7: Return NN's test accuracy, L(x,))Dies [1 (NN(x) = c)]
Interpretation of mpxx) The trained classifier's test accuracy is easily interpretable, reflecting how often can distinguish between the original and symmetrized distributions. If px is already group-invariant, then Px Px and no network can reliably distinguish between samples from the two, So m(px)= ~ 1/2. If in contrast px is canonicalized in a discernable way, then m(px) & 1.
To build intuition for how m(px) interpolates between these two extremes, let us also compute it for the case of a finite group, with a dataset consisting of single orbit {æa g G} := .x..) Parametrize the data distribution as p(x;) 0,, Σ, 0 1. What is the optimal classification accuracy between a uniform distribution over X2,, X, (class 1), and (class O), under infinite samples? For each i, the optimal classifier assigns 1 if ! > 6, and 0 otherwise. The resulting optimal accuracy is m(px) := - _ min(? 0,). For example, for a multimodal distribution with probability mass equally distributed among mn modes, the best possible accuracy is 1 2.- which interpolates between when == r and 2, for perfectly canonicalized distribution. In this analysis, we have assumed infinite samples, an adequately expressive NN, and perfect optimization (although it is accurate for MNIST; see T). In reality. these factors will affect m( px) as discussed in Appendix B although ablations (Appendix D indicate little sensitivity to the size of NN. For an infinite group like SO(3), we instead have m (Px) 1, and rely on validation set to avoid overfitting.
2.1 TASK- DESEEEEPNN METRIC ((x,y)
The value (Px) determines whether there is discernible lack of uniformity over group transforma- tions in the unlabeled dataset. However, does not capture whether that distributional symmetry breaking (e.g. preferred orientations) is correlated with the task labels, such as MNIST 6s/9s. If it does, then we hypothesize that augmenting is likely to be a poor choice, as it discards task-relevant information contained in the orientations (Appendix A.1 formalizes this information loss claim).
Towards this goal, we briefly introduce metric ((xxy) of task-useful distributional symmetry breaking (see Appendix for full details). Let c: 2" - G be canonicalization function, denoting where on each orbit x is. Since data augmentation destroys any information contained in c(x), we wish to understand the dependence between orientations c(x) and labels fxx). A natural way to do this is to predict f(x) directly from c(x), where c(x) is a randomly initialized, untrained equiv- ariant neural network We then compare the test loss L(c(x) + f(x)) to that obtained when the inputs are randomly transformed by elements of the given group, Liot (c(gx) - f(gx), 8 G)
which removes any task-relevant information in the orientations. Thus, (px.r) == Lga is large L << Loo i.c. the symmetry bias is relevant to the specific task at hand.
3 RELATED WORK
Learning symmetry breaking Several works seek to discover functionaa symmetry breaking, where the task may be only partially, rather than fully equivariant, i.e. there are some and g such that f(gx) * gx(x (Wang et al. 20244:; Finzi et al. 2021 McNeela [023 Hofgard et al. 2024 Smidt al. 2021 [Urbano & Romero 2024ab. We distinguish this (more common) notion from our focus, distributional symmetry breaking (p(x) * p(gx)), which Wang et al. 20233 2024c) showed can harm the performance of equivariant models. Indeed, several works proposing equivariant methods
This is closely related to the concepts of V-information Xu & Raginsky 2017 and the information bot-
tleneck (via the canonicalization) (Tishby al. 1999.
4

---

Published as a conference paper at ICLR 2026
have noted that the improvement of their method relative to baselines relies applying test-time augmentations (Cohen et al. 2018 Kaba et al. 2023). This motivates our method.
Learning how to augment Learning augmentation distribution is one way to address either kind of symmetry breaking. [enton et al. (2020a) address functional symmetry breaking by learning an augmentation distribution. For example, Miao al. (2023) encode an input using invariant network, then use this encoding to sample from a learned distribution, feeding randomly transformed inputs into classifier. [rbano & Romero (2024b) pursue a similar goal in self-supervised setting, and show their method can be used to canonicalize data, or detect when an input is transformed out of distribution. Learning to predict transformations applied to data, which is possible only with distributional symmetry breaking, was proposed for representation learning by MyGiidiii et al. (2018). We explore how c(px) can be used for selective augmentation in Appendix D.2.4
Detecting distributional symmetry In the unsupervised setting, Desai et al. (2022) and Yang et al. (2023) train discriminative networks for symmetry discovery in similar way to our binary classifier, but do not produce a quantitative measure of distributional asymmetry on benchmarks. [Chiu & Bloem-Reddy (2023) consider non-parametric hypothesis tests for distributional symmetry, and use the distance between the group-averaged and original distributions as the test statistic. [Soleymani ctt al. (025) devise a robust kernel test for invariance, where a witness g E G must be provided to prove p is non-invariant. Charvin et al. (2023) propose information theoretic framework for detecting distributional equivariance (rather than invariance, as we consider here).
Pros and cons of invariant methods Our theoretical work follows up on Elesedy & Zaidi (021): 1);
Chen et al. (200)) who show that when Ps is invariant, symmetrization or data augmentation im- prove risk. Most existing work that studies the benefits of invariance in over-vvrannnnn wer-parameterize settings similar to ours also assumes invariant Px (Mei et al.] 2021 Bietti et al., (021). On the limitations of invariant methods. [Shao et al. (2024) established that any equivariant algorithm applied to extrin- sically equivariant data, under certain assumptions on the hypothesis class, cannot obtain optimal sample complexity in terms of PAC learnability. [Lin et al. (2024); Huang et al.] (2025) also study unexpected effects of data augmentation, although not focusing on the effects of symmetry.
THEORY: INVARIANT REGRESSION UNDER DATA ASYMMETRY
To exhibit the subtleties of distributional cummetry-breaking we analyze high- figh-dimenssonal ridge regression under non-symmetric covariance. We show that even when the ground-trute function is invariant, data augmentation and symmetrization can be harmful when invariant and non-invariant features are strongly correlated. This is intuitive: a non-invariant feature is useful for invariant task it correlates well with an invariant feature used by the ground truth function, and augmentation renders such a non-invariant feature unusable. Perhaps surprisingly, data augmentation is always helpful in the under-paree parameterized regime, while in the over- weerparammttrized regime it can be harmful even when data is fully symmetric.
Suppose G < (d) acts linearly on Rd, and let = x B & for wi.iddd.o a and e ground truth B (i.e. gß = B for all g E G). Importantly, we do not assume gEg" (so Py may not be invariant). Given data {( }/_ and > O, we consider the ridge regression
problem, B,_ = arg ming 1 I|ly - x B1| + 2 2| = (L + 21) Syx where & = x X /n and E X y/nn for X E Raxd the matrix of samples and y E R" the label vector.
There are several natural approaches to enforeing invariance. Under the standard inner product, Rd decomposes into two orthogonal subspaces Vi and V1, where Vo is the dy-dimensional set of vectors invariant to G. In the first approach. Elesedy & Zaidi] (202T) consider test-time symmetrization, Fg[gβ]= = P.ß, where Po is the orthogonal projection onto Vo. The second approach is to use only the invariant features in the data, Br. in arg argmin ! IIy - (XX)) 21 + = (É..v + -1)) Eyy.inv where E.m (X)) XPo/n = PÓEA, and Eyxinv (XP) y/n = PoL... In the linear setting, this turns out to be the same as ridge regression when (1) restricting B to be invariant, or (2) under infinite data augmentation, i.e. the model sees (gx,,y Vg E G (Appendix A.3).
While such a B must be 0 if G = O(d). the same is not true for G S O(d). egg. G a subgroup of the permutation group that acts only subspace of Rdd
5

---

Published as a conference paper at ICLR 2026
For any estimator B, are interested in its generalization enror (or risk) on unseen data. Condi- tioned on the input data X, it takes the form Rx (B) Fx,e(( ((x B --x' x B)] E. [llB [
where || - B TBB. Elesedy & Zaidi 2022)) prove that when Px is invariant, one can always do
better by symmetrizing at test time: Ex [xx (XoB)] & Ex Ex(Rx(6 (B) (even for non-linear predictors). We
study the behavior of non-invariant Ps under two settings.
Under-pun paramerrrrrmetrieed ridgeless regime: When d <-n-1 1 and 2 + 0, correlations between in- variant and non-invariant features can drive Ex [Ry (P. the risk of test- test-time symmetrization, to infinity. But surprisingly, data augmentation always helpful, even if pa is not invariant.
paramaaaredd regime: This setting captures some of the behavior of real-world neural networks (D Ascoli et al. 2020; Atanasov et al., 2023; Jacot et al., 2018). When d > n, we use minimal model to show data augmentation be harmful when there are strong correlations between invariant and non-invariant features, particularly when they lie in space of dimension significantly smaller than d.
4.1 THE UNDER-PARAMETTE PARAMETERIZED RIDGELESS REGIME
Using straightforward expressions for the vias-variance decomposition (see Lemma 2 Appendix), we show that data augmentation always improves generalization when d < n and +00
Theorem 1. In the dld R-d-1
ridgeless setting, assuming Σ is full-rank, E[Rx(B)] =
E(xxx (Bm::) da 1. so augmentation helps. In contrast, for test-time symmetrization we
d'
have (Px()) n-d-I Tr(E-'Eim)
0²
ddda n-d-I
with equality when px is invariant.
While Elesedy & Zaidi (2021] Theorem 7) prove a non-negative gap for test-time symmetrization when Ps invariant, we see its risk can be much larger than that of regular (unconstrained) linear regression, when E-1 does not "align" with Ein. (This is illustrated in an example in Appendix
4.2 THE OVER PARAMETERIZED REGIME
We next consider d > n, taking the regime n, + and d/n to get deterministic estimates of the risk (assuming Σ has bounded spectrum). It is known that as y + 1, the test risk of the usual ridgeless estimator B diverges. (Hastie et al. 2022). Similarly, one may easily show that data augmentation can lead to a divergence in risk even when p. is perfectly symmetric. Namely, at the interpolation threshold do/n + 7 1, effective dimension (consisting of invariant features) equals sample size.
Fact 1. For identity covariance E= 1, in the 2 + 0 limit we have asymptotic risk R R(Bin)) = (1
1 ) + 0' /(1 111-)1)) '))) which explodes as % + 1.
To try to isolate the effect of the interpolation threshold from that of bymmetry- we suppose
% > 1, i.c. there are many possible invariant features to choose from. For tractability, we consider a minimal model for the covariance. Letting de < min (do.d d do) be the number of strong "coupling letE====- modes," let = (a.- G,) Ed_ 4 + G,,I, where a. > a, are the coupling and weak (or "white")
eigenvalues, and ugg = (Vo. + v..K)// /2 are perfect superpositions of orthogonal basis elements of Vo and _1. We consider G - 0 as the limit of strong correlations. Using random matrix theory (Atanasov et al. 2024b; Bach 2024b). we characterize the asymptotic risk for data augmentation (Appendix A.7 We find that data augmentation is guaranteed to perform worse when # correla- tional modes < ambient dimension (Figure 3 while in other settings the story is more complex.
Theorem 2. Let de/n - Y and consider the ridgeless limit 2 + O, and n,d -> co, In the limit of strong correlations: (i) if Ye < both methods are unbiased and data augmentation has larger variance; (ii) For Ye > 1, both methods have bias C( 1)/2y, where C(B) (eq. (56)) an explicit constant measuring how much of B lies in the coupling subspace, and if moreover % - y./2 < 1/2, then data augmentation has larger variance for small ow > 0.
"This involves version of the "two-point" deterministic equivalence studied by Atanasov et al. 2024a), of which we provide different proof.

---

Published as a conference paper at ICLR 2026
Mals
Fau varsece
62.5
44
53.
00 34 00
64
24
22
Figure 3: Left: Results in our minimal model for the parameetettrr regime, with mean and standard deviation across 200 trials. Small G, corresponds strong correlations between invariant and non-invariant features, and large oμ to no distributional symmetry breaking. In agreement with our theory. data augmentation Appendix [A.222 for details on the group and hyperparameters. Right: Corresponding values of m(px) at varying G,, accurately reflecting the change distributional symmetry-breaking.
harmful in small G, setting, the excess risk is higher. See
Comparison of train/test augmentation, group-averaged, and equivariant models across datasets. Augmentation: TT train+test, TF train only, FT = test only, FF = none. MNIST uses C4 group-averaged model: other datasets use stochastic croupaveeaaging MAE reported for QM7b/QM9; equivariant baselines from 3nn Geiger et al. 2022). Best overall in bold, best within augmentation underlined. CNN used for MNIST, graph transtormer for point clouds (Shi et al., 2022 Ying et al. (2221) See Figure
Table 1:
X for results relative to FF. Results are averaged over 3 random seeds and + standard deviation is reported subscript.
Setting Dataset Units
QM7b j
QM7b Øsso 10 10 (L)
QM9 C,
QM9 j
QM9 Ae mcV (L)
MNIST
ModelNet40 % (1)
10 a.u. (4)
mcal/mol K ())
10
3 D ())
% (1)
Equivariant
41+2
6.840.3
11045
131
151
18.00204
60.06=0.3
Group Averaged
4342
5.320.4
128
217+8
166_3
98.0=0.4
61.87
TT
5345
5.5
14743
243
171
97.640.2
62.39 +0.4
FF
100+
7.2
14346
275
181
98.8...1
78. 78.14 +0.3
TF
5245
5.5
14643
245
172
97.6
61.97
FT
161
12.5
210
471
278
40.7.0.2
16.91
m(px) (%)
89.66.1.02
98.340.1
86.9
93.9
5 EXPERIMENTS
Our theoretical analysis suggests that equivariant methods can be detrimental under distributional symmetry breaking; we now investigate this phenomenon widely-used datasets. The experiments serve multiple goals. First, we validate m(px) for quantifying distributional symmetry breaking by synthetically transforming subsets of datasets, verifying that m(px) has the correct behavior (Fig- ure 4. Second, we compute m(px) to investigate the degree of distributional symmetry breaking in several benchmark datasets across domains, and detect high levels of symmetry bias, distri- butional symmetry breaking. To understand downstream implications, we then compare equivari- and non-equivariant methods on the datasets' associated regression tasks, testing the predictive power of our theory (Table I We expect that, due to the distribution shift induced by augmentation on highly canonicalized datasets, training augmentation will hurt performance. However, this not the case for molecular datasets QM7b and QM9. These counterintuitive results motivate further in- vestigation of task-dependent (Table 2 and local (Figure distributional symmetry breaking, after which more coherent picture emerges (Figure 6). See Appendix D for further experimental details and results. We now discuss each dataset in turn.
TThe goal of this experiment is not achieve state-of-the-art performance, but study the impact of different augmentation settings across varied datasets and symmetry conditions.

---

Published as a conference paper at ICLR 2026
We start with MNIST (Deng 2012). where digits should intuitively be mostly canonicalized with respect to 90" rotations (C4). m mlpx) verifies this, showing that transformed and untransformed sam- ples can be distinguished with nearly optimal (1 (1-33 1 = 87.5%) accuracy, matching the calculation 2-4
from Section 2 We further sanity check m(px) by rotating p-fractions of the dataset (Figure [2], where it achieves nearly optimal accuracy at intermediate levels of canonicalization, too. This is a relatively easy task, so there not large difference between augmentation settings, although the no augmentation (FF) setting does perform slightly better (Table
Moving from 2D images 3D shape classification, ModelNet40 (Wu et al. 2015) provides more complex benchmark dataset, consisting of 12,311 CAD models across 40 common object categories. The version most commonly used in recent works is a pre-aligned variant (Sedaghat et al.) 2016), as confirmed by high m(px), ). We also apply the metric per class (Figure 9 indicating that certain classes are more canonicalized than others. Consistent with our intuition, the FF setting outperforms other augmentation strategies, suggesting that training augmentation can destroy useful information contained in orientations.
Shifting to molecular property prediction, QM9 consists of 133k small stable organic molecules with < heavy atoms, together with scalar quantum mechanical properties (Ramakrishnan et al.)
2014
Wu et al. 2017). m(px) shows that QM9 is highly canonicalized with respect to rotations; see also Figure 2 The molecular conformers were generated using the commercial software CORINA
Wu et al. 2017), which contains options align SMILES strings by default (Sadowski al. 1994
Schwab, 2010: Molecclar Networks Altamira), an example of user-defined canonicalization (as in Figure I where we do not have direct accces to the canonicalization function. Analyzing the deci- sion boundary of m(px) ) allows for fine-grained analysis of this unknown canonicalization, and can be used probe the canonicalization for discontinuities (see Appendix D.4). We find that the de- gree to which equivariance is beneficial varies per property (also seen in e.g. Liao & Smidt (2022); see Table all properties), for nearly all properties, training augmentation/eqiivriinn helps (or does not significantly hurt) performance, even on the original test ser! This is akin to com- puter vision, where augmentations like blurring induce distribution shift yet still aid performance. We next consider molecular dataset with non-scalar labels to further investigate.
still
QM7b is a 7,211 molecule subset of GDB-13 (a database of stable and synthetically accessible organic molecules) composed of molecules with < 7 heavy atoms (Blum & Reymond 2009 Mon- tavon et al. 2013). We use a version of the dataset (Yang et al. 2019) containing non-scalar material response properties to explore how distributional symmetry breaking affects prediction of non-scalar geometric quantities. m(px) demonstrates a high degree of distributional symmetry breaking. which believe follows from pre-processing steps reported in Yang et al. (2019), such as using kemel- based similarity metric to arrange atoms. We find that equivariance and augmentation are particu- larly beneficial for predicting the vector dipole moment (ū), more so than for scalar properties in the dataset (see Figure B nevertheless, augmentation again improves performance for both types of properties. Thus, we see discrepancy between ModelNet40 and QM9/QM7b, in agreement with the theory that equivariance can be helpful or harmful depending on the dataset.
Before diving into this phenomenon, we finish quantifying symmetry bias in additional materials science datasets (including an LLM dataset) to demonstrate the utility of m(px).
We explore two large scale materials datasets for predicting molecular energies and forces: rMD17, containing 100k struc- tures from molecular dynamics simulations, and OC20, con- sisting of adsorbates placed on periodic crystalline catalysts (Christensen & von Lilienfeld [2020] CCanussot** et al.) (02T). Interestingly, the degree of distributional symmetry breaking varies widely between molecules in MD17 (Figure 4 see Fig- ure 35 for all molecules). We hypothesize that this is both due to the initial conditions for the simulation, and the differing physical structures of each molecule. For OC20, both the ad- sorbate and the adsorbate + catalyst are highly canonicalized,, likely due to the catalyst's alignment with the xy plane.
Aerurary vs. Augmented fortt on MD13, ove, 0r28
wath
Augrassad Tedtoon
Figure 4: Test accuracy vs ro- tated fraction for aspirin and ethanol from rMD17, OC20 sur- face+adsorbate, OC20 adsorbate, and QM9. Finally, we explore an LLM materials dataset, as there is growing interest in training large lan- guage models (LLMs) on diverse datatypes, including molecular data. To this end, Gruver et al.
8

---

Published as a conference paper at ICLR 2026
N
m(px) (%, runs)
10
55.6% 8.44
50
72.7% 4.51
100
81.5% 4.79
(%) Local QM9
Global QM9
200
88.8% + 0.62
Original
65.9 0.3
98.3 +
500
93.6% + 0.39
Rotated
50.1 +: 0.1
50.0 + 0.3
700
96.5% 0.12
Canonicalized
99.8 + 0.04
99.9 + 0.02
1024 (global)
96.6% 0.04
Figure 5: Left: The local QM9 dataset (top) and results (bottom). Right: Local ModelNet40 results.
No
Use equivariant method
Is the blas correlated with the
No
Use equivariant method
Is there symmotry bias?
No
py) >>50%?
Yes
Do you want generalize OOD?
Equiveriant method may be suboptimal; caution
Yes
Use equiveriant method
Yes
Figure 6: Practical flowchart describing, if the unknown function is equivariant, what to expect from equivariant methods as function of mpx)) and (px,y).
2024) convert crystals into text format, which requires listing their atoms in some ordering, and then train an LLM to generate new crystal structures. The authors independently noted that per- mutation augmentations hurt generative performance, even though the task is ostensibly permuta- tion invariant. We postulated that this phenomenon was due to distributional symmetry-breaking ie. conventions in the generation of atom order. We thus trained classifier head pretrained DistilBERT transformer to distinguish between permuted and unpermuted datapoints, and found m(px) 95% accuracy. (Indeed, Figure 2 in Gruver al. (2024) reveals clear ordering in the atoms; but with thousands of datapoints, a systematic test is uselul for quantitative verification.)
In summary, our experiments thus far show that many benchmark point cloud datasets actually quite aligned We emphasize that m[px) is an easy-to-train metric that MI. practitioners can use to detect and quantify distributional symmetry breaking, which is often not known priori. Surpris- ingly, even though all datasets have high degrees of symmetry bias, the relative performance of data augmentation varies by task: train-time augmentation on Nede and MNIST hurts test-time performance on the original test set ("TF") relative to training without augmentations ("FF"), while wrai-tiim augmentation on QM9/QM7b does hurt, and often even helps (for some properties), on the unaugmented test set ("TF")) We now explore two hypotheses for this differing behavior.
5.1 TASK- SES--DEPNNDENT METRIC
As discussed in Section 5.1 if the symmetry bias in a dataset is predictive of a particular task's labels, then training augmentation might be particularly harmful. We use ((xxy) ) assess this hy- pothesis. To validate (Px.y), we first explore an artificial QM7b dipole canonicalization. This is constructed example of a very task-relevant canonicalization, where molecules are aligned to place their dipole moments on the axis. This makes it trivial for non-equivariant model to predict dipoles, while an equivariant model cannot exploit this alignment. Table D provides empirically confirmation: the FF augmentation setting outperforms equivariant model. r(px.r) is large and captures this behavior (Table Similarly, ModelNet was another case where equivariance harmed performance, and (px.y) y is significantly greater than 1. In contrast, for the QM9 properties shown, the metric shows a relatively small signal. Thus, (px.y) is generally larger for tasks where equiv- ariance hurts performance, leading to the recommendations of Figure In the future, we aim to more rigorously determine whether there exists a threshold above which equivariance is unlikely to provide benefits.
Although high mx(x) does not precisely mean the datasets are perfectly canonicalized, particularly for infinite groups like SO(3), does mean that datapoints have clear, sparse preferred orientations.

---

Published as a conference paper at ICLR 2026
Table 2: Task-dependent metric r(xx,)) Accuracy (ModelNet) or MAE (QM7b/QM9) of predict- ing f(x) from c(x), versus a random baseline. r(xx,r shows how L deteriorates under rotation: Ltoa/ / for loss, and L Liot for accuracy. Values are averaged over five seeds.
Dataset
L
Ltot
Relative Improvement L Compared to Liot
QM7h Dipole μ (4)
0.128 + .006
0.39 0.04
3.54 + 0.64
QM7b Orig (1)
0.38 + 04
0.39 0.04
1.03 + .03
ModelNet (t)
12.5
8.9 = 0.5
1.41 0.08
QM9 C, (L)
3.07 1 .03
3,22 + .02
1.05 1.01
QM9 [i] (L)
1.14 011
1.17 = 0066
1.02 +.01
QM9 Ae (L)
1.02 1 0.01
1,073 = .006
1.048 + 0.005
5.2 LOCALITY EXPERIMENTS
(Px,y) thus signals when the symmetry bias of a dataset is correlated with its labels, providing an explanation for why equivariant methods might hurt. However, this doesn't 't explain why equivariant methods can perform well datasets where r r( is low. One hypothesis is that the features are equivariant functions of their receptive fields, meaning equivariant CNNs and GNNs naturally have local equivariance (Musaelian et al. 2023). It may be useful to compute locally equivariant features, e.g. fcaturizations of small. recurrent chemical motifs in molccules, rather than just globally equiv- ariant features (Du et al. 2022 Lippmann et 2025). This provides plausible explanation for the effectiveness of equivariant methods on highly canonicalized datasets such as QM9. Moreover, augmenting inputs local (e.g. pessage-passing architecture implicitly conveys bias towards local equivariance. While it is challenging to establish causal link, we can use m(px) to at least quantify the hypothesis that local motifs are comparatively more isotropic orientation.
Concretely, we generate the local QM9 dataset by extracting local neighborhoods (by bonds) from each QM9 molecule (see Appendix D.6 for details). In Figure 5 we compare m(px) between local and ordinary QM9 in three settings: the original datasets (exploration). and under random rotation and manual canonicalization (as sanity checks, which should and do yield 50% and 100%). We find that the detection accuracy is much lower for local QM9, indicating lower degree of local distributional symmetry breaking! For ModelNet40, we analogously constructed a local dataset by randomly selecting one point from each original, 1024-point point cloud, and then collecting its N nearest neighbors. When the number of sampled points is small, the metric drops significantly. indicating that local regions of the point clouds are not inherently canonicalized; this effect reduces with the size of the neighborhood. These findings suggest that the symmetry bias present in some point cloud benchmarks is weaker at the local scale, which may partially explain the success of equivariant methods. Investigating this hypothesis further is an interesting direction for future work.
6 CONCLUSION
In this work, we aimed to provide both empirical and theoretical analysis of distributional asym- metry and its implications for learning. Our interpretable metrics quantify the degree of symmetry- breaking present in dataset without using any specific knowledge of the domain, thus providing practioners with a simple diagnostic for detecting distributional symmetry breaking (symmetry bias) in their datasets. Experiments revealed a high degree of symmetry- -rreakii in every benchmark dataset, yet augmentation only really impeded (test) performance for ModelNet40.
Overall, these findings have intriguing implications for equivariant learning. First, they affirm that evaluated only on in-distribution validation data, non-equivariant models may appear accurate, yet fail generalize under transformations. Assessing whether this is problematic requires domain expertise: see Figure [ for a practical flowchart. Moreover, applying canonicalization to data has been proposed as flexible method for making black-box models globally equivariant (Kaba al.) 2023). However, if molecular datasets are already nearly canonicalized yet still experience benefits from equivariance, this suggests that equivariant networks and augmentation may provide some ad- ditional, possibly domain-specife benefit beyond global equivariance that is currently unexplained. Finally, data augmentation is often considered universally beneficial for invariant tasks, yet we show that it can sometimes hurt performance on the test set.
10

---

Published as a conference paper at ICLR 2026
Acknowledgments We thank Julia Balla, Ameya Daigavane, Oumar Kaba, Teddy Koker, Mit Kotak, Yi-Lun Liao, Ryley McConkey, Ankur Moitra, Samuel Stanton, Behrooz Tahmasebi, Yan Zhang, and the anonymous reviewers for useful discussions and feedback.
R.W. would like to acknowledge support from NSF Grants 2442658 and 2134178. This work is supported by the National Science Foundation under Cooperative Agreement PHY-2019784 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, i iaifi... saiff..rg)................)))ii) H.L. is supported by the Fannie and John Hertz Foundation. E.H was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC0024386. This research used resources of the National Energy Research Scientific Computing Center (NERSC), Department of Energy User Facility using NERSC award ACSR-ERCAP0033254
This report was prepared as account of work sponsored by agency of the United States Gov- ermment. Neither the United States Government nor any agency thereof, nor any of their employees, makes any warranty, express or implied, or assumes any legal liability responsibility for the ac- curacy, completeness, or usefulness any information, apparatus, product, or process disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, or service by trade name, trademark. manufacturer, or otherwise does not necessarily constitute imply its endorsement, recommendation, States Government or any agency thereof. The views and opinions of authors expressed herein do necessarily state reflect those of the United States Government any agency thereof.
favoring by the United
Reproducibility Statement We describe experimental and model details in Appendix D Our code is also publicly available github. com and can be used to reproduce our experiments or try a new dataset.
/distt selistssyman sas-s-ymmm-eeeeee breaking.
REFERENCES
Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In Advances in neural information processing systems (NeurIPS), 2019.
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. In The Eleventh Interna- tional Conference on Learning Representations, 2023. URL httph: httips://ppenrriee..ccr
Alexander Atanasov, Jacob A. Zavatone- Veth, and Cengiz Pehlevan. Risk and cross validation in ridge regression with correlated samples, 2024a. URL httpp::/aarxivv.org//bb//// 04607
Alexander Atanasov, Jacob A. Zavatone- Veth, and Cengiz Scaling and renormalization in high-dimensional regression, 2024b. URL btttps:/arxxvv.rr// wttp://///.aa.a..a/aaaa/ 00592|
F. Bach. Learning Theory from First Principles. Adaptive Computation and Machine Learning series. MIT Press, 2024a. ISBN 9780262381369, URL/attps://books.ggoo co. uk/
books?id-B TEEAAAQBAJ
Francis Bach. High- figh-dimenssooal analysis of double descent for linear regression with ran- dom projections. SIAM Journal Mathematics of Data Science, 6(1):26-50, 2024b. doi:
10.1137/23M15588711 URL ///T//////////...////0////////11
Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Phys.
Rev. B, 87:184115, May 2013. doi: 10. R0.1103/Pyyseev..87.... 8...........1111.... URL https://link 1101 ors.orrg/ooi// 1et.eo1111//ppeeee. Pwt.///hhhs.....ee
184115
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew G Wilson. Learning invariances in neural networks from training data. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information PrciiionoPProoooo Processing Systems, volume 33, pp. 17605-176111 Curran Associates, Inc., 2020a. URL ht /tt.r//ooo/./e/////ooo..ooo.....
files/

---

Published as a conference paper at ICLR 2026
Gregory Benton, Marc Finzi, Pavel Izmailov, and Andrew G. Wilson. Learning invariances in neural networks. In Advances Neural Information Processing Systems 33 (NeurIPS 2020), pp. 21732-21 21744, 2020b. URL arXiv:2010.1 1882.
wttp://ppgpbbba....ba//o 11882
arXiv preprint
Alberto Bietti, Luca Venturi, and Joan Bruna. On the sample complexity of learning under geometric stability. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.). Advances in Neural Information Processing Systems, 2021. URI.httss: http::::///pperreiiiiee net net/forum
FeEv1f02TKa5L.
L. C. Blum and J.-L. Reymond. 970 Million Druglike Small Molecules for Virtual Screening in the Chemical Universe Database GDB-13. Journal of the American Chemical Society, 131:8732- 8733, 2009. doi: 10.1021/ja902302h.
A. Caponnetto and E. De Vito. Optimal rates for the regularized least-squares algorithm. Foun- dations of Computational Mathematics, 7(3):331-366) 2007. doi: 10. URL /ettps:/ddoi.orr 1tt.//1111.1100000 000.:///112000000 006 00900000-00088--1111110010---
Lowik Chanussot*, Abhishek Das*, Siddharth Goyal*, Thibaut Lavril*, Muhammed Shuaibi*, Mor- gane Riviere, Kevin Tran, Javier Heras- Deeras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 2021. doi:
10.1102/acsccttal00005...
Hippolyte Charvin, Nicola Catenacei Volpi, and Daniel Polani. Towards information theory-based discovery of equivariances. In NewrlPS 2023 Workshop Symmetry and Geometry in Neural Representations, 2023. URL
wttp://wwww.eeetttttoeeoo------ 1051
Shuxiao Chen, Edgar Dobriban, and Jane H. Lee. A group-theoretic framework for data augmenta- tion. Journal of Machine Learning Research, 21(245):1-71, 2020. URL |tthtttp:// org/
paaers/v21/20-22 html
Kenny Chiu and Benjamin Bloem- Mloem-Reddy. Non-parametric hypothesis tests for distributional group symmetry. In NeurlPS Al for Science Workshop, 2023.
Stefan Chmiela, Alexandre Tkatchenko, Huziel
Sauceda, Igor Poltavsky, Kristof Schütt, and Blaus-Robert Müller. Machine learning of accurate energy- molecular force fields. 2017. doi: 10.1 112//ccaadv.1603000 URL httl:::tt/////// wwwpwwwwwwww.w.wwwwwwwwwwwwwwwwp ortp://eerr//o///..... Publisher: American Association for the Advancement of Science.
sci adv. 1.tvv.6600 1603015
Anders Christensen and O Anatole von Lilienfeld. On the role of gradients for machine leam- ing of molecular energies and forces. Machine Learning: Science and Technology, 1(4): 045018, oct 2020. doi: 10. 10.188/2632-2113/abbe 1088/2632-21 21133322215/b/baaa6
URL https://dxx doi org/10.
Deborah Cohen, Amit Daniely, Amir Globerson, and Gal Elidan. Learning flrstfirst classifiers. In AISTATS, Proceedings of Machine Learning Research. PMLR, 16-18 Apr 2019.
Taco S. Cohen and Max Welling. Steerable CNNs. arXiv preprint arXiv:1612.084400 2016a.
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International Conference Learning Representations, 2018. UR//htppss//pppnerriie..ttttt
id=Hkbd5x2Rb.
T.S. Cohen and M. Welling. Group equivariant convolutional networks. In Proceedings of the International Conference on Machine Learning (ICML). 2016b.
Stéphane D' Ascooi,, Maria Refinetti, Giulio Biroli, and Florent Krzakala. Double trouble in dou- ble descent: Bias and variance(s) in the lazy regime. In Hal Daumé and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. https::/proceedinng wtw.e/tngsssssssss////eo..s.s.ss
PMLR, 13-18 Jul 2020. URL www.prrs/vee. 9/ aaco 120a httl
12

---

Published as a conference paper at ICLR 2026
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J Guibas. Vector neurons: A general framework for so (3)eqqivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 12200-12209, 2021.
Li Deng. The mnist database of handwritten digit images for machine learning research. 1EEE
Signal Processing Magazine, 29(6):141-14)) 2012.
Krish Desai, Benjamin Nachman, and Jesse Thaler. Symmetry discovery with deep learning. Phys- ical Review D, 105(9):096031, 2022.
Edgar Dobriban and Yueqi Sheng. Distributed linear regression by averaging. The Annals of Statis- tics, 2018. URL https: / /appi http:/pp....aaaitiieeememe htt.://ccooo...oooooooooooiaa... :8:511p
Weitao Du, He Zhang. Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie- Yan Liu. SE(3) equivariant graph neural networks with complete local frames. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 5583-5608. PMLR, 17-23 Jul 2022. URL https: htttps:/pproeeeeeeee http:///oocceddeeccooooooo ngs. mlr. wwx.ppees/vll22. du22e html
Nadav Dym, Hannah Lawrence, and Jonathan W Siegel. Equivariant frames and the impossibility of continuous canonicalization. In Proceedings of the 41st International Conference on Machine Learning. ICML'24. JMLR.org. 2024.
Bryn Elesedy and Sheheryar Zaidi. Provably strict generalisation benefit for equivariant models. In International conference on machine learning, pp. 2969, PMLR, 2021.
Carlos Esteves, Yinshuang Xu, Christine bien-Blanchette, and Kostas Daniiiidis. Equivariant multi- view networks. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1568-1577, 2019.
Mare Anton Finzi, Gregory Benton, and Andrew Gordon Wilson. Residual pathway priors for soft cquivariance constraints. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: //openreview. net/ di=kmmmdddkkkK {M2ww
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice. Kos- tiantyn Lapchevskyi, Maurice Weiler, Michal Tyszkiewicz, Simon Batzner, Dylan Madisetti, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanborn, Mingjian Wen, Josh Rackers, Mar- cel Rød, and Michael Bailey. Euclidean neural networks: e3nn, April 2022. URL https:
01115581//eee 2wt22////eeooooo 6459381
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. In International Conference on Learning Representations, 2018. URL
https::/cpeezevvll
wwoee//ouummmmmm wiw.ummii-Siliiii S1v4N210
Nate Gruver, Anuroop Sriram, Andrea Madotto, Andrew Gordon Wilson, C. Lawrence Zitnick, and Zachary Ward Ulissi. Fine-tuned language models generate stable inorganic materials as text. In The Twelfth International Conference on Learning Representations, 2024. URL https:
//openrevis net. net/fform?? d
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J. Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. The Annals of Statistics, 50(2):949 986, 2022. doi: 10.1214/21-A0S2111
httthttts: wrw.oioorrg//e/ wtt.oo/g//..........1... -2114/221-A02211
Elyssa Hofgard, Rui Wang, Robin Walters, and Tess E. Smidt. Relaxed equivariant graph neural networks. arXiv preprint axXiv:2407.20471 2024.
Kevin Han Huang, Peter Orbanz, and Morgane Austemn. Gaussian and non- gaussaan universality of data augmentation, 2025. URL hwtph:://ttsttete http://awwwww...abb/b/b/a 2202. 09134
Ilia Igashov, Hannes Stärk, Clément Vignac, Ame Schneuing, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffu- sion model for molecular linker design. Nature Machine Intelligence, pp. 1-11, 2024.
13

---

Published as a conference paper at ICLR 2026
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018 URL ht httpps: htt.tpss//prrcceeddnggoo ngg. wwwhhnggn...uuuuuuueeee 4be
neur i cc.cc/pperrrrrpe pdf
#11es/ppaer/2221/1111
AkkourOumar& Kaba, Amnb Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In International Conference on Machine Learning, pp. 15546-15566. PMIR, 2023.
David M. Klee, Ondrej Biza, Robert Platt, and Robin Walters. Image to sphere: Learning equiv- ariant features for efficient pose prediction, 2023. URL httttpp:://px...oo//////tat//e/t
13926
Joseph Kleinhenz and Ameya Daigavane. e3tools. peascient-designee
URL https://github. com/
mmiin
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. 1990, 2022.
Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. arXiv preprint arXiv:2306. /2059, 2023.
Chi-Heng Lin, Chiraag Kaushik, Eva L. Dyer, and Vidya Muthukumar. The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective. Journal of Machine Learning Research, 25(91):1-88, 2024. URL http: http///mlr.c/// org. wtwpp//pppp///......
22 15-1312.har htm
Peter Lippmann, Gerrit Gerhartz, Roman Remme, and Fred A. Hamprecht. Beyond canoni- calization: How tensorial messages improve equivariant message passing. In Proceedings of the International Conference on Learning Representations (ICLR), 2025. URL https://
Poster, arXiv preprint arXiv:2405.15389k3
David Lopez-Paz. and Maxime Oquab. Revisiting classifier two-sample tests. In International Con- ference on Learning Representations, 2017. URL. https::/
net/ /tt/ffrrmm/
id SJKXEE5xX
Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Reddy. On the benefits of invariance in neural networks, 2020. Ut///ttpss///axxxvvorr///////b// 001 78
Daniel McNeela.
Almost equivariance via lie algebra convolutions.
arXiv preprint
arXiv:2310.13164 2023.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Mikhail Belkin and Samory Kpotufe (eds.), Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 3351-3418. PMLR, 15-19 Aug 2021. URL
ngs mmlr.
press/ /
hhmll
Ning Miao, Tom Rainforth, Emile Mathicu, Yann Dubois, Yee Whye Teh, Adam Foster, and Hyun- jik Kim. Learning instance-specific augmentations by capturing local invariances. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scar- lett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 24720-24736. PMLR, 23-29 Jul 2023, URL https::/pproeeaiinn
Molecular Networks Altamira. 3D Structure Generator CORINA Classic. http: www. www.mmmn-aa
com
MN-AM, Nuremberg, Germany.
Grégoire Montavon, Matthias Rupp, Vivekanand Gobre, Alvaro Vazquez- Vazquez- Mayagoitia, Katja Hansen, Alexandre Tkatchenko, Klaus-Robert Müller, and O Anatole von Lilienfeld. Machine learning of molecular electronic properties chemical compound space. New Journal of Physics, 15(9): 095003, 2013. URL _R..ttpp: /tttt::/stack.... /ttakkaa.kk.... htt.a//w....oaao//aa////////////
14

---

Published as a conference paper at ICLR 2026
Kieran Murphy, Carlos Esteves, Varun Jampani, Srikumar Ramalingam, and Ameesh Makadia. implicit-pdf. Non-parametrii representatiioonnn of probability distributions on the rotation manifold, 2022. /tRL
05965
A. Musaelian, S. Batzner, A. Johansson, L. Sun, C. J. Owen, M. Kornbluth, and B. Kozinsky, Learn- ing local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14( 14(1))579, February 2023. ISSN 2041-1723, doi:
Long Nguyen, Xuan J. Wainwright, Martin and I. Jordan, Michael' On suiTO-
gate loss functions and f-divergences. 904, 2009. doi: 10.1214/8-A0SS595
The Annals of Statistics, 37(2):876 URL htttp:://projccteeeeete http:////voouuee///sssseuee On surrogate loss- www.tooonttoossaaaaa and ott.r//ooo.aaaaaoeeeeeeeeeeeee.d
08 AOS595. full
Pratik Patil. Facets of regularization in high-dimensional learning: Cross-validation, risk monotonization, and model complexity. Phl) thesis, Carnegie Mellon University, 12
2022. URL httts:://kittute
cmu cmu. edw httpp///.c.///rrrreee////ee/////
regularization_ing high Srgssaaee..aaaalaoaaaaaaaiarennn risk risk_ wtw._/omooooiiitioonnnnddado modell _cmpp_exitty/111222222
Cross ott. validation
Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and Anatole von Lilienfeld. Quantum
chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014.
Francisco Rubio and Xavier Mestre. Spectral convergence for general class of random matri- ces. Statistics & Probability Letters, 8 81(5):992-602, 2011. ISSN 0167-7152. doi: https://doi. org/10.1106// 0tg///0..010///sp12111......0 URL http////////////////ooooaoao.oo.
atticiee/11/50017715110000
J. Sadowski, J. Gasteiger, and G. Klebe. Comparison of automatic three-dimensional model builders using 639 x-ray structures. Journal of Chemical Information and Computer Sciences, 34: 14:0 1008, 1994. doi: 10.1021/c0000200003,
C. H. Schwab. Conformations and 3d pharmacophore searching. Drug Discovery Today: Technolo- gies, ((4)=245-e255) 2010. doi: 10. ddto://.ddecc.2100....
Winter 2010.
Nima Sedaghat, Mohammadreza Zolfaghari, Ehsan Amiri, and Thomas Brox. Orientation-booste voxel nets for 3d object recognition. arXiv preprint 1604.v660.003333 2016.
Han Shao, Omar Montasser, and Avrim Blum. A theory of pac learnability under transformation in- variances. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 978171387 1088.
Yue Sheng and Edgar Dobriban. One-shot distributed ridge regression in high dimensions. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 8763-8772. http://a.......eee/////eeeeeeeee
PMLR, 13 18 Jul 2020. URL
html
Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie-Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint arXiv:2203.0444 04810, 2022. URL httpp::ts://arre org abs
/btaaabb//222222 2203. 04810
Tess E Smidt, Mario Geiger, and Benjamin Kurt Miller. Finding symmetry breaking order parame- ters with euclidean neural networks. Physical Review Research, (1):2000 2021.
Ashkan Soleymani, Behrooz Tahmasebi, Stefanie Jegelka, and Patrick Jaillet. A robust kernel sta- tistical test of invariance: Detecting subtle asymmetries. In Yingzhen Li, Stephan Mandt, Shipra Agrawal, and Emtiyaz Khan (eds.), Proceedings of The 28th International Conference on Artifi- cial Intelligence and Statistics, volume 258 of Proceedings of Machine Learning Research, pp. PMLR, 03-05 05 May 2025. URL https: / /ttps////rcceeeennge. mlr. press/v258/ so eyman 125.. html
15

---

Published as a conference paper at ICLR 2026
Bharath K. Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert R. G. Lanckriet. On integral probability metrics, O-divergences and binary classification. arXiv preprint arXiv:0901.2698, 2009. URL (otLbb::///axite owg./iii.rr//. abs / 00000.0...... 2698
Behrooz Tahmasebi and Stefanie Jegelka. The exact sample complexity gain from invariances for kernel regression. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems, volume 36, 2023. URL bt http:sss/pppoooeeesss.......sas.
htm
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and ienollation-qquvariiann neural networks for 3d point clouds. arXiv preprint arXiv:1802.008219, 2018.
Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method. In Proc. of the 37-th Anmual Allerton Conference on Communication, Control and Computing, pp- 368-377, 1999. /taL:ttpee//aaxvv.or///aa//////
Alonso Urbano and David W. Romero. Self-supervised detection of perfect and partial input- dependent symmetries. arXiv preprint arXiv:23/2.12223, 2024a.
Alonso Urbano and David W. Romero. Self-supervised detection of perfect and partial input- dependent symmetries. In Sharvaree Vadgama, Erik Bekkers, Alison Pouplin, Sekou-Oumar Kaba, Robin Walters, Hannah Lawrence, Tegan Emerson, Henry Kvinge, Jakub Tomezak, and Stephanie Jegelka (eds.), Proceedings of the Geometry- Representation Learning and Generative Modeling Workshop (GRaM), volume 251 of Proceedings of Muchine Learning Re- search, pp. 113-131. PMLR, 29 Jul 2024b. URL https: httpss////ccceeeeengeeeeo mmr. .press/ v251/urbano24B.
Dian Wang, Jung Yeon Park, Neel Sortur, Lawson LS Wong, Robin Walters, and Robert Platt. The surprising effectiveness of equivariant models in domains with latent symmetry. In The Eleventh International Conference on Learning Representations, 2023.
Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, and Robert Platt. Equivariant diffusion policy, arXiv preprint arXiv:2407.01812 2024a.
Dian Wang, Xupeng Zhu, Jung Yeon Park, Mingxi Jia, Guanang Su, Robert Platt, and Robin Walters. A general theory of correct, incorrect, and extrinsic equivariance. Advances in Neural Information Processing Systems, 36, 2024b.
Dian Wang, Xupeng Zhu, Jung Yeon Park, Mingxi Jia, Guanang Su, Robert Platt, and Robin Walters. A general theory of correct, incorrect, and extrinsic equivariance. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2024c. Curran Associates Inc.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. In International Conference on Learning Representations. 2021.
Rui Wang, Elyssa Hofgard, Robin Walters, and Tess Smidt. Discovering symmetry breaking in physical systems with relaxed group convolution. arXiv preprint arXiv:2310.02299 2024d.
Zhenqin WU, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S. Pappu, Karl Leswing, and Vijay Pande. MoleculeNet: benchmark for molecular machine learning. Chemical Science, 9(2):513, October 2017. doi: 10. 10.10/77cc02264a.. URL htRpst 7/pmc. ncbi. n.m. nin gov hwtphhhho//aataaleeee/ceeeeeeee/
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920, 2015.
16

---

Published as a conference paper at ICLR 2026
Aolin Xu and Maxim Raginsky. Information thheoreti analysis of generalization capabillit of learn ing algorithms. I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 30. Cur- ran Associates, Inc., 2017. /ttp:ttttpp//ppooeeeeirrrreeee
neurissneuupp.s. wwtppppp.cppppp...pppppppppppp fise// /ttp://00.......222222226666666.
Jianke Yang, Robin Walters, Nima Dehmamy, and Rose Yu. Generative adversarial symmetry dis- covery. arXiv preprint arXiv:2302.002333 2023.
Yuezhi Yang, Ka Un Lao, David M Wilkins, Robert A DiStasio Jr, and Alexandre Tkatchenko. Quantum mechanical static dipole polarizabilities in the qm7b and alphaml showcase databases. Scientific Data, 6(1):152, 2019. doi: (0.1038/441597-019-111
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? In Thirty-Fisth Conference on Neural Information Processing Systems, 2021. URL URR/https://ooeneviiee net / htt.:/oooomdddddoooooooo--o---
Jacob A. Zavatone-Veth Veth. Harvard APMTH 226: Lecture notes on the inductive biases of high-dimensional ridge regression, 2024. URL
17

---

Published as a conference paper at ICLR 2026
APPENDICES
Theory
19
A.I
Gradient descent on equivariant networks is invariant under per-datapoint transfor-
mations
20
A.2
Review of the generalization gap Blesedy Zaidi (2021)
21
A.3
Equivalence of training symmetrization and data augmentation
21
A.4
Bias-variance decompositions
21
A.5
Proof of Theorem
22
A.6
Permutation example
22
................................
A.7
Deterministic equivalents for bias and variance
................................
23
A.7.I
Bias term
23
................................
A.7.2 Variance term
24
A.8 Analysis of the minimal model for covariance
25
A.8.J Proof of Theorem
25
A.8,2 2
Simulation details
26
Learning theoretic context for m px
27
Task-dependent metric
28
Derivation and explanation
28
Theory
C.2.1
Task dependent metrics for finite and infinite groups
30
Experiments
D.1
ModelNer40
31
D.I.J
Classification results on Transformer model
D.1.2
Relation between the degree of canonicalization and accuracy FFTFF aug- mentation settings per class
32
D.1.3
Task-Deepnnden Metric
32
D.1.4
Training Curve
32
D.1.5
Model details
33
D.1.6
Training details
................................
33
D.2
MNIST
35
................................
D.2.1
Training and model details
35
D.2.2
Task independnt and dependent metrics
35
D.2.3
Loss curves
35
D.2.4
Selective augmentation using the trained classifier
35
D.3
Swiss Roll
39
D.3.1
Training and model details
39
D.3.2
Results
39
18

---

Published as a conference paper at ICLR 2026
D.3.3
Loss curves
39
D.4
QM9
D.4.1
Dataset, Models, and Training Details
41
D.4.2
Task Dependent Metric
D.4.3
Interpretability of Classifier for Distributional Symmetry Breaking
43
Ablations
D.3.I
Hyperparameter ablations
48
D.5.2 Direct Task Dependetndd Metrie Ablations
48
D.6 Local QM9
51
D.7
QM7b
53
D.7.1
Dataset Details
53
D.7.2
Model and Training Details
53
D.7.3
Task-Relevant Canonicalization
54
D.7.4
Loss Curves
54
D.8
rMD17
55
D.9
Open Catalyst Project 2020 (OC20)
56
D.10 p-values for m(px)
57
D.10.1 QM9 p-values
57
D.10.2 rMD17 p-values
A THEORY
In this section, we elaborate the theory of the main paper, including both more context and proofs of results. Most of the earlier parts are dedicated to our study of ridge regression in Section
Appendix A.I we formalize the motivating intuition for the paper, i.c. that "equivariant methods lose information". In particular, we prove that under an invariant loss function, gradient descent on an equivariant architecture yields exactly the same learned parameters, regardless of the group transformations applied to the training data.
Appendix A.2 we review the result of Elesedy & Zaidi (2021), which says that when data is invariant in distribution, test symmetrization always improves generalization error. (Unlike our analysis, this holds even in the non-linear setting.)
Appendix A.3 we show that using invariant features, or equivalently, restricting to invariant estimators which call training symmetrization equivalent data augmentation in the ridge regression setting.
Appendix A.4 record the bias-variance decompositions of risk for vanilla ridge regres- sions, test-time symmetrization, and data augmentation.
Appendix A.5 we prove Theorem which generalizes Elesedy & Zaidi 2021
Theorem
7) to non-invrriant data, and to data augmentation.
Appendix(A.6 we demonstrate in explicit example that test-time symmetrization can ar- bitrarily increase risk when the data distribution is not invariant (in the inder-parameterizer regime).
Appendix A.7 using random matrix theory, istic equivalents") for the bias and variance of cach of our three estimators in the over- parameterized regime. Fact direct corollary.
we derive asymptotic expressions ("determin-
Appendix(A.8] we analytically study our minimal model of covariance, proving Theorem We confirm our results empirically, shown in Figure
19

---

Published as a conference paper at ICLR 2026
A.1 GRADIENT DESCENT ON EQUIVARIANT NETWORKS INVARIANT UNDER PER- DETAPPOTP TRANSFORMATIONS
The guiding motivation for this paper is that symmetry bias (i.e. distributional symmetry breaking) contains information, which is lost by an equívariant method. This is self-evident if that equivariant method consists of data augmentation (since the orientation information literally thrown away) or a network that operates on invariant features, but is slightly more subtle for a generic equivariant function. Here, we formalize our assertion by proving that, when using gradient descent learn the parameters of an equivariant network under group-invariant loss, the gradient update at each step agnostic to the exact orientation of each input datapoint. In other words, an equívariant network trained by gradient descent learns the same solution with symmetry-biasee dataset and with fully- augmented version of that dataset; thus, it too is "losing" information contained in the distribution asymmetry. (In fact, this is true not just for gradient descent, but for most optimization methods, since we'll show the loss landscape is unchanged.) We now formalize this below.
To begin, let G be group acting linearly input space Vo and output space Vz. We write g g.x for the action of g E G on a vector in either space.
Let NNg Vo V2 be G-equivariant neural network, differentiable in θ O C RP:
NNe(g xx - g NNo(x) Vg G, Vx E Vo. Let the loss V x Vz + IR be differentiable in its first argument and G-invariant:
(1)
{gg u, g.v) = E(u,v)
Vg G, A (u, E V2 x Vz.
(2)
Let Si ((x))))) be a training set. Fix arbitrary 81.....88 E G and define S ((g Xis 8r-)))
with empirical risks
1 (3) :- n Σ ANNN (xF (8) (1)))
(K)),
k (1,2}.
First, we have the relatively obvious statement of loss invariance. Lemma (Per-example loss invariance). For each and all 9,
8 (NNo (g)-x;,, vr) {(NNo ya).
(3)
Proof. By equation and equation 2
((NNgggi (igiixx)) (g): 8: y;) (gi NN (x:), gt.); y:) ((NN yt).
Using Lemma we have that the two loss functions are equal, and so their gradients are too.
Corollary (Equality of risks and gradients). For all O,
(0) Ls and
V (0)= (0).
At this point, the key result fairly clear, but we state below for completeness.
Proof. Lemma 1 gives pointwise equality of per-example losses as functions of O. Average for the risks; differentiate for the gradients.
Theorem (Identical optimization trajectories). Fix (ne and 0g O. Define
e
0 (k)
-n: Vo Ls, (8(t).
ke (1,2),
with 6 = = 0g. Then 8 0, for all 2 0.
Proof. By induction. The base case holds by construction. If e!") = o.2), then Corollary 1 gives equal gradients, so the updates coincide.
is important to note here that the gi are arbitrary per-example; no common global transformation is assumed.
20

---

Published as a conference paper at ICLR 2026
A.2 REVIEW OF THE GENERALIZATION GAP OF ELESEDY & ZAIDI (0221
Consider data (x,y) E Rd x R generated as y (x) + forx Px. ground-truth invariant function f², and independent mean-zero finite-variance variance noise e. Considering L2 loss, the excess risk of a given function f (say the result of learning some fixed training datasct) is
R(ʃ) = [[(y
E((((x) - (
(4)
since -
We define a new inner product on functions, (fi.f2) (X)(())))))
The excess risk of is then |lf f 1pp. with the norm induced by this inner product.
Let f(x) = [g(((( be the symmetrization of f with respect to uniformly random g G. We can think of this as test-time augmentation. We can ask what the difference is between the excess risk of f and f.
(5) Elesedy and Zaidi show that when x is invariant in distribution, f- f is orthogonal (in the inner product defined above) to invariant functions, and thus in particular f ---f.. In this case A(f.f) 2 0, meaning for any f one can always achieve better generalization using f. When p. is not invariant, however, the inner product might make the overall expression negative. This case thus warrants further investigation.
-
+
A.3 EQUIVALENCE OF TRAINING SYMMETRIZATION AND DATA AUGMENTATION
We consider the estimator obtained by infinitely many augmentations,
Bx.aug argmin
Σ Egl(O - (gx)) B) 2|1B11 (Laug A1) - Lna
(6)
B
where Eaug (gīg J. As linear map Rd + Rd, Eaug is equivariant, so by Schur's lemma it block-diagonal in Vo, Vi (or more generally, in irreps). Since Lyx,in non-zero only in the Vo component, (Eoug + 21) - Lyrrinv = (Linv + - Eyx.inn and thus BR.ang = =Ba.inv We therefore refer to B2.iov interchangeably using data augmentation invariant features.
A.4 BIAS-VARIANCE DECOMPOSITIONS
The risk of any estimator / has variince decomposition Rx = +VxB)+xx(BBB with
Bx (x
E 8
x
B B
=x Tr(CoVB x)E)
In the case of vanilla ridge(less) regression, the expressions above have well-known and easily de- rived forms (Bach) 2024a] (Hastie et al. 2022). We list the equivalent expressions for test-time symmetrization and data augmentation below. The proof is standard, being simple expansion of definitions. One may notice that the expressions are the same as the vanilla case except (I) test-time symmetrization replaces with Σnn = PoEPo, and (2) data augmentation replaces E with E
Lemma 2. For unaugmented ridge(lesx) regression, the bias and variance terms are standard:
Vx( σ²
Bx = =
Tr( -2)
(8)
② Tr Σ)
Bx( B "IIEII
vx(
(9)
where - I- -++ projects onto the null space of X. Test-time symmetrization gives
σ² n
Bx = 2.3 (É+A1) Ein (L (É BB211
(P Tr(E(E+ "
(10)
6² Vx(P,B)= Tr(E-tim)) Emm)
(P = B 1 IIB
(11)
21

---

Published as a conference paper at ICLR 2026
whereas for invariant features and data augmentation, obtain
σ² Tr Tr (iiv +211
Bx (B( = BB (Einv + -I)) + B Vx (Bx,in))))
(12)
σ² Vx (Bm))
Bx(Binv) - "IinnyIIin
(13)
n
where Dimy Po ( In latter case, we note that every instance of Σ can equivalently be replaced with Linv, being multiplied "on both sides" by invariant objects.
A.5 PROOF OF THEOREM 1
When d <n- 1, the matrices and Iim defined Lemma 2 are almost surely equal to the zero ma- trix. Thus the vanilla, test-time symmetrization, and data augmentation estimators are all unbiased, and we compare only their variances.
Having assumed X; M0,, Σ), the empirical covariance is scaling of Wishart-distribteed matrix: E W (I, Since Σ is full rank, the standard form for the expectation of the inverse Wishart i W [E .n) gives
mΣ
T2 ΣΣ
o'd
E[Vx(B)] n
(14)
1-1
For test-time symmetrization, note that the trace term in the variance only depends on Vo components of the inverse empirical covariance: Tr(É-'Ennw) diagonal sub-matrices of unverse-Wishart matrices are inverse- -iirsert certain form. Letting be the change of basis matrix into Vo. V1, so that any matrix M can be written as
Tr(PpÉ-1PpEin.. We thus use the fact that
Moo
Moi
VTMV
M
M.T).
(15)
we have nV ÊV W EV,, and 1/ ( Don W ((E-) '0o. di) where d dim V We thus have
(P 0² Tr((2 o00 TrT (n d1) -o
(16)
d-1
For invariant features, the relevant trace term is Tr((Sim)) ) ) Einv) lows the same logic as in the vanilla case: nLoo ~ W (Eoo, n), and thus
-Loo). The result fol-
n nEoo))
E[Vx(Pmmm)) O² do
(Eoo))
6
E((200)-11
(17)
do
do
A.6 PERMUTATION EXAMPLE
Consider the case of G S3 acting on three-dimensional inputs E R by permuting coordinates. Let V be the change of basis matrix into the G-invariant subspaces Vo, Vi, and write Moo for the (Vo, block of V MV. We then consider covariance
2
p
VTEV=
(I
000 )00
(18)
P
1
o
202 1+t
such that Tr(E-'E (Inw)
20² of (L+c) This term is large when lp). the correlation strength
between invariant and non-invariant features, is large compared to the invariant signal o In par- ticular, we have E[Rx(8) (B)] < ([x)(()))) + 0o as 2p² grows from of (1 + 7) to o (1
While do not do here, this example be extended to general G and Σ by using the Schur complement formula for (I )on. in which case the "size" of Σg in the Loewner order plays the role of the correlation p.
22

---

Published as a conference paper at ICLR 2026
A.7 DETERMINISTIC EQUIVALENTS FOR BIAS AND VARIANCE
In the proportional asymptotic regime, where n,d + oo and d/n + 7, we leverage the notion of deterministic equivalence of possibly random matrices An and Bn. In partieular, write A,, = B, when for any matrices C of bounded trace norm,
Tr((An, [(, -.
(19)
In our derivations below, we take advantage of the "calculus of deterministic equivalents" devel- oped by Dobriban Sheng (2018); Sheng & Dobriban (2020), as well as proof techniques of Hastie al. (2022). which in turn rely on the generalized Marchenko-Passe
theorem of Rubio & Mestre
(2011) We also utilize the notions of first- and second-order degrees of freedom, df' and df used
by Atanasov et al. (2024b); Bach 2024bb and introduced by Caponnetto & De Vito (007 as
"effective dimension."
Our goal is to find deterministic equivalents for the matrix products appearing in the bias and vari- ance expressions in Lemma (for À >> 0). which all take the two forms
a² n
Buv = 2BBTTBB (Ep + 2112211-11 E, (2, + BB
Vuv =
22)
(20)
where μ. V run over empty or inv subscripts. Here, we assume the setting of the generalized Marchenko-Pastur theorem -at--nnn namely, that we have the deterministic equivalence
2L4n +1) (Ep + K21)
(21)
where Ky is the unique positive solution to
2
1 Tr((EE kl) k
Ky 1- (Kg)
(x) dr) (x)
(22)
and can be seen as the effective or renormalized ridge parameter; this includes the setting of i.idd Gaussian data, but extends much further, to the Gaussian universality regime (Hastie et al. 2022 Zavatone- Veth 2024). We prove the deterministic equivalences
<2 Cuv BT (Iµ + k + x2, (En + EE,(E + kxl)=
Buy 1 aup
(23)
Vuv δ² 1- Oyu
Ouv
(24)
1-0yp
where we define generalized second-order degrees of freedom,
duv n
dr2 (Kp)
( Tr((En k) EvE.
(25)
We do claim our calculation of these deterministic equivalents novel; indeed our presentation closely follows the notes of Zavatone- Veth (2024), and the "two-point" equivalences we consider were recently analyzed in the context of cross validation by Patil (2022) and Atanasov et al. (2024a).
To our knowledge, however, our work is the first to apply these techniques to study invariant learning.
A.7.1
BIAS TERM
Note that
2 +2 -I ( 21 2xE,)
(26)
=-2 2.1 +1 ATE,)
(27)
r=0
We find a deterministic equivalent for the expression inside the derivative. First, note
(/1++ 2 + 21- = 2/)-1 +E
(28)
Note that the notion used by Atanasov al. 2024b is scaled by 1/d with respect to that of Bach 2024b
the latter convention.
23

---

Published as a conference paper at ICLR 2026
where we define
(29)
That is, our expression is a product of deterministic matrices with the matrix ridge resolvent for scaled version of the empirical covariance. This resolvent thus has the deterministic equivalent
2 + 21)
(30)
where Σ. is the population covariance Σn scaled in the same way as ΣS is, and K₂ is the unique positive solution to
Kr 1 (KT))
(31)
We thus have
2 2L2x + +1 2 = Kr + k,tE,)
(32)
Under the assumption that Σn is trace class, one exchange the n, - limit and the derivative to obtain
a 22n + 21, + 21) - Ke + k,/ KTT,,
(33)
===0
Let us write 8 = = and note that Kr /r|r=n- = Kp. We get, first using the matrix identity àM
MM 1 (MM)M and then combining terms,
-O(E, + kal) + Ky (Iμ + Kg (81+ (b) kyl) 8I (Eş knl) 2 x2) (I kul) k ky(!)
(34)
(35)
remains evaluate 8, which we do by differentiating the fixed fixed-point equation t 0,
K; KrT₁ = 2
8T; =-
ky (T+)
=0. 0.
(36)
r-0
Recognizing 1- Te|==0 as 2/xy (from the ixee-point equation at r 0), we get
x
x 8Tr(( + K2l) Exu)
8 (T)))) ==0
#
Auv
(37)
àn
Since (En + kal) = (I + (En + (k()) /E/)/=
1
1
à
Tr((E (
(Trlr=0 auu)
1
Cuup
(38)
n
Ky
Ky
Subsequently,
(1
--
à
Cup
KH
auv
-
KHy (1 apuI - -1)
Kj
Cyv-
(39)
2
Ky
2
We can then solve,
ky
auv
8 =
auv = - K2 Opu
(40)
Ky (1 duu) 2
A.7.2 VARIANCE TERM
We begin by noting that
σ² Tr + AT) Ey --
0² - or Tr( + -1) - 0,2 21)
-1,). (41)
By assumption 2 (Ln, + 2/1-11 = (En + xp()- Exchanging limits and the derivative (which is justified by the assumption of bounded trace norm) get the deterministic equivalent
1 -02 Ky (r( + kyl) k,/))'_,) 'vv) - ( Tr((En + K2l) - Ev) = 22 (22(k))) (Ky))a.v.
(42)
To find the derivative, we differentiate the fixed-point equation
Ky Kp =
2g d₂ (u) 22 (kx) - =1,
(43)
where in the first equality on the right used the chain rule and dxy (T) aup. Thus, d₂ (Kμ) 1/(1 Upu), which plugged into the expression above proves the result.
24

---

Published as a conference paper at ICLR 2026
A.8 ANALYSIS OF THE MINIMAL MODEL FOR COVARIANCE
Recall that our minimal model is
d
Σ 0
Uglli + Gw
Σ k=
(44)
where the first de eigenvectors Mj (Vo. + / v1,))/ 22 represent coupling modes, and the remaining Mk complete the orthonormal basis. Note that when the coupling and weak directions have the same strength (o, - G.), reduce the isotropic case. We thus are interested what changes o ow grows, which we study by taking the oμ + 0 limit.
The fixed-point equation for
( 1
O
GN on G
1
+ (y-
=2.
(45)
o
This has the same solutions as cubic in K. One option is to directly study the large Ge limit, in which case the equation becomes independent of o... We instead take the over-vvrameee pve-pparameterized ridgeless limit (2 - () and 1), where κ solves
G
% o + (7-%)
(46)
G +
Comparison to training symmetrization must be delicate. The effective ridge parameter solves
Gw
Kinv
1 - %
(7o - Yc) G, Kjnv
(47)
σ + Kinv
we obtain an effective ridge of Kinv 0. This is intuitive: do < n we are back in the ordinary least squares regime once we project the data down into Vo. If do O(1) (i.e. the number of invariant features in the problem is finite) then are in the 70 - O regime of Theorem [ where training symmetrization helps. However, as in the isotropic example, y but 0 < % < then R(Pin) can still grow arbitrarily large as we approach the new interpolation threshold. (This also means the order one takes the 76 - and y oo limits matters; taking the latter first, for example fixing do and n and taking d - still shows harmful effects for training symmetrization.)
where 6= (a. -0,6--)w/// 0/2. We again take 2 + O. When %6 <
A.8.1 PROOF OF THEOREM 2
We consider the case 7o and 2 -> 0, in which
=
ō
+ (%
Y.)
G
(48)
o Kinv
o Kinv
This has the same form as the equation for k. Indeed, can write x(Ge, Ow, y) and Kinv x(ō, Ow, 16) where x(s, w, solves the quadratic system
x(s, w) + b(s, w. (.) w.g) + c(s, b(s,w,g) (s+ 7.5 1g-7) W/ -g)
(49)
(50)
c(s, ((s,wg) = (1 (1-g)sv g)sw
(51)
For 8 > %6, one can observe k(s,w.g) w. is increasing in its arguments. Thus, Kmn S K training symmetrization has smaller effective regularization. In the "strong correlation" limit O, have ((s,0.g) = max(0, - 1)). That is, we have a new threshold, corresponding to when the model is over-parameterized with respect the number correlational modes.
Similarly, a a (( ow:?) 7) and Cinviinv = ai Ow, 76) where
2
W
w, Yo (+
(g 76)
(52)
W w K(s,
The second term and approaches (1 + K' (s,00 as w + 0 when Ve < 1, where we use to de- note differentiation with respect w. Evaluating the derivative, in the regime of % < have
(1-x)² and thus Canv.inv (.)) 7. g-k
25

---

Published as a conference paper at ICLR 2026
In the correlationally wver-parameterized regime % > 1, the second term vanishes - 0, and we get (s,, p) - Y (using our result for k(s,0,g)). which we note is independent s and g. So, in the limit of strong correlations, so, we again find that Cnnyinv >> a in a neighborhood of w = 0 when 70 - (x/2) < 1/2, i.e. when a good portion of the invariant features are captured in correlational modes.
Oanv,inv We thus examine the derivatives a' (s, 0, p). Doing
Since x/(1 - x) is monotonically increasing, the above behavior fully describes how Vx
compares asymptotically to Vx (Binv).
Understanding the biases
Bx( 1-a x))+
x²
OF
Σ
de
[4g B)
On
Σ (uk B)2
(53)
(ow + x)2 Amdet1
(0 (10+ Kiyy)
do
Bx (Binv)
inv
EvEB))
Ow
Σ (v B)}
(54)
1 Cinv inv
(o,
Kin))? k=de+1
requires handling the dependence on B. Since ß is assumed invariant, and thus 11B11 -x (v
B)?
d
de
(1 (1- (1-C(
C(B)
[n = =1 Σ
(55)
2
where we define the coupling factor
C(B) Σ (v. B)1111112
(56)
The expressions for biases become
Bx(B)=
x² 111111
6 C(B)
Ow
(1- C(B) )
(57)
1-a a
(g. k)2 2
(o + x)2
x2 ||
Bx Bx (Bir))
C(B) +
Ow
C(B))
(58)
Oanv.inv
(ō + Kinv)
(w + Kinv)
The y. regime is straightforward reusing our previous calculations, giving
Bx( = Bx (Bim) = 0.(% - I)C(B)IBB1
(59)
2%
at G 0, while for % both go zero with G, O.
A.8.2 SIMULATION DETAILS
We describe the simulations used to obtain Figure 3 settings, and then the group symmetry.
We first describe the hyperparameter
We fix values of 100, d 5n, dg 2n, and d, = |n/2], which puts in the parr-parameterized regime, but where the model is "correlationally under-parameterize theory predicts that as Gw + O, both methods become unbiased, but data augmentation has higher variance. This is indeed what we observe in simulations. To generate Figure
(% < 1). In this setting, our
3 we set a, = 1, so
that we examine GN as a fraction of o The noise generated as E Nf (0, with a = 0.5, and we set nominal regularization À 10-8 to approximate the ridgeless setting.
In order to set do, we use the permutation group G = S + acting on the first do + coordinates of Rd. The invariant space Vn then consists of a one mie-iimesioonal subspace of R (the one with equal entries) together with the remaining d - (d dg + 1) do coordinates unaffected by G. We therefore indeed get do total invariant directions.
For the rightmost panel of the figure, we trained a 3-layer MLP with hidden dimension 256 as the classification network.
26

---

Published as a conference paper at ICLR 2026
B LEARNING THEORETIC CONTEXT FOR m(px)
To formalize the intuition associated with m(px), we piece together some learning theoretic context below.
First, what would the classifier metric actually measure, if we had infinite data? In other words, if & our class of neural networks, & is a binary classification loss function, p our data distribution and p is its randomly symmetrized version, want to understand the following:
1 2. ((((s) 1 I. ((((),
min fe
(60)
Since these are continuous integrals rather than finite-sample approximations, this is an "infinite- data" setting.
In fact, this expression has been studied extensively in various settings (sce Sriperumbudur et al. (2009), [Lpez--ao & Oquab (2017) and references within), which usually vary in their choice of & and É (and with replaced by more general second distribution q). For example, when f is the set of all measurable functions and l the 0/1 loss, eq. (60) is the optimal Bayes risk, classically known to be simple affine transformation of the total variation norm.
When instead using the loss function E(az0) 0) = a and E(a, 1) = a. Sriperumbudur et al. (2009) showed straightforwardly that eq. (60) is exactly the integral probability metric associated with :::
max EE3
(61)
Integral probability metrics recover range of familiar quantities for different choices of $ For example, when & consists of all functions with Lipschitz norm bounded by I, eq. (61) reduces to duality. Similarly, when the functions in I are norm-bounded in some Hilbert space, the resultant integral probability metric the kernel maximum mean discrepancy (MMD), as used by Chiu & floem-Reddy (2023).
the Wasserstein metric via Santorovich-Rubinge
For other loss functions and with unrestricted class F. Nguyen et al. (2009) showed that there exists some convex such that eq. (60) is equal to the corresponding / diieergencce For example, when é is the exponential logistic loss and I is all measurable functions, eq. (60) reduces known divergences (Hellinger and x² respectively).
Thus, while we cannot precisely articulate eq. (60) as depends on both the exact class of neural networks and the loss function, it reduces to reasonable and familiar quantities in special cases.
We now turn the finite sample complexity aspect of the metric. Since we cannot truly compute eq. (60). we follow standard learning practice and use a dataset of finite samples to approximate the continuous integrals. Standard generalization bounds apply, restricting how much our finite- sample approximation to eq. 460 can diverge. In particular, let E (0, 1), let be drawn with equal probability from p or P with denoting the corresponding binary class, and am (f F) be the Rademacher complexity. i.c
am(loF)= = ((s,on) Ea;
sup re.r
Σ [f(()x))))))))) m
(62)
Then standard generalization bound (see e.g. Appendix A Cohen al. (2019) tell us that, with probability 2 1- - 8 over the randomness of the training set, the following holds ff E F:
8 8n(
E(ay) ([)))))(())))))
Σ
((((x)))))+2
(63)
S
(XL.9) Nmn
optimized in
60
what actually optimize
With these tools in hand, we now interpret what learning theory has to say on the impact of the dataset size and model family on m(x)).
27

---

Published as a conference paper at ICLR 2026
Impact of dataset size The dataset size with which we train metric affects only the general- ization bound from finite to infinite sample loss, not the underlying pseudometric itself. According to the generalization bound above, the the dataset size contributes inversely to an additive term in the overall generalization error, and should scale roughly with log(1/8).
Impact model family f The choice &, i.e. the family of neural networks, affects both the generalization bound eq. 633 and the pseudometric itself eq. 660). In particular, as the architecture class becomes more expressive, the pseudometric eq. (63) grows "stronger" (i.c. can better discern subtle distributional asymmetries, with the TV nomm between distributions the limiting case), but the generalization bound eq. (63) also grows larger (i.e. the risk overfitting is greater). In practice, of course, generalization bounds not perfectly predictive of modern deep learning, and more nuanced bounds are possible (e.g. which take into account the complexity of the data distribution itself, rather than just the hypothesis class).
C TASK-DEPENDENT METRIC
In this section, we walk through the task-dependent metric more detail, including its derivation and variant that used binary classification (which was ultimately less performant).
C.I DERIVATION AND EXPLANATION
m(x)) does not capture whether distributional symmetry breaking contains useful information for the task at hand. If does, for example in cases of inherent symmetry breaking as in Figure [ then we predict that performing full-group data augmentation is poor choice, as it discards task- relevant information contained the exact position within the orbit (Appendix [A.]]. However, the distributional symmetry breaking is superficial in the sense that it has relation to the task of interest, then it is more subtle whether full- full-group data augmentation will hurt performance, as shown in Section 5 As such, seek to refine m(px) from Section 2 to produce stronger signal for when data augmentation is harmful.
Intuitively, we wish to capture how much information about the task labels are captured by the non- uniformity in the data points' orbits. Let c: 2 G be canonicalization function, such that c(x) denotes where on each orbit x information contained in e(x), wish understand the dependence between orientations c(x) and labels fx).
(8 Since data augmentation and invariant featurizations destroy any
We explore two potential task dasendepend metrics. A standard cntemattion-hhoorette quantity for measuring dependence the mutual information, which the KI-divergence hetween the joint distribution and the product of the marginals
MI(c(),,()) : KL (((x))), (())))/)))((()))))))))))) (((()))((())))
Here, x and x' are independent draws from Px- However, the KL divergence inefficient to ap- proximate with finite samples, places stringent requirements on the distributions' supports, and does not capture any notion of ease of learnability or computability. Thus. a potential task- dependent metric is replacing this divergence between product distributions with the classifier dis- (cx))
tance detass (x)x), f(x)), referred to as the musk-dependent detection metric below. In other words, we train a small network to classify whether pairs of group elements and labels are mis- matched. Note that the task-independent and task-dependent metric are not necessarily correlated; one can be high while the other is low, and vice versa, as verified in Section 5
One complication for this metric is that it depends the choice of canonicalization c(-). We assume there some "natural" choice of ((.), i.c. which is easily computable by a neural network, as we care about the implications of distributional
-rreakin downstream learning tasks. To give
8 Formally, all this requires is that c(gx) gc(x). Although such a map is not well-detined for objects x with self- symmetry, we ignore this issue for the sake of exposition. Note that c(-) makes an arbitrary, but hopefully logical (barring unavoidable discontinuity (Dym et al. 2024) choice of which x to assign to the identity element of the group
28

---

Published as a conference paper at ICLR 2026
Original Dataset
Extract c(x)
0.5
197
6r
111
Small, unconstrained
Class Label
Class Label 1: mamasshre (aov
p-1
ruiched pain (ctouoo
Cuadly
3
madel
(15.4)
(190. 60
6F.13
(145.2)
Predict
Figure 7: (Left) The swiss roll dataset (Wang et al. 2024b) provides varying levels of dependence between the canonicalization (with respect to the group of discrete vertical translation), and the task. Augmentation by vertical shifts destroys the useful canonicalization of the = dataset, forcing the network to learn a complicated rather than a simple decision boundary. (Middle) For the classification version of the task-dependent metric, input datapoints assigned orientations by a small equivariant network c(-), and we create pairs of orientations with labels that may or may not match. A binary classifier network then tries distinguish between matched and mismatched (orientation, lahels) pairs, as proxy for how informative the orientation for the task. (Right) For the direct prediction version of the task-dependent metric, which we focus on the main text, input datapoints are also assigned orientations by a small equivariant network c(-). These are used directly to predict the labels, and compared the label prediction of random group element input.
an example of a "bad" choice, imagine a c(-) which discontinuous in x, or which so complex to compute that even if it correlates very well with f(),, it would be difficult for network to compute it as a feature. Therefore, we parametrize c(-) via a small equivariant network, which can be either trained alongside the binary classifier for detass or just randomly initialized.
Another choice for the task-dependent metric the accuracy of predicting fxx directly from c(x) (referred to as the direct tesk-depnndnnt metric (px,y) in our experiments). In Appendix C.2
closely related. This is in turn closely related to the concepts of V-information (Xu & Raginsky (017) and the information bottleneck (via the canonicalization) (Tishby et al.) 1999). We find empirically that the direct task-dependent metric t (Px.r) (results presented in 5.1 1s casier to optimize and the performance is reasonable on test cases.
show that r (px.y) and detuss (())
C.2 THEORY
For simplicity, throughout this section we assume c(-) is not learned, e.g. coming from small cquivariant network with frozen weights as in the experiments.
We introduced the task-dependent metric as a measure of the dependence between the canonicaliza- tion c(x) and the label f(x). Instead of using the mutual information, as below,
M((c = KL ((o(s)) (c())))())))
used the classifier distance (from the mask-independen metric) between the joint distribution and the product of the marginals. Recall the classifier distance:
detass (po.Pi)) = br-BBem Er-pp (NN(xx = =b]
Specializing to our distributions, where po is the joint and pı is the product of marginals: mi (c(()) detass (e))))
- F. [=NNN((())))()) (
1
+
2 F., [I(NN(((s.)
In other words, assess classifier (NN)'s ability distinguish between pairs of canonicalization and label that are matched, vs mismatched. In practice, we of course train NN on a training set, and then approximate this expectation via held-out test set.
29

---

Published as a conference paper at ICLR 2026
However, another natural measure of the dependence between c(x) and f(x) is assess how predic- tive c(x) is of fx), that is: how well can a neural network predict f(x) directly from c(x)? If there is no dependence between them, then it can do no better than random. Letting & be loss function, we define
mz(c()).f()) =- =., [(N - E, (<(NN
Here, the second term captures how well c(x) can be used to predict f(x), while the first term regularizes/caliirates by how well NN performs with independent inputs.
Intuitively, my and my are quite related to each other, and it is a straightforward exercise to make this precise. Letting é be the 0/1 loss (i.e. O if NN' (c(x)) = f(x and otherwise), one can obtain one direction of a bound between my and my by using NN' to define NN. particular, define NN(c.f) == 0 if (C) f, and 1 otherwise. Then,
m2(c()),(()) = E.,< -
= f(x)) - (N-1) = f(x)) = E, ( -x, [((NN(c(()) - f(())
Ex (L)) 0]] -L. (()) =0]
=E, (INN(((((()) =0] - (()
-2m((())(())-
In the other direction, we could start with NN and define (c) to be any f such that NN f) = 0. Therefore, when optimizing independently over NN and NN', m2 is at least 2m; - 1, while at the same time, my is at least " The two quantities are thus related by an affine transformation - at least under a certain choice of loss (and optimization practicalities notwithstanding).
These quantities are also very related to V-information (Xu & Raginsky (017). In particular, when É in the definition of my is the cross-entropy loss, my is essentially the predictive V-information from c(.) to f(-) (Xu & Raginsky 2017). In subsequent experiments, when we report m2 (the "tesk-dependent direct prediction" metric), we report only the latter term E, [[(((NN(c(()))
C.2.1 TASK-DEPENDENT METRICS FOR FINITE AND INFINITE GROUPS
When G is finite and |G| is much smaller than the number of class labels, then it is clear that c(x) can not be expected to predict f(x) perfectly (hence the role of the first term in the expression for m2). In the case of MNIST, for example, the group of 90º rotations has 4 elements. while there are 10 digits to classify (which occur with equal probabilities). Therefore, directly predicting the digit label from c(x) is impossible, regardless of the dataset distribution; one can only associate one label to each of the four elements of c(x). (Indeed, Fano's inequality can provide a lower bound on this probability of error.)
On the other hand, when G is infinite and the dataset is essentially canonicalized according to our main, task-independent metric, an element c(x) E G can information theoretically capture exactly the identity of i.e. such that can be recovered from c(x), for sufficiently expressive choice of c. The extent to which this generalizes will depend on the smoothness of the canonicalization, but not necessarily on how "task-correlated" it is (if NN is powerful enough to simply recover from c(x) and use it to predict fxx)). Thus, considering the version of the task-dependent metric that aims to predict (x) directly from c(x). overfitting is concern, and may explain some of the sensitivity to architecture choice from Table 6 and Table
30

---

Published as a conference paper at ICLR 2026
D EXPERIMENTS
For all experiments, we provide further details on the training, model specifications, and results. To reproduce experiments, we include README. md file in the supplemental codebase, specifying the exact commands for each dataset and setting. We also summarize the task-independent metric in Table As shown, every dataset experiences distributional symmerry--reakinn. varying (but all high) degrees.
Dataset
Test Acc.
rMD17 Aspirin
97. .688
rMD17 Ethanol
79.834
OC20 Surface+ Adsorbate
99.280
OC20 Adsorbate
96.529
QM9
97.6
Local QM9
67.6
QM7b
89.93
MNIST
87.50
ModeINer40
92.45
Local ModelNet40 10
55.6
Local ModelNet40 N = 100
81.5
LLM Materials
95
Table 3: Task-independent metric on selected datasets (omitting the toy Swiss Roll dataset and ModelNet40 reported per class in Figure T0).
Performance relative to FF (with #1 std propagated)
90-
Reletive improvenent
-25
-50-
QM7b
GM76 aisa
QW9 Cv
QM9 [+]
QM9 dz
MNIST (%)
ModelNen40
Model
Equivariant
Group Averaged
Figure 8: To complement Table [ we show the percentage improvement relative to the FF setting for the different augmentation settings (TT, TF, FT), equivariant, and group averaged models. We see that the relative improvement is largest for the vector quantity üü, and there is no improvement for MNI50.
D.1 MODELNET40
D.1.1 CLASSIHCATION RESULTS ON TRANSFORMER MODEL
To show that the results in the main text are not specific to the Graphormer architecture, we also experiments with transformer architecture. We train transformer with the four different augmentation settings and report the test accuracies: TF=76.778% TT=75.723% FT=7.86% and FF=84.49%. Thus, for this dataset, FF setting is better which aligns with the results in the main text.
Figure shows per-class results for m(xx)-
31

---

Published as a conference paper at ICLR 2026
Pitecton Acurace Orribudo
Cataggyy Coue
Acouracy
Figure ModelNet40 m(xx) histogram over classes
D.1.2 RELATION BETWEEN THE DEGREF OF CANONICALIZATION AND ACCURACY ON FF/TF AUGMENTATION SETTINGS PER CLASS
We show the scatterplot of the FF/TF test accuracy vs. the degree of canonicalization per class in Figure IO As above, FF indicates no augmentation at train or test time and TF indicates augmen- tation at train but not test time. We do not notice trend in the relative improvement between FF and TF as a function of task-independent metric, but is interesting to note weak positive trend between the task-independent metric and both test accuracies. It is unclear why this is the case. We hypothesize that perhaps certain classes are defined by simple features, which would then tend to re- sult in both the task-inde inkeidddenndent metric and the test accuracy being higher but further exploration needed to truly explain this phenomenon.
between Independent metric (detection accuracy) Accuracy FF/TF settings
Scal Seal Beffrman
Test Accuracy
independent metric (detection accuracy)
Figure 10: Relation between task independent metric and accuracy on FF/TF augmentation settings.
D.1.3 TASK-DEPENDENT METRIC
We use vector neuron network (Deng al. 2021) for canonicalization. For the direct prediction task-dependent metric (px.y), we use a four-layer MLP.
1.1.4 TRAINING CURVE
We show the training curve ModelNet40 classification task with different augmentation settings in Figure [I] The training curve shows that the FF setting achieves the best performance all the time, while TF and TT settings achieve similar performance, and FT setting achieves the worst performance. It is interesting to contrast this result with QM9, where the best-performing setting the unaugmented test-set is to still augment (TF). This suggests a fundamental difference between ModelNet40 and QM9.
32

---

Published as a conference paper at ICLR 2026
Test Accuracy on each setting
Test Accuracy
Tamm/uu lascrase
Tatn/hve Taa:
epoch
Figure 11: Test accuracy curve of ModelNet40 classification task with different augmentation set- tings.
D.1.5 MODEL DETAILS
For training the eek-inddeenneet metric, we use a four layer transformer architecture with four attention heads and 128 hidden dimensions the backbone. The number of parameters in the model 793k. For the task fask-dependent detection metric and task-dependent direct prediction metric, we use three layer vector neuron network (Deng al. 2021) as the canonicalization network, and four layer MLP as the prediction backbone, respectively. The number of parameters for the vector neuron network is 6.1K, and for the MI.P is 14.2K for direct prediction and 13.8K for detection. To ensure fair model comparison, we use the Graphormer model used in QM9/QM76 for the classification task in the main text [Shi al. (2022). In this setting, each point cloud is uniformly downsampled to 512 points. We run the base Graphormer setup with 4 blocks, transformer layers, 8 heads, an FFN width of 256, and distance encodings using 32 Gaussian kernels. The regularization applied is dropout=0.1 on attention and on the final layer only no dropout is used on inputs or intermediate activations. The number of parameters in this model 822K.
For the equivariant counterpart, we build on e3nner et al. 2022 Kleinhenz & Daigavane). Graph edges are defined via k-nearest neighbors with k 15. Point embeddings are first lifted to the mixed representation irreps hidden = 64x0e + 16x1o. Edge attributes derive from relative offsets within a cutoff distance (max radius = 5.0) and are then expanded in spherical harmonics with irreps sh Ix0e 1xlo to encode angular structure. We stack three equivariant convolutional stages with gated nonlinearities and include linear self-interaction terms. The head performs global node pooling and equivariant MLP, producing outputs with irreps out 40x0e. The number of parameters in this model is 772K.
For the transformer baseline, the architecture eight attention heads and 256 hidden dimensions as the backbone. The number parameters in this model is 4.7M.
is six layer transformer architecture with
1.6 TRAINING DETAILS
The data split for all ModelNet40 experiments is 80/10/10 for training/aalidationnessnnng, The train- ing details of classification experiments in the main text are as follows: Graphormer uses Adam with learning rate le-4 and batch size 16. The e3nn uses Adam 1e-3 with batch size 16. To make the comparison fair, we implement stochastic, group-averaged Graphormer: at each forward pass we sample n = random rotations from S3(3), run the network on each, and average the resulting predictions. For transformer models, trained each setting for 300 epochs with batch size 128.
We train the task- indepeddeenn metric for 30 epochs, the deskndeeend detection metric for 1200 epochs and the direct prediction task-dependent metric for 300 epochs. All models trained on one NVIDIA GeForce RTX 4090. The training time for the musk-independent metric was about 20 minutes, for the task-dependent detection metric was about 9.5 hours, and the direct prediction
33

---

Published as a conference paper at ICLR 2026
deseddepnndent metric about 8.5 hours. The classification task took about two hours and half for each setting for transformer model. For graphormer model, it will take about 8 hours for each setting. For the equivariant model, it will take about 16 hours for each setting.
34

---

Published as a conference paper at ICLR 2026
D.2 MNIST
D.2.1 TRAINING AND MODEL DETAILS
All experiments were run on single NVIDIA RTX A5000 with batch size 128, 50 epochs, and standard 3e-4 learning rate for the Adam optimizer, which took roughly 30 minutes each. For the data splits, we split the original training set of 60k images, and split into 60% 60%/209 /00% /00 for train, validation, and test. MNIST training runs with/without augmentation used as a base network a basic 421k-parame -aaraaaaeere CNN with two convolutional layers, followed by two-layer MLP. For the iask-indeeendent metric, the training hyperparameters and model architecture were the same, with only the final number of model outputs modified from to 2. The group-averaged model used this base architecture taking the average over the C4 group at each forward pass. Note this is equivalent to an equivariant model over a discrete group.
For the task-dependent detection metric, we parameter 90º rooation equivariant classifier outputting a four-dimensional vector, corresponding to the four elements of G (90trrotaioons). (The classifier essentially applies 2-layer CNN to all four rotations of the input image.) To obtain a single element of G from this vector, we simply apply softmax with low temperature (1e-3), effectively setting it to be one-hot at the index of maximum value. For the network that predicts a binary class from pairs fx)), of these inputs into 4-layer Saraameeee MLP, with 64 hidden features per layer.
the canonicalization network c(.) 19k-
), we simply concatenate all
D.2.2 TASK-INDEPENDENT AND -EEPENDENT METRICS
lask iel-inddeenndent Classification Metric
2.5
0.2
5.4
0.5
3.0
Transtorm
Figure 12: MNIST task-independee metric as function of proportion randomly rotated.
In Figure 2
we report m(px) as a function of what fraction of the dataset was randomly rotated, and recover the predicted optimal accuracies.
we report the task-dependent metrics. They all performed somewhat poorly, possibly as result of on how much information can be encoded in a canonicalization with respect to a group of only four elements (see Appendix C.2). or possibly because the orientation not practically that informative for most of the digits (despite the 6/9 toy example), which is consistent with the fact that all augmentation settings (outside of "FT") worked very (and nearly equally) well.
In Figure 13 and Figure [4]
D.2.3
LOss CURVES
See Figure 15
We can also analyze the test accuracy per class as in Figure 16
D.2.4 SELECTIVE AUGMENTATION USING THE TRAINED CLASSIFIER
We have focused on evaluating whether or not dataset with equivariant labels has distributional symmetry breaking, and what the implications are on whether or not to employ group augmentation. However, augmenting versus not augmenting does not have to be a strict dichotomy - one can apply restricted set of augmentations, or only certain datapoints. Indeed, this is the approach taken by methods which learn how to augment, including e.g. Augerino (Benton et al. 2020b). We now explore how might selectively augment using trained classifier, interpolating between full data augmentation (when there is no distributional smmmetry-breaking and no data augmentation (when there is significant distributional symmetry-breaking
35

---

Published as a conference paper at ICLR 2026
Task dependent Detection Metric
untrained
50.15
50.10
Test Accuracy
50.05
$0.00
49.90
49.80
0.0
0.2
0.4
0.6
0.8
1.0
Transform
Figure 13: MNIST task desendepndd detection metrics.
Task-dependent Prediction Metric
untrained
0.1650
0.1645
Test Accuracy
0.1660
0.1635
Transform
Figure 14: MNIST task-dependent prediction metrics.
Train/Tes /lest Accuracy Across Runs
100
96
Accuracy
92
Grounfuv (trein)
GrousAve (test)
(brin)
[tesi)
TT Itrin))
(test)
86
(train)
TF [tess)
40
Steg
Accuracy over the course of training for the MNIST classification task, in different augmen- tation settings ("TF" augmentation for training, no augmentation for testing, etc).
Figure 15: MNIST loss curves. Dashed lines indicate losses, while solid lines indicate train losses. We omit the FT setting as the test accuracy was around 40%.
36

---

Published as a conference paper at ICLR 2026
Class
Tlat: Accurasy
Tlc: Accurncy
70
Step Gooz
Seep Class
080
Accurocy
Pezt Nicurnccy
00
40
10 30 50
30
20 35
50
Figure 16: MNIST loss curves per selected classes. It is interesting to note the variability in the FT setting, which corresponds to whether test-time augmentation destroys useful information not. For example, FT much worse than the other settings for 6/9. However, for more symmetric shapes like 0/8, FT performs better.
As first step towards this goal, we use the binary classifier (that was trained to compute m(px' )) as part of a mechanism for deciding how to augment an inpul datapoint. For now, we use simple heuristic proof of concept, but more advanced methods of harnessing the classifier should be possible. In particular, we threshold at 0.45 the probability the trained classifier placed on each rotation of the input being from the original (non-rotated) half of the dataset. We then only augment from the set of per-datapoint elements which pass this threshold. The threshold a hyperparameter, but our rationale for choosing 0.45 based on the behavior of an optimal classifier on a simplified dataset with finite orbits, as evaluated at the end of Section 2 In particular, for an input x, let o= (gx: g E G} be the orbit of x, and let S C 6g denote the subset of the orbit of on which the natural data distribution places probability mass. For simplicity, we assume each element of S equally likely under the data distribution. The log odds between the original dataset and the rotated dataset, for each datapoint, is T/7,
1/m a : r. Setting Th, we obtain T Thus,
for inputs x in the support of the data distribution, the optimal can range between when all elements of the orbit appear under the data distribution (no distributional symmetry breaking) and 10.1 when only one element of the orbit is natural to the data distribution. Thus, we threshold [0.++1
near 0.5, augmenting only on inputs that the classifier believes are likely to fall in in-distribution. In an ideal world, this method should automatically interpolate between "no augmentation" and "full augmentation": if the classifier is always 50/50, then we always augment, and if the classifier is very confident that datapoint is or is not natural, we do or do not augment accordingly.
In Figure 17 try this out with three variants of MNIST: with the original dataset ("None"), rotated version ("All"), and a partially rotated version with only certain digits rotated ("345"). To make the task harder and therefore benefitting from augmentation, 1,000 images to use. The selective augmentation method described above automatically interpolates between augmentation and full augmentation, performing near the best of the two in each case. These are promising first results, and can be investigated further in future work.
randomly select subset of
37

---

Published as a conference paper at ICLR 2026
Subset Size: 1k
0.95
Test Classification
0.85
0.80
Augmentation Method
None
0.75
Always Rotate
Selectively Rotate
0.70
All
345
None
Rotation Subset
Figure 17: Classification results with CNN using selective augmentation, averaged over trials. In each trial, new data split was used and new detection network was trained.
38

---

Published as a conference paper at ICLR 2026
D.3 SWISS ROLL
As shown in Figure the swiss roll dataset consists of two interleaved spirals (Wang et al. 2024b). The spirals have distinct values, they are easily separable by horizontal plane. However, there also a more complex function fitting the data that is invariant to z-shifts (the group Z2). Following (Wang et al. 2024b), we create continuous family of datasets in which only p-fraction of one spiral are separated vertically. This creates a spectrum of tasks, where p = 1 is canonicalized in task-useful way, whereas p 0 is not. We find that augmentation of simple MLP indeed hurts performance on this task, with the effect increasing along with p (??). This is captured by the task- dependent metrics, which increases along with p. However, the task-independent metric cannot capture the dataset canonicalization, as this would nearly require solving the hard spiral task itself!
To elaborate, we can think of the p= distribution as a perfectly canonicalized dataset. The reason that task-dependent metric does not pick up this, and instead has only 50% accuracy, is essentially that the canonicalization was not simple in fact, it effectively solved the prediction task (i.e. given point (x, ). 2 was set based on spiral - class(x.y)). So, it is hard for small network to detect on its own whether an input is canonicalized or not.
When p (see Figure 7) the classifier knows that any datapoint with 1 came from the trans- formed distribution, since the original p = 0 distribution always has Z = 0. The classifier can guess the label corresponding the original dataset when z O, and this will achieve 75% accuracy as shown. Moving from 0 to why the task-independent metric drops. In a sense, as p goes from 0 to we continuously trans- form from a task-useless canonicalization to a task-useful canonicalization, which is reflected more accurately by the task- deseddepnndent metries (see Figure 18)
simply interpolates between these two scenarios, and this is
D.3.1 TRAINING AND MODEL DETAILS
All experiments were run a single NVIDIA RTX A5000. In work, we modified the swiss rolls of their extrinsic equivariance setting by adding hyperparameter to denote how much of the data on top spirals randomly chosen be flipped the bottom spiral.
All experiments were trained with the Adam optimizer with learning rate 3c-4, batch size 100, and for 150 epochs. The model is 3-layer 67k-parameter MLP for the original classification task and for the detection task; for the lask-dependent metrics, a Cz-equivariant network with 3k parameters used to canonicalize, composed with 4-layer 13k- 3akparrmeterr MLP to perform the final prediction. The dataset consists of 1, .000 examples, split randomly as 60%//0%//20% crain/vaidatiiontess
D.3.2 RESULTS
Tan Aug RattAg
Matri
fuse TusFales
Depaidert idrect
Teal Depandent (dttect
Thet 02
Deprie of Tase Depandor: Caroniadicaaann.p
Degroe of Tape Taae--aaaaunntt Dopunden: Caronizalzzaann
Figure 18: Swiss roll augmentation performance and metrics as functions of dataset canonicaliza- tion.
D.3.3 LOSS CURVES
See Figure 19 and Figure 20
39

---

Published as a conference paper at ICLR 2026
(a) p= 0
(b) p=
(c) p 1.0
Figure 19: Accuracy over the course of training for the swiss roll classification task, in different augmentation settings ("TF" = augmentation for training, no augmentation for testing, etc). Dashed lines indicates test losses, while normal lines indicate train losses.
p-0
p-0.5
Mr
70
85
40
80
100
120
140
Epoch
Figure 20: Accuracy over the course of training for the task-independent detection metric on the swiss roll dataset, at different levels of task-comreatednnes (p). Dashed lines indicates test losses, while normal lines indicate train losses.
40

---

Published as a conference paper at ICLR 2026
D.4 QM9
We describe setup details, as well as further experiments on (1) the task-dependent metric, (2) the symmetry-breaking of local motifs, (3) interpreting the learned classifier.
D.4.1 DATASET, MODELS, AND TRAINING DETAILS
We obtain the QM9 dataset (Ramakrishman et al. 2014) from https://doi.ozz 6084/m9.. fi gagggshare... 978904.9.... v5 The original dataset has 133885 molecules, 3054 of which are uncharacterized, as found in madonnoodde:Eilees335
figshare.. com/ 404 The uncharacterized molecules are removed during prepro- cessing. We follow (Anderson et al. 2019) and split the dataset into training/validaaionnnsst partitions consisting of 100k, 18k, and 13k molecule examples.
For the task-independent metric, train a generic transformer architecture with 812k parameters for 20 epochs and the Adam optimizer at learning rate le-5 and batch size 128. For the task- dependent metries, used 3-layer, 28.5k parameter e3nn canonicalization network, using Gram- Schmidt orthogonalization to turn the 2x1e outputs into a proper rotation matrix, and a basic 4- layer, 13.8k parameter MLP for the final task-dependent predictions. We used the Adam optimizer with learning rate 3e-4, 50 epochs, and batch size 128.
For the regression tasks, as we are studying the need for data augmentation, the principal model used should be seleguvvariinn/nonivvvairi
We note that many of the recent top performing models on QM9 are equivariant or invariant, so we use a slightly older Graphormer [Shi et al. (2022); Ying iet al. (2021) architecture from 2021. We include an embedding depending on the position of each atom (not solely the relative position) so the model is not invariant. Each node in the graph thus has scalar feature (the atom number) and 3D position associated with it. We use an embedding dimension dembed = 128 for both the atom positions (embedded with a learnable linear layer) and for the atom types. The edges between atoms are encoded using set of learned Gaussian radial basis functions. We adopt the following parameters of the Graphormer base architecture: blocks, 8 layers, 32 attention heads, feedforward dimension of 128, and 32 Gaussian kernels for distance encoding. Regularization uses a dropout rate of 0.1 for both attention and final layer dropout, with input or activation dropout. We train separate model for each property with different data augmentation settings (TT = train/test augmented, FF = none, TF = train-only. FT = test-only). For training Graphormer, we use the Adam optimizer with a learning rate of 5.
For QM9 property regression, we compare to a simple equivariant convolutional neural network ar- chitecture using e3nn(Geiger et al. 2022 Kleinhenz & Daigavane), as by equivariance, predictions should not change whether train/tes, faun/testaan are augmented or not. The network uses a learnable embedding (embedding dimension = 32) for atomic species and lifts the atom embeddings into a mixed repre- sentation i irreps hidden = 64x0e between atoms within a cutoff radius (max radius to spherical harmonics transforming as irreps sh = 1 1x0e + 1x1o, capturing the angular de- pendence. Radial dependence is captured via Gaussian radial basis functions applied to interatomic distances. We then use 3 layers of equivariant convolutions with gated non-linearities and linear reff-interactions. The final layer pools over nodes and uses an equivariant MLP to return the final output as irreps out 1x0e (a scalar for example for predicting one of the QM9 properties). The E3ConvNet model is trained with the Adam optimizer and a learning rate of le 14-4. 4. For an apples to apples comparison, we implement a stochastic group-averaged variant of Graphormer, in which n 5 random rotations of the input are sampled from SO(3) at each forward pass, and the corresponding outputs are averaged to produce the final prediction. While neither of these archi- tectures are near state-of-the art for QM9, for our studies it suffices to use smaller models (each with approximately 800k parameters) to understand how augmenting impacts results. For both the e3nn model and the Graphormer model with augmentation settings for each property we train each model for 150-200 epochs (depending on property) on a NVIDIA RTX A5000, which takes 2-3 hours. Minimum test MAE values and test MAE curves are reported in Table 4 and Figure 22
16xlo. Edge features are computed through relative 5. 5). These features are then projected
D.4.2 TASK DEPENDENT METRIC
We describe further details of the task-dependent metric for QM9. Both the detection metric detow and the direct task-dependent metric ((xx,)) from the main text were trained for QM9 properties.
41

---

Published as a conference paper at ICLR 2026
Table 4: MAE on the QM9 dataset for Graphormer under different data augmentation settings (TT train/test augmented, FF none, TF train-only, FT = test-only). We include an e3nn convolutional neural network model with similar number of parameters for comparison. The best- rest-performing model is in bold, and the best performing-man monell within the augmentation settings underlined. Results are averaged across random seeds and + I standard deviation is reported.
Target
Unit
TT
E3NN
Group Avg
и
0.243=n.mm
0.275=n.010
0.245
0.07
0.131=0a00
0.217
a
0.480a0.115
2.482
0.482200009
200
0.3033
0.41
HOMO
1.095=a.000
0. 0.108
0.095=0.001
0.165
0.099
0.09442.08
LUMO
eV
g0.20
0.1 10.00
0.121
10.002
0.2120
0. 0.10220.004
0.11010000
g0:000
0. g00006
0.1721 00.00s
0.278
20.00
0. 10
R²
4.7gm.2
4.46 2.46
4.4
10.48
4.6ga.1
3. 3.0
ZPVE
eV
0.0123
0.01 100
0.0121
0.01 10.0003
0.0091
0.0080.20.007
Up
eV
6.620.5
5.612.1
6.740.
6.7.29
12a5
4.0
U
eV
6.1.
6.3114
6.1
8.3
11
5.8120
H
eV
6.1.n.
5.2g2.1
6.2g0.8
8.9.22
11
5.6
G
eV
:0:5
5.5
7.03
8.2.2.
12.6:4.4
5.6
Cv
cal/mol K
10.005
0.143
0.147
0.21
0.11010.005
0.128
proparty)
0.5
650
1000
1240
50
-90
8
Figure 21: Bar plot visualization of Table 4 where the scale is normalized by the maximum MAE per-property across all six methods compared.
42

---

Published as a conference paper at ICLR 2026
Test MAE Curves Across Properties (Graphormer, E3NN, GroupAve)
0.6
Test MAVE
0.75
Test MAE
Test MAE
1.5
0.50
1.0
0.2
0.4
0.25
0.5
0.1
0.2
50
100
50
100
50
100
50
100
Epoch
Epoch
Epoch
Epoch
150
Test MAE
0.6
Test MAE
60
0,075
Testt MAE
100
40
0.050
0.4
50-
20
0.025
0.2
0
0
50
100
0
50
100
0
50
100
50
100
Epoch
Epach
Epach
Epoch
150
Test MAE
100
Test MAE
Test MAE
MAE
1.0
100
100-
50
0.5
01
0
50
100
50
100
50
100
50
100
Epoch
Epach
Epach
Epoch
groupave
Figure 22: Test MAE per epoch for QM9 predicted properties using the Graphormer- architec- ture. We show the augmentation types train/test aug=TT, train aug only=TF, test aug only=FT, no aug=FF. We also show an E3NN network with a similar number of parameters and a group-averaged graph transformer.
We observe that both task-dependent metrics have small signals (see Figure 23 This aligns with our observation that equivariant models/data augmentation perform well for QM9 as there may be lack of significantly task relevant information contained in the QM9 canonicalization.
D.4.3 INTERPRETABILITY OF CLASSIFIER FOR DISTRIBUTIONAL SYMMETRY BREAKING
A primary motivator for using the classifier distance for distributional asymmetry detection is be- cause of the opportunity to explore and interpret the trained classifier. As a first step, focus on the task insk-independent classifier trained on the QM9 dataset with Anderson splits as outlined in the previous section. To probe the decision boundary, we evaluate the classifier predictions set with no augmented rotations (e.g. all have label and are from the original dataset). It is thus easy to interpret which molecules are "hard" for the classifier to distinguish as being from the orig- inal dataset. We apply PCA to the learned embeddings (i.c., the layer immediately preceding the final output layer) and visualize them in Figure 24 revealing that the misclassified examples tend to cluster together in PCA space.
the test
We evaluate the sigmoid of the classifiers logits on a discrete grid of 3D rotations (representing the probability that the given sample has label 0 or is from the original dataset rather than the augmented version). In order to visualize the probabilities over SO(3), rotations with non-negligible probability are plotted as dots using Mollweide projection (Murphy et al. 2022, Klee et al. (023), with rota- tions orthogonal to the sphere encoded as colors and the size of the dot representing the magnitude of the probability. We explore correctly classified molecules, incorrectly classified molecules, and samples that are close to the decision boundary and show examples of each. We also investigate the stability and robustness of the classifier's 's decision boundary by identifying rotations of given sample that lead to a change in its predicted label. To probe the stability of the classifier, we identify pairs of rotations that are close together yet lead to large changes in the classifier's output logits. Given two rotations represented by quaternions p,4, the distance between rotations is
0 = 2 aarcoos| p.9>
(64)
43

---

Published as a conference paper at ICLR 2026
10.
FF
10"
crand
mg
Cg
MAE
20.2
10-1
homg
I
10-2
53
54 55
56 57
58
Task-Dependent Detection Accuracy
(a) dask-dependent detection metric using trained vs, untrained small equivariant network for c.
(b) Normalized MAE for QM9 vs the task- dependent detection metric, with untrained.
(c) MAE FF/TF improvement vs QM9 direct task-dependent metric.
Figure 23: QM9 daseddennddn detection metric figures with MAE table and plots.
3D PCA of Embeddings Colored by Predicted Probability
Prob(Class
0.9
0.8
PC3 22
0.5
0.4
-8
0.1
PC2
PCI
Figure 24: PCA of learned embeddings for the task-independent classifer for QM9 applied to the test dataset with no rotations. The misclassified examples thus are shown red.
44

---

Published as a conference paper at ICLR 2026
Example Correctly Classified with High Probability. We select sample classified correctly with high probability as being from the original dataset and investigate the classifier outputs per rotation angle, as shown in Figure 25
Original: p(from
Rotated: pffrom dataset) 0.99
C2H4O, gdb_11
dataset) 0.94
P(from dataset) for differing rotation angles
4900.
Rotated: p(from dataset) =
180
270
90
Probabilities per rotation angle for sample that is classified correctly. The colors correspond to rotations orthogonal the sphere and the size of the dots corresponds to the probability value. The original molecule in the dataset is shown on the upper left. The lower left shows rotations that cause prob (original) to be zero. The right shows rotations that cause prob (origine) to be high.
Inetabiiiles: Large Aonah for Small Angle
Trequaney
p(from dataset) 0.03 to 0.95 with 0.05 radian angle between rotations
19
9.5
7.5
10.0
Instability in the decision boundary (left): two nearby rotations cause large change in predicted probability. Histogram (right): certain examples exhibit such instabilities more frequently.
Figure 25: Visualizations of classifier outputs for example classified incorrectly.
Example Incorrectly Classified. We select a sample classified incorrectly with low probability as being from the original dataset and investigate the classifier outputs per rotation angle, as shown in Figure 26
Example Close to the Decision Boundary. augmented test set that the model assigns 50% probability of belonging to the true dataset (cor- rectly classified but close to the decision boundary).
We select also select example from the non-
interesting to note that for each example, we find instabilities in the decision boundary (rota- tions that are very close together but correspond very large changes in the classifier output). This demonstrates that our method could perhaps be used probe the instabilities of given canonical- ization know that each canonicalization has such instabilities (Dym al. 2024). although
45

---

Published as a conference paper at ICLR 2026
Original: p(from dataset) 0.008
Rotated: p(from
C5H7N30, gdb _9421
dataset) = 0.99
P(from dataset) for differing rotation angles
Rotated: p(from dataset) =
270
Tit 90
Probabilities per rotation angle for a sample that is classified incorrectly. The colors correspond to rotations orthogonal to the sphere and the size of the dots corresponds to the probability value. The original molecule in the dataset shown on the upper left. The lower left shows rotations that cause prob (orig high, with arrows pointing to corresponding (starred) points on the Mollweide projection plot.
to be zero. The right shows rotations that cause prob (origing miriinall to be
instabilties: Large sprob for Small Angle
10"
303
20F
10
p(from dataset) = 0.07 to 0.76 with 0.05 radian angle between rotations
rangie
Instability in the decision boundary (left): two nearby rotations cause large change in predicted probability. Histogram (right): certain examples exhibit such instabilities more frequently.
Figure 26: Visualizations of classifier outputs for example classified incorrectly.
46

---

Published as a conference paper at ICLR 2026
Original:
p(from dataset) 0.5
Rotated: p(from
C5H10N202, gdb_56276
dataset) 0.99
P(from dataset) for differing rotation angles
Rotated:
p(from dataset) 0
180
270 Til 90
Probabilities per rotation angle for sample on the decision boundary. The colors correspond to rotations orthogonal to the sphere and the size of the dots corresponds to the probability value. The original molecule in the dataset shown on the upper left. The lower left shows rotations that cause prob (original) to be zero. The right shows rotations that cause prob (original) to be high, with arrows pointing to corresponding (starred) points on the Mollweide projection plot.
Instabilties; Large Aprob Small Angle
Trequency
p(from dataset) 0.07 to 0.734 with 0.05
between rotations
radian angle between rotations
Instability in the decision boundary (left): two nearby rotations cause large change in predicted probability. Histogram (right): certain examples exhibit such instabilities more frequently.
Figure 27: Visualizations of classifier outputs for example close to the decision boundary.
47

---

Published as a conference paper at ICLR 2026
Architecture
Architecture
Transformer Architecture
80
LR
Test Accuracy
0.0001
0.0003
70
0.001
large
medium
50
small
10000
20000 30000
10000 20000 Subset Size
30000
10000
20000
30000
Subset Size
Subset Size
Figure 28:
Ablations of m(px) on MNIST with respect architecture type, architecture size, learn- ing rate, and dataset size. Curves are averaged over two random seeds.
cannot guarantee that the instabilities the model's predictions arise for this reason (and not e.g. due to a failure to learn). Nonetheless, the models probed achieved very high test accuracy, lending confidence that the identified example instabilities are genuine.
D.5 ABLATIONS
As shown in Table [] the task iedeenndeenddnt metric on QM9 robust the choice of architecture.
Table The classifier network is a transformer architecture, in which we vary the depth, number of heads, and hidden dimension. The task-independent metric is robust with respect to the architecture size.
Setting tiny
Depth
Heads
Hidden Dimension
Test Accuracy
Parameters
64
98.3
110,000
small hidden
64
98.6
210,000
large hidden
256
98.4
3.2e6
many_ heads
4
8
128
98.5
810,000
shallow
128
98.3
420,000
micro
32
98.0
27,000
few Leadd
4
128
98.7
810,000
deep
128
98.7
1.6e6
D.5.1 HYPERPARAMETER ABLATIONS
We further ablate sample complexity, together with architecture and learning rate, for MNIST in Figure 28]and QM9 in Figure 29 To vary sample complexity (on the x-axes), we randomly select subset the original dataset the specified size. As shown, there little sensitivity these except when using a learning rate that is too large for the transformer architecture (but this is easily detectable from the loss curves). The greatest variation is visible at lowest-sample regimes, but the classifier performance in all cases is still significantly better than random, which would be 50% accuracy (and recall that the optimal accuracy for MNIST is 85.7%, since the group of 90" rotations only has four elements). Overall. the metric is remarkably sample-efficient coumetry-breaking is quite pronounced, shown by the high accuracies, m(px), ), for each. (In the language of Appendix B of the hypothesis class/network family as being relatively simple; thus, the networks can be trained with even small fractions of each dataset.)
factors,
these two
datasets. This matches our intuition that the distributional
think
D.5.2 DIRECT TASK- SER--DRPEDDNNT METRIC ABLATIONS
As motivation for using the direct desk-depnndent metric (px.y) instead of the task-dependent de- tection metric d we note that d is overly sensitive to the choice of canonicalization and classifier network architectures (see Table 6 and Table 7). In contrast, (px,r) is relatively robust
48

---

Published as a conference paper at ICLR 2026
6891
2000 12300 MCDI
10N6 4000 BCDO
2000 2000 M0000000
X00 2200 ROC--0
Ablations of m(px) on QM9 with respect to architecture size, learning rate, and dataset size. Curves are averaged over two random seeds.
Figure 29:
to architecture choice of both the canonicalization network and the classifier network: shown in Figure 30 and Figure 31 the value stays close to and does not predictably increase or decrease as function of network hyperparameters.
Table 6: Ablations the canonicalization network architecture for the task-dependent detection metric (not (()x,))), averaged over 2 independent runs on QM9. As shown, there indeed varia- tion in the task-dependent metric with the architecture of the canonicalization network. Although we expect some variation-since the task-dependent metric is supposed to pick on a "simple" canonicalization,
did intend restrict the maximum size of the network-the ablations below surprisingly demonstrate that the highest accuracies were achieved by the smallest networks. We note that the loss curves were fairly unstable, possibly pointing optimization difficulties with ten- sor product equivariant networks that might be alleviated for smaller/shallooer networks. For fair comparison practice, one should fix an architecture size, and only compare accuracies computed with the same architecture.
Layers
Irrep Dimension
Test Accuracy
Parameters
acc mean
acc std
param float
4
16x0e 4xle
3
38,000
72
3.0
38,000
2
16x0e 4xle
69 = 0,9
19,000
69
0.9
19,000
2
32x0e 8xle
72
74,000
72
2.0
74,000
4
32x0e + 8xle
76
150,000
76
2.0
150,000
3
32x0e 8xle
74
110,000
74
1.0
110,000
3
16x0e + 4xle
29,000
86
2.0
29,000
2
8x0e 2xle
89 0.1
5,400
89
0.1
5,400
4 8x0e + 2xle
89 0.4
10,000
89
0.4
10,000
3 8x0e + 2xle
22
7,900
88
2.0
7,900
Table Ablations on the classifier network architecture for task-dependent detection metric (not (). The classifier network an MLP, for which vary the number of layers and the hidden dimension.
Depth
Hidden Dimension
Test Accuracy
Parameters
32
68.2
3,800
4
128
88.1
52,000
2
64
88.2
5,200
64
70.0
14,000
8
64
87.6
31,000
2
32
74.9
1,600
49

---

Published as a conference paper at ICLR 2026
A010t on Caroniaa lastion lastion Newrr Archiecture
Ticle 3ds
Natuor Layers
Figure 30: Ablating canonicalization network architecture when computing r(xx,) for the Up value of QM9. The variation ratio between losses low.
Abration
VEP Nerwork Archtecture
1,32
offect (ulbe
2.06
1.000
Figure 31: Ablating classifier (MLP) network architecture when computing t(px.y) for the Up value of QM9, The variation ratio between losses remains low.
50

---

Published as a conference paper at ICLR 2026
D.6 LOCAL QM9
We have shown that the QM9 dataset is highly canonicalized, yet data augmentation and equivariant methods both perform well even on the original, canonicalized test set. This behavior is distinct from ModelNet40, where train-time augmentation impedes performance on the (also canonicalized) test set. This poses a question: why are equivariant methods so helpful for QM9, even though it's already canonicalized? We explore the hypothesis that locality is an important factor impacting performance (which would be captured by equivariant methods, but not canonicalization) In particular, we seek to understand whether local graph motifs in QM9 are less canonicalized - i.e. more likely to appear in a variety of rotations - than the full molecules. If true, then augmentation and equivariant architectures might both benefit from exposing the network to full group orbits of local motifs.
Concretely, our question is: do the local motifs present in QM9 graphs experience distributional symmetry breaking? To address this, we create new dataset from the original QM9 dataset by randomly selecting three nodes from each molecule, and creating a new molecule fragment out of only each node and its neighborhood (as determined by its edges/bonds). As shown in Figure 32 this often includes repeated neighborhoods. This creates a dataset of size 392k. We first simply apply the maskiindependent detection metric, asking network to distinguish between rotated and unrotated motifs. (All experimental and model details are preserved from the ordinary QM9 setting). As shown in Figure(5] the local dataset has lower accuracy than the QM9 dataset. However, this does not provide maximally fine-grained distinction between different kinds of distributional symmetry breaking. For example, suppose a molecule always appears in one of two possible canonicalizations. With an infinite group like SO(3), this detection problem is still likely to be perfectly solvable, as two orientations are still only a measure zero set (3))). Yet, this case distinct from the perfectly canonicalized case.
To assess whether a dataset is truly canonicalized, we train a network to predict g from gx, where g drawn randomly from the Haar measure. Solving this task to high accuracy is only possible when the distribution is truly canonicalized (only one element per orbit appears). We use the same transformer architecture to output 9 values the entries of rotation matrix, and trained it accord- ing to the MSE. (Neither backpropagating through a Seham- procedure to make it a proper rotation matrix, nor training equivariant architecture, rotation error instead of the MSE, were as effective as this simple method, which also circumvents that would be required output a group element on symmetrie inputs (Smidt let al. 2021).) .) As shown in Table 8 there is a discrepancy between the best test accuracy achieved on the original QM9 dataset, and that achieved on the local neighborhood version. Therefore, it ap- pears that the original QM9 dataset is more canonicalized, whereas the local motifs presented in the QM9 dataset can appear in several orientations (although still far from uniform over SO()). This provides some evidence for the hypothesis that methods which involve equivariance to local motifs- including data augmentation and equivariance, but not canonicalization may be providing an additional advantage QM9.
backpropagating through the angle of
the symmetry-breakiin
Consistent with Figure 33 also took much longer to train these models (500 epochs took hours on the original QM9 dataset, and nearly 15 hours the local QM9 dataset, likely due to slower dataloading), which contrasts with the efficient convergence (around 30 minutes) our main task- independent detection metric.
QM9
Local QM9
Test Error (degrees)
13.5
53.7
Table 8: Predicting g from gx.
51

---

Published as a conference paper at ICLR 2026
8
8
8
8
Figure 32: Local QM9 dataset visualization.
Local GMS (trai)
Local 0MB (test)
Original OMD (trant) Original 0M9 (tea))
Figure 33: QM9 and local QM9 dataset training curves for predicting from gx.
52

---

Published as a conference paper at ICLR 2026
D.7 QM7B
D.7.1 DATASET DETAILS
In (Yang et al. 2019), density functional theory (DFT) and linear-response coupled-cluster theory
including single and double excitations (LR-CCSD) is used to compute vector and tensorial molec-
ular response properties for the 7,211 molecules in the QM7b database (Blum & Reymond] 2009
Montavon et al. 2013). LR-CCSD generally more computationally expensive (scaling O(N
with the number electrons No) yet more accurate than DFT, which scales as O(2)). As the QM7b dataset is composed of small molecules, computing material response properties with LR-CCSD is feasible. Quantities computed include the dipole vector Ji, polarizability a, and quadrupole moment Q. The molecular dipole polarizability a describes the tendency of molecule to form an induced dipole moment in the presence of an external electric field (Yang et al., (019). can be computed by taking the second derivative of the electronic energy U with respect to an applied electric field E: a²U
04j
(65)
dE Scalar polarizability response quantities are the isotropic polarizability asso and the anistropic polar- izability aaniso
0sso
5 + Dsy + Oze))
(66)
v2 (asu - dgy) 2 (dhy 0nz)² -kyy)
Caniso
+ +(dzz axx) 2 + 6 dny + ob + +GGz)) 1/2
(67)
The dipole moment the first derivative:
aU
j
(68)
DE
The quadrupole moment Q is a rank-2 tensor that characterizes the second-order spatial distribu- tion of the molecular charge density, capturing deviations from spherical symmetry and providing information about the shape and anisotropy of the electron cloud beyond the dipole approximation:
Qij Σ qa (3ruiaj--8pp
(69)
qa is the charge of particle a, fai is its i-th coordinate operator relative to the molecular center of mass, and ,, is the Kronecker delta.
www.mmatiii... Data can be downloaded from httts:://rcch: lttp://rrcciee.ttttrrllllll//// swwaterriass.ood..oo//reeeooo 2019.0002/10
For our studies, we use the most accurate level of theory available in the dataset respoose coupled cluster with single and double excitations (LR-CCSD)-2 in com- bination with the d-aug-cc-pVDD (daDZ) basis set, to reduce basis set incompleteness error (Yang al., 2019) (specifically, the file CCSD daDZ.tar. available the link above). The data is then converted from XYZ format into a torch.geometr... dataset.
D.7.2 MODEL AND TRAINING DETAILS
For the task-independent metric, we use the same generic transformer used for QM9 to find that the dataset is canonicalized. We train for 100 epochs with a batch size of 128 and a learning rate of le-5 with the Adam optimizer. From (Yang al. 2019), this is expected the molecules were reordered using a hemel--ased similarity measure from (Bartók et al., 2013). For the task-dependent metric, we use e untrained, as we found that using trained allowed the network to learn the dipole vector itself. We use the same parameters as for the task-dependent metric for QM9.
For the regression tasks, we use the same graph transformer architecture as described in Ap- pendix DD4 and compare to the same E(3) -qqiivvauva neural network (now with vector or f= 2 output rather than a scalar as in QM99/rrop--vveneee each for 500 epochs with batch size of 128 single NVIDIA RTX A5000. The Net model is trained with a learning rate of le-4 and the Graphormer model trained with a learning rate of 3e-5, both with the Adam optimizer. As anticipated, the E33) eqquivarra model achieves better performance than the Graphormer in predicting dipole moments, owing to its physically consistent treatment of vector-valued (non-scalar) quantities.
network with sampled rotations. We train
53

---

Published as a conference paper at ICLR 2026
D.7.3 TASK-RELEVANT CANONICALIZATION
To investigate the impacts of task-relevant canonicalization, we run further experiments on the QM7b dataset. Consider aligning molecules such that their dipole moments coincide with the z- axis, filtering for molecules with non-zero dipole moments. This canonicalization clearly makes easier for non-equivariant model to solve the task, whereas an equivariant model will be unable to use this information. We test different data augmentation settings to illustrate the impacts of the task-useful canonicalization (train/test aug=TT, train aug only=TF, test aug only=FT, no aug=FF). Values reported the table below are the MAE across the dipole vector components atomic units (a.u). For the task-useful canonicalization, the FF setting (train/test fully canonicalized by dipole, no augmentation) outperforms the equivariant model (shown bold). In the original dataset without canonicalizing based the dipole, FF does not outperform the equivariant model. These results provide an interesting avenue for future work/for testing the deseedeenndnnt metric.
Table 9: Dipole prediction MAE for different datasets and models. Lower values are better. For the task-relevant dipole canonicalization, only FF is reported; other augmentation settings are changed, Equivariant model and group averaging is also unaffected by canonicalization. Note the FF setting with task-relevant canonicalization outperforms equivariant methods.
Dataset
FF (Dipole Canon)
c3nn
Group Avg 5 rot
Dipole Canon
0.037 - 0.003
0.041 + 0,002
0.043 1 0.002 0.043 = 0.002
Original Dataset
0.100 = 0.001
0.041 = 0.002
D.7.4 Loss CURVES
Test MAE Across Runs Alpha Isotropic QM7
Test MAE Mean Component Across Runs Dipole QM7
10
Test MAE (mean component)
TF
0.30
0.25
Grouptiee
Test MAE
GroupAve
0.15
10a
0.10
250
500
750
1250 Step
1000
1500
1750
2000
2000
Step
Figure 34: (left) Test MAE per epoch for predicting the dipole moment QM7b (with the c3nn model/group averaged shown for reference). (right) Test MAE per epoch for predicting the isotropic component of the a tensor.
54

---

Published as a conference paper at ICLR 2026
D.8 RMD17
We use the revised MD17 has high level of numerical noise (Chmiela et al. 2017). The revised MD17 dataset was cal- culated with a more accurate DFT functional/comvergennce criteria than the original MD17. For this dataset, we currently have explored the task indeepndend metric. We use the provided five train/test splits from https: /ttp:://Eigsharee com com/articces.. /ttaaatttteeeeeeeee /tw.ssseeeeeiiieeee__eee and train a separate model for each molecule. Note is not recommended to train a model on more than 1,000 samples from rMD17 Christeensnnnnen & von Lilien- feld (2020), even though the dataset has 100,000 conformers for each trajectory. We train a generic transformer with 812k parameters for 50 epochs the train/test splits provided with the Adam optimizer at learning rate 1e-5 and batch size 128.
& von Lilienfeld (020). as the original MD17 dataset
As seen in Figure 35 all molecules canonicalized. However, the task- inskinndeenddent metric of test accuracy yields significantly different values per molccules. For example, aspirin has a test accuracy of 97. 97.869% but ethanol yields 79.834 %. In Figure 35 ethanol and malohaldehyde have noticeably lower degree of canonicalization.
Test Accuracy vs. Augmented Fraction for MD17
100
naphthalene
aspirin
90
salicylic
Accuracy (%6
paracetamol
80
malonaldehyde
toluene
azobenzene
70
benzene
ethanol
60
50
0.0
0.2
0.4
0.6
0.8
1.0
Augmented Fraction
Figure 35: Test accuracy vs. augmented fraction for molecules in rMD17. Note the difference between the more canonicalized molecules and ethano//malonalcchyd..
55

---

Published as a conference paper at ICLR 2026
D.9 OPEN CATALYST PROJECT 2020 (OC20)
For our study, we use the 200K subset from the structure to energy and forces (S2EF) task, available https::/fairr-chemniiittt struucure-too andrrg--add-ooore----- github.
For this dataset, we have explored the task-dependent metric. It would be interesting in the future explore other tasks (e.g. Ini- tial Structure to Relaxed Structure) and larger dataset sizes, as the OC20 dataset training set alone has 20 million structures. We use the preprocessing pipeline provided Shnttts://fairraee github. 1o/corr////// dwtp//aaaassttt.eeeeeseaaa oc20. html Positions for each catalyst+adsorbate are tagged with catalyst surface, I: catalyst sub-surface, and adsorbate. The unit cell for the catalyst repeated twice in the direction, twice in the y direction, and once in the direction, leading to the slab's alignment with the xy plane. This alignment most likely trivially causes our metric to report distributional symmetry breaking. We also expect the adsorbate alone be slightly less canonicalized than the combined catalyst surface- surface-adsorbate system (as the adsorbate alone is not periodically repeating slab). This is supported by the test accuracy, which is 96.529% for the adsor- bate alone compared to 99.280% for the surface plus adsorbate system. It would thus interesting future work to consider how to treat periodic crystalline systems.
56

---

Published as a conference paper at ICLR 2026
D.10 P-VALUES FOR m(px)
Following (Chiu & Bloem-Reddy 2023), our computation of m(px) can be extended to yield p- values via Algorithm 2 To construct null distribution, repeatedly sample train/test splits from the original dataset, apply a random rotation to all points each split, and record the resulting pairwise distances. We then compute test statistic by sampling a second collection of train/test splits, applying random rotations to only subset of points in each, and taking the mean distance across these splits. The p-value the proportion of null distances exceeding this test statistic. Here, the distance function is the classifier test accuracy used to compute m mpx).
We also consider Maximum Mean Discrepancy (MMD) an alternative distance metric for compar- ing distributions over point clouds (Chiu & Bloem-Reddy. 2023). MMD measures the discrepancy between two distributions without requiring explicit density estimation or distributional assump- tions, making it well-suited to high-dimensional settings. The advantage of our method is its flexi- bility and interpretability, but we show here that the symmetry biases we detected are strong enough to also be detected by MMD, dependent we use an unbiased empirical estimator with three different kernel choices reflecting natural notions of distance between point clouds. The most naive choice, which we call the Mean/Covar kernel, computes distances between the means and covariances of two point clouds; while simple and effi- cient, it fails capture local geometric structure. The Chamfer distance kernel instead measures the average nearest-nett nearest-neighhor distance between two point clouds, capturing local structure but ignoring global shape distribution and density. Finally, the Hausdorff distance kernel replaces the averaging in Chamfer distance with maximum operation, making it more sensitive to outliers and worst-case geometric discrepancies.
the choice of kernel. To compute MMD in practice,
Algorithm 2 p-value Computation
1: Input: Training set Zuain test set Peea. calibration distances sample size n1, actual distances sample size n2, distance function Distance(:.)).
2: Output: p-value
3: actual dists t-
4:
culbration_diste <
>> Compute calibration distances under null hypothesis
5: i= to ny do
6: Sample training set Luain and test set Hest from Puain and Prest-
7: Apply rotation transformation all data
8: d, f Distance 2uiin 2rea)
9:
calibration aitiooonddists.att (ppend))
10:
end for
Compute actual distances
11:
for i= 1
n2 do
Sample training set T trann and test set Vest lest from Duain and Brett
13: Apply rotation transformation to subset of data
14:
da f Distancee 9 main L' test)
15:
ectual.dists.appentt 16: end for
17: da + 1 En actual-
Compute mean of actual distances
02
18: count + ld, E calibration dists d, > doll
19: p-value < T44
1+couutt
return p-value
D.10.1 QM9 VALUUS
Figure 36 demonstrates the values used in our computation of the p-values for each method (on row) and different levels of augmentation in the detection dataset (column) for QM9, The p-value plots were computed using 20 samples (for each histogram) of size 1k, trained for 20 epochs (in the case of the classifier metric). All methods exhibit the expected behavior: the augmented fraction increases -ie. as the distribution becomes more similar to the reference, perfectly symmetrized distribution the distance decreases.
57

---

Published as a conference paper at ICLR 2026
Augmenttd Indien 00
Aupmented Frartion 3:55
Augmented farter:0.
Augmerrn Racten
Augmented Fratten
Precuery
Procaacy
%y
0.54
5.01
0.005
0,210
6.015
0.002
0.004
3,004
5,005
01030
Augmented Froetio 0.D
Aupmented Fraction 5.25
Aaymmetted Faatt on
Auomerme Foction
Auomented Frccion
Prequencs
01005
60000
0.015
0.002.0
00/30 2300
2/94
34006
0.033
0.004
0.0030
5,0035
Augmertzd Frodton 0.D
Augmented Faacion 9.25
Augmented Hract on
Augmerizt Fradton 0.35
Augmented Fraction 3,0
Prog ePoy
0.00 3,505, 9.30 0.015
2.2250 Rtst: 9275
t,0t25 MMst 0,20
9.001 BLOZ 1822 0.06
0.001
2,022
2,005
Merrowwr
Meartowwr
Augmented Tractor: 0.0
Augmentted Tracton: 3:25
Augmanted Pract 0r:
Augmentes Practio: 0.75
Aupmented Fraction 3.0
60 Cussilirr
30
50
70
50
90
42
20
55
60
40
49
35
55
Chass
Figure 36: Distance metrics for different methods, and different levels of augmentation (i.e. dif- ferent levels of underlying distributional similarity).
03077
0.0
0.5
055 Augmantac
Augmarrid
MeaniCosar
0.9122
0.3000
Augrersad Tracion
Au gn Pract
Figure 37: Different distance metrics from perfectly symmetrized distribution, a function of the degree of synthetic augmentation of the QM9 dataset. (Higher augmented fraction indicates a greater similarity the symmetrized distribution.)
D.10.2 RMD17 P-VALS
Figure 41 demonstrates the values used in computation of the p-values for each method (on row) and different levels of augmentation in the detection dataset (column) for one the molecules in rMD17 (benzene). The p-value plots were computed using 20 samples (for each histogram) of size Ik corresponding to the given train/test splits, trained for 20 epochs (in the case of the clas- sifier metric). As shown, all methods separate the calibration distances from the actual distances.
58

---

of the size), convergence is Published as a conference paper at ICLR 2026
MMD Chamfer
MMD Hausdorff
Mean/Covar Classifier
pvvouee
20-4
Augmented Fraction
Figure 38: p-values for different methods, and at different levels of augmentation (i.e. different levels of underlying distributional similarity).
0.72
tr
0.70
0.40
03
255
5.0
:0.3 Epoch
12.5
27.5
0.0
25
7.5
10.0
12.
35.0
Speer
Figure 39: Left: the loss curve from of the 20 training used to compute the classifier distance in the p-value computation, on 1K examples. Right: the loss curve from a training run used compute the classifier distance over the full dataset. As shown, the loss converged much faster for the full dataset, whereas with only 1k examples (one much slower.
59

---

Published as a conference paper at ICLR 2026
resulting in identical, statistically significant p-values. As the tests asked to distinguish between increasingly similar datasets (moving from left to right), the histograms gradually move closer to- gether, until they overlap. For case of visualization, Figure 40 plots the mean distance computed from each histogram for benzene (excluding the calibration distances). We also plot the p-value the augmented fraction Figure 42 The Chamfer and Hausdorff kernels exhibit similar trends the classifier, and the naive mean/ covar kernel exhibits less reasonable behavior. This illustrates the importance of choosing a good kernel and provides relative advantage of method. All other molecules in rMD17 exhibted similar trends for the p-values.
MMD Chamfer
MMD Hausdorff
Distance
Distance
0.020
0.10
0.010
0.00
0.0
1.0
0.0
0.5
Augmented Fraction
Augmented Fraction
MMD Mean/Covar
Classifier
0.00403
Distance
80
0.00402
0.00401
0.00400
0.0
0.5
1.0
0.0
0.5
1.0
Augmented Fraction
Augmented Fraction
Figure 40: Different distance metrics from perfectly symmetrized distribution, as a function of the degree of synthetic augmentation of the rMD17 dataset for benezene. (Higher augmented fraction indicates greater similarity to the symmetrized distribution).
60

---

Published as a conference paper at ICLR 2026
Callbration
Augmeris froddar:
Aupmented Fracson: 2.25
Augmented 0.
Asgmeriss Phddlam
Auprenteed Fracten
Calicrtton
Chamler
Augmentaa fracden:
Augmented Fracson: 3.25
Augmenter on:
Augmersss Predtem:
Augrented
Augmented fracdor: 0.0
Aupmented Fraczon: 3.25
Augmented Praction
Augmennes Preciem
Augrented Fraction
0.2042
0.204011 04400111-0400000744000
Augmented frucder
Aupmented Fracton: 3:25
Augmented Prattio:::
Augmerres Precder
Augmented
Figure 41: Distance metrics for different methods, and at different levels of augmentation for ben- zene (i.e. different levels underlying distributional similarity).
p-value vs Augmented Fraction Benzene
0.5
MMD Chamfer
MMD Hausdorff
0.4
Mean/Covar
Classifier
povalue
0.3
0.2
0.1
0.0
0.2
0.4
0.6
0.8
1.0
Augmented Fraction
Figure 42: p-value augmented fraction for benzene rMD17.
61


Reviews:
[{"review_id": "y6xlLmwCR8", "reviewer": "Reviewer_oMVy", "summary": "This paper investigates the effect of an often overlooked issue in datasets of equivariant tasks, namely that many training data distributions are not actually invariant to the transformation that the trained equivariant model is. To this end, the authors first propose a simple metric to measure the distributional symmetry breaking of a given dataset. The metric uses a single classifier that tries to detect whether an input came from the original training distribution or from an augmented version of the same distribution. In a perfectly invariant data distribution, such a classifier can only do random predictions since the two distributions are identical. The authors propose to use the performance of such a classifier to define a quantitative measure of the degree of symmetry breaking. Along with the definition of the metric for measuring the distributional symmetry breaking, the authors use ridge regression to provide conditions under which data augmentation may actually hurt performance. Finally, they evaluate the proposed metric and theoretical insights across a large range of datasets, and show that many commonly used datasets for equivariant learning exhibit substantial distributional asymmetry.", "questions": "- How the number of training data and the parameter count of the classifier can affect the proposed metric. Is there any proposed heuristic for tuning such parameters, given a dataset we want to investigate?\n- How can a practitioner decide whether to use the task-dependent or the task-independent metric to measure the symmetry breaking of a given distribution?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper's topic is quite timely, since there is an open discussion regarding the benefits of equivariants compared to the most commonly used data augmentations. The proposed metric provides a practical and easily implementable tool for testing the distributional invariant of a distribution. This can help researchers better understand when and why equivariant models or data augmentation are preferable.\n- The theoretical analysis using ridge regression is clearly presented and provides formal examples that illustrate how augmentation can either hurt or improve performance. Although this analysis is performed in a simplified linear setting, it allows the readers to build intuition for the behavior of the more complex neural networks\n- The authors perform an extensive evaluation, covering a large range of tasks and showcasing interesting behaviors (even when they might not explicitly align with the main hypothesis of this work).", "weaknesses": "- There is limited discussion and investigation about a task-dependent symmetry-breaking metric. This limitation is crucial since in a lot of different settings, even in cases when the input data distribution is invariant, the ground truth labels break the symmetry. Without being able to effectively model this task-specific symmetry breaking, the proposed metric can be less useful since it only shows a partial image of the problem.\n- The results in the molecular datasets, such as QM9, contradict the main point of the paper. Although this is an important result that should be reported, it creates a more ambiguous interpretation regarding the utility of the proposed metric. \n- There is limited discussion and ablation studies about how the expressivity of the classifier and also the availability of training data can affect the proposed metric.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s9g0cOZxqY", "reviewer": "Reviewer_uapy", "summary": "This paper investigates the often-overlooked issue of distributional symmetry breaking, situations where data transformations (e.g., rotations, permutations) are not equally likely under the data distribution, violating the common assumption behind data augmentation and equivariant architectures.\nThe authors propose a simple yet powerful metric m(pX) based on a two-sample neural classifier test to quantify how far a dataset deviates from perfect symmetry. They show how this measure can diagnose anisotropy in various datasets such as MNIST, ModelNet40, and QM9.\nThrough theoretical analysis on invariant ridge regression, they demonstrate that data augmentation can harm performance under asymmetric covariance when invariant and non-invariant features are correlated. Extensive experiments across 2D, 3D, and molecular datasets confirm the theory, revealing that many commonly used benchmarks are surprisingly canonicalized.\nOverall, the paper argues that understanding when and why equivariance helps requires diagnosing symmetry biases in the data itself.", "questions": "Is there an empirical relationship between m(pX) and the observed gain/loss from equivariant methods? For example, could a threshold on m(pX) predict when augmentation is beneficial?\n\nCan the proposed metric be used during training to adaptively modulate augmentation strength or enforce “symmetry regularization”? This could make the method more practical.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces the notion of distributional symmetry breaking as a fundamental factor explaining inconsistencies in the effectiveness of equivariant or augmentation-based methods. This is an important conceptual contribution.\n\nThe proposed metric m(pX), implemented via a two-sample classifier, is elegant, interpretable, and easily applicable across data modalities (images, point clouds, molecules). It can serve as a general diagnostic tool for symmetry analysis.\n\nThe ridge regression analysis under asymmetric covariance provides deep insight into why augmentation may hurt, connecting symmetry assumptions to bias–variance trade-offs and data geometry.", "weaknesses": "While m(pX) quantifies symmetry breaking, the paper does not provide a formal correlation analysis between this metric and downstream model accuracy or robustness — the connection remains mostly qualitative.\n\nSince m(pX) depends on a trained neural network, results may vary with architecture, capacity, and hyperparameters. The paper claims low sensitivity, but further ablation or calibration would strengthen this claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qDtg815IUV", "reviewer": "Reviewer_RG36", "summary": "The paper explains that data augmentation and symmetrization can be harmful when invariant and non-invariant features are strongly correlated, and proposes a novel metric that measures the distributional symmetry breaking in a dataset.", "questions": "The paper’s framework relies on specific assumptions about the data distribution and the augmentation operator family. Could the authors clarify how restrictive these assumptions are in practice? For example, do they hold for common image augmentations such as Mixup or Cutout?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper provides a principled analysis of when data augmentation benefits or harms generalization.", "weaknesses": "While the paper presents a compelling framework for selective data augmentation, it could be strengthened by a deeper exploration of its practical limitations and boundary conditions. For instance, the theoretical analysis assumes access to accurate estimates of distributional alignment and bias-variance components, but the paper provides limited discussion on how these quantities can be reliably approximated in high-dimensional, real-world settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5lFsCU6KHB", "reviewer": "Reviewer_BM93", "summary": "This work proposes a metric based on neural network (NN) binary classifier training to evaluate the degree of distributional symmetry breaking in data. The authors provide both theoretical analyses and empirical validations on several datasets, including MNIST (2D images), ModelNet40 (3D image sets), QM9 and QM7B (molecular data), and large-scale materials datasets.", "questions": "Minor comments\n\n1. Term Clarification\n\nTerms such as anisotropy, symmetry, equivariance, distributional symmetry-breaking, distributional symmetry, perfectly symmetric, and invariant are used in many places, but they should be introduced only after clear definitions and used more consistently.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "It provides a comprehensive demonstration of the proposed ideas.\nIt includes theoretical efforts to clarify the impact of data-specific invariant and equivariant properties in terms of generalization.", "weaknesses": "1. Unclear Problem Definition\n\n It is unclear whether this paper aims to propose a new metric, present derived theoretical findings, or introduce a new concept (distributional symmetry breaking). If the goal is to propose a metric, the paper should clearly demonstrate how and why it functions effectively as one. If the focus is on presenting theoretical findings, it should explain why those findings are significant from both theoretical and experimental perspectives. If the work intends to overcome limitations of the underlying assumption, it should show what problems actually occur and that problem is solved by removing the assumption directly. Overall, the paper seems to mix these three goals without providing sufficient support for each direction as follows. \n\n2. Weak Justification of the Impact of Motivation \n\nOn page 1, line 42, p(x) is similar to p(gx) is described as the main cause of the problem to be tackled. However, it seems likely observed only in Elesedy & Zaidi’s work. To justify why it is worth tackling, the author needs to provide more explanation and literature works, because equivariance is used in various ways, whether full or partial, implicit or explicit, in practical neural networks, this assumption could have a large impact. However, this assumption is not explicitly made in many existing works.\n\n3. Loose Connection of the Theoretical Works to Contribution\n\nIn my view, conclusions such as the one on page 5, line 216, are not novel at all, which the author also notes as intuitive; the same applies to page 6, line 289. If clarifying the conclusions as a theoretical result, it requires justification of the novelty and impact of the \"conclusions\".\n\n4. Weak Validation As as a Metric\n\nThe proposed metric is a binary classifier NN–based approach. Even if it can be used as a test setup (Lopez-Paz & Oquab, 2017), it is a training-based metric, so many hyperparameter conditions affect the variance of test results, which is a significant risk for a metric used in fair comparisons. To make a contribution as a metric, it is necessary to demonstrate robustness to the training environment and sensitivity to actual equivariance changes. (Honestly, it may not be possible to provide data-agnostic standards. Even in this case, to use it as a metric, the risks should be analyzed to determine whether they are within an acceptable range.)\n\n\n5. A Flaw in Motivation\n\nThe distributional symmetry breaking is measured on page 3, line 142. This part seems flawed, because gx may be another observed sample x', which should have p(x'), but it is integrated to be matched to p(x) again. It means that the metric tries to represent one point of the data manifold with a region of the data manifold including it, which reduces the total density of the manifold. This confuses what distribution p(x) the metric is expected to represent at an optimum, given the broken probability mass.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)