Efficient Learning and Symmetry Discovery under Exact Invariances

Data augmentation is a simple and model-agnostic approach for exploiting known invariances in learning problems. Given a group acting on the input space, one augments the training set with transformed copies of each sample. Because it exploits symmetries without modifying the underlying learning algorithm, data augmentation can be applied broadly across learning methods. However, this universality comes at a computational cost: when the group is large, full group-sized augmentation quickly becomes computationally infeasible. This raises a fundamental question: Can partial data augmentation achieve the same statistical benefits as full augmentation in terms of generalization and sample complexity? We develop a general framework for investigating this question using Fourier analysis and the representation theory of finite groups. We show that, for a broad class of classical learning problems, partial data augmentation based on a randomly sampled subset of group elements achieves the same minimax rates as full augmentation, up to an approximation error that vanishes as the subset size increases. Our results provide a theoretical explanation for why partial augmentation can retain the statistical benefits of full augmentation despite enforcing symmetry only approximately, and shed light on a recently raised question in learning with symmetries: whether statistically optimal learning under general group invariances can be achieved using computationally scalable methods. Moreover, we prove a complementary impossibility result: enforcing exact invariance via data augmentation requires averaging over the entire group, and cannot be achieved by any strict subset when the hypothesis space is sufficiently expressive. Together, these results provide a unified perspective on full and partial data augmentation, as well as exact and approximate symmetry enforcement.

Grokking occurs when a model achieves high training accuracy but generalization to unseen test points happens long after that. This phenomenon was initially observed on a class of algebraic problems, such as learning modular arithmetic (Power et al., 2022). We study grokking on algebraic tasks in a class of feature learning kernels via the Recursive Feature Machine (RFM) algorithm (Radhakrishnan et al., 2024), which iteratively updates feature matrices through the Average Gradient Outer Product (AGOP) of an estimator in order to learn task-relevant features. Our main experimental finding is that generalization occurs only when a certain symmetry in the training set is broken. Furthermore, we empirically show that RFM generalizes by recovering the underlying invariance group action inherent in the data. We find that the learned feature matrices encode specific elements of the invariance group, explaining the dependence of generalization on symmetry.

Continuous symmetry discovery is typically formulated as a separate estimation problem for each function or dataset. We instead investigate whether the map from a function to its symmetry structure can be learned across a family of functions from finite observations. Using the infinitesimal differential characterization of continuous symmetries, we formulate this map as an inverse symmetry operator, analogous to an inverse PDE problem. For global Lie-group symmetries, we establish a representation-theoretic population characterization: the symmetry Lie algebra is the nullspace of a positive-semidefinite operator constructed from representation coefficients and the infinitesimal group action. We also show that the inverse symmetry operator is equivariant under coordinate changes through the adjoint action. This characterization separates the geometric symmetry map from finite-sample estimation, motivating a hybrid approach that amortizes statistically difficult components while retaining analytic structure. On SO(5) benchmarks, amortized equivariant estimation of harmonic coefficients improves symmetry recovery from sparse observations, whereas ordinary least squares suffices for SL(2,R). We extend this viewpoint to local symmetries, where symmetry directions form smoothly varying tangent distributions rather than a fixed Lie algebra. Under the Lie-bracket consistency condition, Frobenius' theorem guarantees local integral submanifolds along which observables remain unchanged. We recover this geometry using a hybrid estimator that learns observation weights while preserving analytic differential and geometric recovery. Across global, non-compact, local, and physically motivated synthetic benchmarks, our results support learning the statistically difficult finite-sample components while preserving the known geometric structure of symmetry recovery.