Technical Summary: Equivariant Score-Based Generative Models Provably Learn Distributions with Symmetries Efficiently

Problem Statement

Generative modeling often faces challenges in data efficiency and computational cost, particularly when the target data distribution possesses intrinsic structural symmetries (e.g., group invariance found in physics, images, or molecular simulations). While empirical studies suggest that incorporating symmetries into generative models improves generalization and sampling efficiency, a rigorous theoretical understanding of these improvements in the context of Score-Based Generative Models (SGMs) has been limited. Specifically, there is a lack of theoretical guarantees quantifying the benefits of equivariant inductive bias versus standard data augmentation, and a formal analysis of how SGMs handle group-invariant distributions.

Methodology

The authors employ a combination of optimal transport theory, partial differential equation (PDE) analysis, and stochastic control theory to analyze SGMs under group symmetry.

1. Theoretical Framework:

   • Wasserstein-1 ($d_1$) Distance: The generalization error is measured using the Wasserstein-1 distance, which allows for meaningful bounds even when distributions lack densities or are supported on submanifolds.
   • Symmetrization Operators: The paper defines symmetrization operators for functions ($S_G$), probability measures ($\bar{S}_G$), and vector fields ($S^E_G$). These operators project functions and measures onto the space of $G$-invariant or $G$-equivariant counterparts using the Haar measure of the group $G$.
   • Hamilton-Jacobi-Bellman (HJB) Theory: The authors connect SGMs to Mean-Field Games (MFGs) and stochastic optimal control. They utilize the HJB equation to characterize the optimal velocity field (score function) of the generative process.

2. Score-Matching Objectives:
   The analysis covers Denoising Score Matching (DSM), Explicit Score Matching (ESM), and Implicit Score Matching (ISM). The authors investigate the behavior of these objectives when restricted to $G$-equivariant vector fields versus general vector fields, and when applied to augmented versus non-augmented datasets.

3. Error Quantification:
   A Model-Form Uncertainty Quantification (UQ) perspective is adopted, decomposing the generalization error into four sources:

   • $e_1$: Measurement of non-equivariance (Deviation from Equivariance, DFE).
   • $e_2$: Score-matching error with symmetrized vector fields.
   • $e_3$: Sample complexity bound of $d_1$ under group symmetry.
   • $e_4$: Errors from early stopping and time horizon.

Key Contributions

1. Improved Generalization Bounds for Symmetric Distributions

The paper provides the first rigorous error analysis for SGMs learning group-invariant distributions. By combining recent results on the robustness of SGMs to $d_1$ distance with the sample complexity of empirical estimations under group symmetry, the authors derive a generalization bound (Theorem 1 and Theorem 2).

• Result: The bound demonstrates that for $G$-invariant target distributions, the sample complexity improves significantly. Specifically, the expected $d_1$ distance between the empirical symmetrized distribution and the true distribution scales with the volume of the quotient space $\Omega/G$ rather than the full domain $\Omega$, leading to faster convergence rates (e.g., $O((|G|N)^{-1/d})$ for finite groups).

2. Equivariant Inductive Bias and HJB Theory

Using HJB theory, the authors prove that SGMs are intrinsically equivariant when the target distribution is $G$-invariant and the drift function is $G$-equivariant (Theorem 3).

• Result: The optimal score function for a symmetric target is necessarily $G$-equivariant. This establishes that the neural network approximating the score should be parameterized to be $G$-equivariant to incorporate the correct structural inductive bias.

3. Equivalence of Symmetrized and Raw Data Training

A central theoretical contribution is the proof that performing score matching with $G$-equivariant vector fields on raw (non-augmented) data is mathematically equivalent to performing score matching on symmetrized (augmented) data (Theorem 4 and Proposition 1).

• Result: For any distribution $\rho$, the Implicit Score Matching (ISM) objective restricted to equivariant vector fields yields the same optimal solution as the objective applied to the symmetrized distribution $\bar{S}_G[\rho]$. Furthermore, the optimal equivariant vector field for ESM/DSM is exactly the score of the symmetrized distribution.
• Implication: This implies that explicit data augmentation is not strictly necessary if the model architecture enforces equivariance; the network can learn the symmetrized score directly from unaugmented data.

4. Quantification of Model-Form Error

The paper quantifies the cost of failing to incorporate equivariant structure. By defining the "Deviation from Equivariance" (DFE), the authors show that the generalization error includes a term proportional to the non-equivariance of the learned vector field (Theorem 5).

• Result: If a non-equivariant vector field is used (even with data augmentation), the error term $DFE$ remains non-zero, potentially leading to worse generalization bounds compared to using an equivariant parametrization on unaugmented data.

Results and Numerical Validation

The theoretical claims are corroborated by numerical simulations on a 2D mixture of Gaussians with $C_4$ symmetry.

• Experimental Setup: Four configurations were tested: (1) Equivariant network, no augmentation; (2) Equivariant network, with augmentation; (3) Non-equivariant network, with augmentation; (4) Non-equivariant network, no augmentation.
• Findings:
  • The Equivariant, not augmented case consistently achieved the lowest Wasserstein-1 distance ($d_1$), often outperforming the Equivariant, augmented case, particularly with limited training data ($N=10$).
  • The Non-equivariant, augmented case did not consistently outperform the Non-equivariant, not augmented case, suggesting that data augmentation alone cannot fully compensate for the lack of equivariant inductive bias in the model architecture.
  • As the training sample size increased, the performance gap between equivariant and non-equivariant models narrowed, but the equivariant approach remained superior or comparable, validating the theoretical bound that $DFE$ contributes to the error.

Significance and Claims

The paper claims to provide the first theoretical analysis and guarantees for SGMs learning distributions with group symmetry. Its primary significance lies in:

1. Theoretical Justification: Rigorously explaining why symmetry-respecting models are more data-efficient, linking this to improved sample complexity bounds in the $d_1$ metric.
2. Practical Guidance: Demonstrating that data augmentation is not a substitute for equivariant parametrization. The authors argue that enforcing equivariance in the vector field (via neural network architecture) is sufficient to learn the symmetrized distribution, potentially avoiding the computational overhead of data augmentation while achieving better generalization.
3. Error Decomposition: Introducing a UQ perspective that isolates "model-form error" (non-equivariance) as a distinct and quantifiable source of generalization failure, distinct from statistical estimation error.

The authors conclude that while data augmentation provides gains in sample complexity, the most effective strategy for minimizing generalization error in symmetric settings is to ensure the score approximation itself is $G$-equivariant, thereby eliminating the model-form error term.