Technical Summary: Symmetries in PAC-Bayesian Learning

Problem Statement

Symmetries in machine learning, such as rotational or translational invariance, are empirically known to improve model performance. However, existing theoretical guarantees explaining these gains are limited by restrictive assumptions. Prior theoretical frameworks, particularly within the PAC-Bayesian setting, typically rely on two key constraints:

1. Compactness: The symmetry group must be compact (e.g., finite groups or rotations), excluding non-compact groups like translations or scalings which are central to architectures like Convolutional Neural Networks (CNNs).
2. Data Invariance: The data distribution itself is assumed to be invariant under the group action. In real-world scenarios, this is rarely true (e.g., an upside-down cup of water is unlikely to occur in natural data, violating translation invariance).

This work addresses the gap in theoretical understanding by extending generalization guarantees to settings involving non-compact symmetries and non-invariant data distributions.

Methodology

The authors utilize the PAC-Bayesian framework, which allows for data-dependent priors and posteriors over hypothesis classes. The methodology proceeds in three main stages:

1. Problem Formulation and Assumptions

The paper formalizes the data generation process where outputs $Y$ are generated from inputs $X$ via a group-equivariant function $f^*$ perturbed by independent noise $\Xi$ (Assumption 2.1).

• Group Action: A topological group $G$ acts on the input space $X$ and output space $Y$.
• Equivariance: A hypothesis $f$ is equivariant if $f(g \cdot x) = g \cdot f(x)$.
• Loss Function: The loss $\ell$ is assumed to be convex in its first argument and $G$-invariant (Assumption 4.1), meaning $\ell(g \cdot y, g \cdot \hat{y}) = \ell(y, \hat{y})$.

Crucially, the authors relax the requirement for $G$ to be compact and for the data distribution to be invariant. They assume the group action on the input space is free (Assumption 3.1) to facilitate the construction of an averaging operator.

2. The Averaging Operator and KL Decomposition

To handle non-compact groups and non-invariant distributions, the authors introduce a generalized averaging operator $Q$. Unlike previous works that rely on the Haar measure (which requires compactness), this operator integrates over the group orbit using a probability kernel derived from the Disintegration Theorem (Theorem 1.1).

• Definition: $Q(f)(x) = \pi_G(x) \cdot \int_G g^{-1} \cdot f(g \cdot \pi_{X_\phi}(x)) \kappa(\pi_{X_\phi}(x), dg)$, where $\kappa$ is the conditional distribution of the group given the orbit representative.
• Property: $Q$ maps arbitrary measurable functions to equivariant ones and acts as a projection ($Q^2 = Q$).

Using this operator, the authors derive a KL divergence decomposition (Lemma 3.4). For any two measures $\mu$ and $\nu$ on the hypothesis class:
$$D_{KL}(\mu \| \nu) = D_{KL}(Q_*\mu \| Q_*\nu) + \text{Remainder}$$
This shows that the KL divergence between the original measures is strictly greater than or equal to the KL divergence between their equivariant projections ($Q_*\mu$ and $Q_*\nu$). The remainder term captures the "non-equivariant" information discarded by the projection.

3. Generalization Bounds

The authors apply this decomposition to McAllester's PAC-Bayes bound.

• Complexity Reduction: By restricting the prior and posterior to equivariant functions (via the pushforward $Q_*$), the complexity term (involving $D_{KL}$) is reduced (Corollary 3.6).
• Risk Reduction: The authors prove that symmetrizing the hypothesis does not increase the true risk and that risks can be computed entirely on orbit representatives without information loss (Theorem 4.3).
• Main Result: The derived bound (Equation 14) holds for non-compact groups and non-invariant data, providing a tighter generalization guarantee for equivariant models compared to standard bounds applied to the full hypothesis class.

Key Contributions

1. Extension to Non-Compact Symmetries: The paper provides the first PAC-Bayesian generalization analysis that accommodates non-compact symmetry groups (e.g., translations, scalings, Lorentz transformations), covering translation-invariant architectures within a unified framework.
2. Relaxation of Data Invariance: The theoretical guarantees do not require the data distribution to be invariant under the group action, addressing a major limitation of prior work where real-world data often violates this assumption.
3. Generalized Averaging Operator: The introduction of an averaging operator based on the Disintegration Theorem allows for the projection of hypotheses onto equivariant subspaces without relying on the existence of a uniform (Haar) measure, which is unavailable for non-compact groups.
4. Tighter Bounds: The work demonstrates that enforcing equivariance leads to a strictly smaller KL divergence term in the PAC-Bayes bound, theoretically explaining the improved generalization performance of symmetric models.

Experimental Results

The authors validate their theory through numerical experiments on several datasets (Toy Example, MNIST, CIFAR-10/100, ModelNet, Top Tagging) with both compact (SO(2), SO(3)) and non-compact (SE(2), $\mathbb{R}^+$, $SO^+(1,3)$) groups.

• Setup: They compare baseline neural networks (non-equivariant) with equivariant counterparts (using libraries like e2cnn and e3nn). The experiments explicitly use non-uniform sampling of transformations to violate data invariance assumptions.
• Findings:
  • KL Reduction: The symmetric posteriors consistently achieved significantly lower KL divergence compared to baselines (e.g., 14.5% to 89.2% reduction across datasets).
  • Bound Tightening: The PAC-Bayesian generalization bounds were substantially tighter for equivariant models (reductions of 10% to 34%).
  • Performance: Equivariant models matched or improved upon baseline test risk (e.g., 5.3% to 41.5% risk reduction in specific cases), confirming that the theoretical tightening translates to empirical gains even when data distributions are not invariant.

Significance and Claims

The paper claims to provide theoretical evidence that symmetric models are preferable beyond the narrow setting of compact groups and invariant distributions. By removing these restrictive assumptions, the work broadens the theoretical foundations of symmetry in machine learning.

The authors state that their findings offer a principled explanation for the empirical success of symmetry-aware architectures (like CNNs) in practical scenarios where data is neither perfectly invariant nor governed by compact groups. They emphasize that their framework applies to general hypothesis classes and follows standard PAC-Bayesian formulations, making the results applicable to a wide range of learning problems involving non-compact symmetries. The work does not claim to solve all generalization problems but specifically targets the theoretical gap regarding the role of symmetries in non-ideal, real-world data distributions.