Equivariant score-based generative models provably learn distributions with symmetries efficiently

## Abstract

Abstract Symmetry is ubiquitous in many real-world phenomena and tasks, such as physics, images, and molecular simulations. Empirical studies have demonstrated that incorporating symmetries into generative models can provide better generalization and sampling efficiency when the underlying data distribution has group symmetry. In this work, we provide the first theoretical analysis and guarantees of score-based generative models (SGMs) for learning distributions that are invariant with respect to some group symmetry and offer the first quantitative comparison between data augmentation and adding equivariant inductive bias. First, building on recent works on the Wasserstein-1 ( 𝐝 1 subscript 𝐝 1 \mathbf{d}_{1} bold_d start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ) guarantees of SGMs and empirical estimations of probability divergences under group symmetry, we provide an improved 𝐝 1 subscript 𝐝 1 \mathbf{d}_{1} bold_d start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT generalization bound when the data distribution is group-invariant. Second, we describe the inductive bias of equivariant SGMs using Hamilton-Jacobi-Bellman theory, and rigorously demonstrate that one can learn the score of a symmetrized distribution using equivariant vector fields without data augmentations through the analysis of the optimality and equivalence of score-matching objectives. This also provides practical guidance that one does not have to augment the dataset as long as the vector field or the neural network parametrization is equivariant. Moreover, we quantify the impact of not incorporating equivariant structure into the score parametrization, by showing that non-equivariant vector fields can yield worse generalization bounds. This can be viewed as a type of model-form error that describes the missing structure of non-equivariant vector fields. Numerical simulations corroborate our analysis and highlight that data augmentations cannot replace the role of equivariant vector fields.

## 1 Introduction

Improving data efficiency and reducing computational costs are central concerns in generative modeling. In the case when the target data distribution has intrinsic structure, such as group symmetry, the task of distribution learning can be made more efficient and stable by leveraging the structure of the data. Various empirical studies such as structure-preserving GANs [], equivariant normalizing flows [, ] and equivariant and structure-preserving diffusion models [, ] have shown that symmetry-respecting generative models can effectively learn a group-invariant
distribution even with limited data. However, theoretical understanding of these improvements is still limited. To our knowledge, the only work that provides theoretical performance guarantees is [] for group-invariant GANs. In this work, we present new rigorous analysis explaining why score-based generative models (SGMs), or diffusion models [, , , ], can more efficiently learn group-invariant distributions by incorporating the underlying symmetry into the score approximation, as empirically observed in [].

#### Our contributions.

We provide the first rigorous error analysis for SGMs with symmetry as well as the first quantitative comparison between data augmentations and incorporating inductive bias of symmetries into generative models. First, by combining recent results relating to the robustness of SGMs with respect to the Wasserstein-1 ($\mathbf{d}_{1}$) distance [] and the sample complexity of empirical estimations of $\mathbf{d}_{1}$ for distributions with group symmetry [, ], we derive a generalization bound for SGMs with group symmetry to explain the sample efficiency gained when using the symmetry explicitly during training. (See Theorem 1 and Theorem 2)

Second, we demonstrate the inductive bias of equivariant SGMs using Hamilton-Jacobi-Bellman theory (see Theorem 3) and show that performing standard score-matching, a crucial step in SGM, with respect to any distribution by a $G$-equivariant vector field is equivalent to score-matching with respect to the symmetrized distribution, and that the optimal vector field is exactly the score of the symmetrized distribution (See Theorem 4 and Proposition 1). This provides insights into how to avoid potentially expensive data augmentation by embedding symmetries directly into the score approximation, typically achieved through a $G$-equivariant neural network. Moreover, we compare the impact of non-equivariant score matching via symmetrically augmented datasets with the use of equivariant score matching via the non-augmented datasets using both theory and numerical simulations.

We adopt a model-form uncertainty quantification (UQ) perspective, attributing errors in equivariant SGMs to the following four sources:

- •
$e_{1}$: Measurement of the non-equivariance of the learned score function;
- •
$e_{2}$: Score-matching error with symmetrized vector field;
- •
$e_{3}$: Sample complexity bound of $\mathbf{d}_{1}$ with group symmetry;
- •
$e_{4}$: Error due to early stopping and time horizon.

We show that the generalization error as measured by the expected Wasserstein-1 distance between the generated and target data distributions is bounded by a combination of these four errors above. A particular novelty of our UQ analysis is the quantification of the model-form error $e_{1}$ of the equivariant structure. This type of UQ perspective was introduced recently for SGMs without structure []. Detailed description and discussion of the derived bounds are found in Theorem 2 and Eq. 21.

#### Related work.

Various symmetry-preserving generative models have been proposed such as structure-preserving GANs [], equivariant normalizing flows [, ], equivariant flow matching [], and equivariant diffusion models for molecule generation []. Theoretical analysis of performance guarantees for such models, to our knowledge, has only been conducted for group-invariant GANs []. In the context of SGMs, the convergence and generalization of SGMs without group symmetry have been well-studied. The quality of a generated distribution for approximating a target distribution is typically measured by probability divergences and distances. For example, [, , , , ] prove generalization bounds for TV, $\chi^{2}$, and $\mathbf{d}_{1}$ by bounding the KL divergence, which is a stronger divergence. Our results, however, cannot be derived from bounding the KL divergence. The direct $\mathbf{d}_{1}$ generalization bounds have been derived in [, ], but [] relies on a particular discretization of SGMs. In [], empirical estimates of the $\mathbf{d}_{1}$ distance on compact domains of $\mathbb{R}^{d}$ are shown to obtain a faster convergence assuming the group is finite. Subsequently, [] extended the $\mathbf{d}_{1}$ bound to closed Riemannian manifolds with infinite groups. Our generalization bound for SGM with symmetry is built on the $\mathbf{d}_{1}$ bounds and UQ perspective for SGMs without structure [] and the convergence of the empirical estimations of $\mathbf{d}_{1}$ distance with group symmetry [, ]. Recent work [] empirically studies diffusion models with equivariance and proposes various implementations. However, it only provides some guarantees to ensure the generated distribution is $G$-invariant, but no further theory is shown beyond numerical experiments to demonstrate the data efficiency.

The rest of the paper is organized as follows. In Section 2, we review score-based generative models, score-matching objectives, and the notion of group symmetry. We present our theoretical results of generalization bounds in Section 3. The inductive bias of equivariant SGMs and properties of score-matching with equivariant vector fields are presented in Section 4. In Section 5, we discuss the importance of equivariant parametrizations for obtaining a better generalization bound and related insights for practical implementations. In Section 6, we provide numerical experiments that corroborate our theory and insights. All the proofs can be found in Section 7. We conclude our paper with a discussion in Section 8. All the proofs can be found in the appendix.

## 2 Background

In this section, we introduce group actions and symmetrization operators, and review the score-matching objectives for score-based generative modeling.

### 2.1 Group actions and symmetrization operators

Let $\Omega$ be the domain, $\mathcal{P}(\Omega)$ the space of probability measures on $\Omega$, and $\mathcal{M}_{b}(\Omega)$ be the space of bounded measurable functions on $\Omega$. A group is a set $G$ equipped with a group product satisfying the axioms of associativity, identity, and invertibility. Given a group $G$ and a set $\Omega$, a map $\theta:G\times\Omega\to\Omega$ is called a group action on $\Omega$ if $\theta_{g}\coloneqq\theta(g,\cdot):\Omega\to\Omega$ is an automorphism on $\Omega$ for all $g\in G$, and $\theta_{g_{2}}\circ\theta_{g_{1}}=\theta_{g_{2}\cdot g_{1}}$, $\forall g_{1},g_{2}\in G$. By convention, we will abbreviate $\theta(g,x)$ as $gx$ throughout the paper.

A function $\gamma:\Omega\to\mathbb{R}$ is called $G$-invariant if $\gamma\circ\theta_{g}=\gamma,\forall g\in G$. On the other hand, a probability measure $P\in\mathcal{P}(\Omega)$ is called $G$-invariant if $P=(\theta_{g})_{*}P,\forall g\in G$, where $(\theta_{g})_{*}P\coloneqq P\circ(\theta_{g})^{-1}$ is the push-forward measure of $P$ under $\theta_{g}$. We denote the set of all $G$-invariant distributions on $\Omega$ as $\mathcal{P}_{G}(\Omega)\coloneqq\{P\in\mathcal{P}(\Omega):P~{}\text{is}~{}G
\text{-invariant}\}$.

In this paper, the domain $\Omega$ is bounded; in particular, we focus on the torus $\Omega=R\mathbb{T}^{d}$ with radius $R$, which is equivalent to a bounded domain with periodic boundary conditions, as considered in []. We make the following assumption on $G$ in this paper.

###### Assumption 1 .

$G$ is a group such that the mapping $g:\Omega\to\Omega$ can be written as $g(x)\mapsto A_{g}x$ for some unitary matrix $A_{g}\in\mathbb{R}^{d\times d}$ for any $g\in G,x\in\Omega$. That is, any $g\in G$ is a linear isometry.

Next, we introduce two symmetrization operators from [], that are useful for our theoretical analysis.

Symmetrization of functions: $S_{G}:\mathcal{M}_{b}(\Omega)\to\mathcal{M}_{b}(\Omega)$,

$$ $S_{G}[\gamma](x)\coloneq\int_{G}\gamma(gx)\mu_{G}(\mathop{}\!\mathrm{d}{g})= \mathbb{E}_{\mu_{G}}[\gamma\circ g(x)],$ (1) $$

where $\gamma\in\mathcal{M}_{b}(\Omega)$ and $\mu_{G}$ is the unique Haar probability measure of $G$.

Symmetrization of probability measures (dual operator): $S^{G}:\mathcal{P}(\Omega)\to\mathcal{P}(\Omega)$, defined for $\gamma\in\mathcal{M}_{b}(\Omega)$ by

$$ $\mathbb{E}_{S^{G}[P]}\gamma\coloneq\int_{\Omega}S_{G}[\gamma]\mathop{}\! \mathrm{d}{P(x)}=\mathbb{E}_{P}S_{G}[\gamma].$ (2) $$

It is shown in [] that both $S_{G}$ and $S^{G}$ define projections. We also abuse the notation that if $P$ evolves with time, then $S^{G}[P]$ means the symmetrization of $P$ at each time.

We say a vector field $\mathbf{s}:\Omega\times[0,T]\to\mathbb{R}^{d}$ is $G$-equivariant if

$$ $\mathbf{s}(gx,t)=A_{g}\cdot\mathbf{s}(x,t)$ (3) $$

for any $x\in\Omega$, $g\in G$. It can be easily verified that if $\rho\in\mathcal{P}_{G}(\Omega)$, then its score $\nabla\log\rho$ is $G$-equivariant. In addition to $S_{G}$ and $S^{G}$, we propose

Symmetrization of vector fields:
$S_{G}^{E}:(\Omega\times[0,T]\to\mathbb{R}^{d})\to(\Omega\times[0,T]\to\mathbb{
R}^{d})$,

$$ $S_{G}^{E}[\mathbf{s}](x,t)\coloneq\int_{G}A_{g}^{\top}\cdot\mathbf{s}(gx,t)\mu _{G}(\mathop{}\!\mathrm{d}{g})$ (4) $$

for any vector field $\mathbf{s}$, which is an extension of formula (12) in [] for finite groups. It can be shown that $S_{G}^{E}[\mathbf{s}]$ is $G$-equivariant for any vector field $\mathbf{s}$. The proof can be found in Section 7.3. By the definition of equivariance, we immediately have $S^{E}_{G}[\mathbf{s}]=\mathbf{s}$ if $\mathbf{s}$ is $G$-equivariant.

###### Assumption 1 .

### 2.2 Score-based generative modeling

Given a drift term or a vector field $\mathbf{f}(x,t)$, we consider the following forward and backward diffusion processes

$$ $\displaystyle\mathop{}\!\mathrm{d}{x_{s}}$ $\displaystyle=-\mathbf{f}(x_{s},T-s)\mathop{}\!\mathrm{d}{s}+\sigma(T-s) \mathop{}\!\mathrm{d}{W_{s}},\quad x_{0}\sim\pi;$ (5) $\displaystyle\mathop{}\!\mathrm{d}{y_{t}}$ $\displaystyle=\left(\mathbf{f}(y_{t},t)+\sigma(t)^{2}\nabla\log\eta^{\pi}(y_{t },T-t)\right)\mathop{}\!\mathrm{d}{t}+\sigma(t)\mathop{}\!\mathrm{d}{W_{t}}, \quad y_{0}\sim m_{0},,$ (6) $$

where $x_{s}\sim\eta^{\pi}(\cdot,s)$. Here, $\nabla\log\eta^{\pi}(x,t)$ is called the score function. It is known from [] that if $m_{0}=\eta^{\pi}(\cdot,T)$, then we have $y_{t}\sim\eta^{\pi}(\cdot,T-t)$. In this work, we consider $\mathbf{f}=0$ and $\sigma(t)=\sqrt{2}$, and the target distribution $\pi\in\mathcal{P}_{G}(\Omega)$.

Score functions are typically approximated by optimizing parametrized vector fields with respect to the discretization of one of several score-matching objective functions. The denoising score matching (DSM) [] objective is defined as:

$$ $\mathcal{J}_{D}(\eta^{\pi},\theta)=\int_{0}^{T}\int_{\Omega}\int_{\Omega}\left |\mathbf{s}_{\theta}-\nabla\log\eta^{x^{\prime}}\right|^{2}\mathop{}\!\mathrm{ d}{\eta^{x^{\prime}}}(s)\mathop{}\!\mathrm{d}{\pi}(x^{\prime})\mathop{}\! \mathrm{d}{s},$ (7) $$

where $\eta^{x^{\prime}}(s)$ denotes the conditional probability from $x^{\prime}$ at time 0 to $x$ of Eq. 5 at time $s$. In addition, we also introduce two other types of score-matching objectives.

The
explicit score matching (ESM) objective [], is defined as:

$$ $\mathcal{J}_{E}(\rho,\theta)=\int_{0}^{T}\int_{\Omega}\left|\mathbf{s}_{\theta }-\nabla\log\rho\right|^{2}\mathop{}\!\mathrm{d}{\rho}(s)\mathop{}\!\mathrm{d} {s},$ (8) $$

and it is obvious that $\mathcal{J}_{E}(\rho,\theta)=\mathcal{J}_{D}(\rho,\theta)$.

The implicit score matching (ISM) objective [], is defined as:

$$ $\mathcal{J}_{I}(\rho,\theta)\coloneq\int_{0}^{T}\int_{\Omega}\left(\left| \mathbf{s}_{\theta}\right|^{2}+2\nabla\cdot\mathbf{s}_{\theta}\right)\mathop{} \!\mathrm{d}{\rho(s)}\mathop{}\!\mathrm{d}{s},$ (9) $$

which is more practical for score-matching. By an expansion of the square of the norm, it is easy to verify that $\mathcal{J}_{D}(\rho,\theta)=\mathcal{J}_{E}(\rho,\theta)=\mathcal{J}_{I}(\rho
,\theta)+4\left\lVert\nabla\sqrt{\rho}\right\rVert_{2}^{2}$ for any $\rho\in\mathcal{P}(\Omega)$. This suggests that the optimal solutions to the DSM, ESM and ISM coincide for the same $\rho$. We also abuse the notation using $\mathcal{J}(\rho,\mathbf{s})$ for a generic vector field $\mathbf{s}$ with an additional subscript on $\mathcal{J}$ when referring to a specific score-matching objective.

## 3 Theoretical result (I): equivariant SGMs have improved 𝐝 1 subscript 𝐝 1 \mathbf{d}_{1} bold_d start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT generalization bounds

The probability distance we use to measure the generalization error in this paper is the Wasserstein-1 distance $(\mathbf{d}_{1})$, defined as:

$$ $\mathbf{d}_{1}(\pi_{1},\pi_{2})=\sup_{\gamma\in\Gamma}\left\{\mathbb{E}_{\pi_{ 1}}[\gamma]-\mathbb{E}_{\pi_{2}}[\gamma]\right\}$ (10) $$

for any $\pi_{1}$, $\pi_{2}\in\mathcal{P}(\Omega)$, where $\Gamma=\text{Lip}_{1}(\Omega)$ is the set of 1-Lipschitz function on $\Omega$.

In this section, we derive a generalization bound with improved sample complexity in $\mathbf{d}_{1}$ for learning a $G$-invariant target distribution.

Let $\pi$ be the target data distribution that is $G$-invariant. In SGMs, the generated distribution is $m(T)$, where $m(t)$ follows the denoising diffusion process Eq. 6 with $\nabla\log\rho$ replaced by $\mathbf{s}_{\theta}$ through score-matching. That is,

$$ $\partial_{t}m=\Delta m+2\operatorname{div}(m\mathbf{b}_{\theta})\,\,\text{in} \,\,\Omega\times(0,T],\quad m(0)=\frac{1}{\text{vol}(R\mathbb{T}^{d})}\,\, \text{in}\,\,\Omega,$ (11) $$

where $\mathbf{b}_{\theta}(x,t)=\mathbf{s}_{\theta}(x,T-t)$.

In practice, we only have access to finite samples drawn from $\pi$, denoted by $\{z_{i}\}_{i=1}^{N}$. Thus, the score-matching or the DSM objective Eq. 7 is often approximated when $\eta^{\pi}(t)$ is replaced by its kernel density estimate $\eta^{N}(t)$, where $\eta^{N}(0)=\pi^{N}\coloneq\frac{1}{N}\sum_{i=1}^{N}\delta_{z_{i}}$. Since the kernel estimate does not have a well-defined score at $s=0$, the DSM objective is often integrated only for $s\in[\epsilon,T]$, an example of early-stopping in SGM []. More specifically, this is equivalent to score-matching for the mollified distribution $\pi^{N,\epsilon}=\pi^{N}\star\Gamma_{\epsilon}$, where $\Gamma_{\epsilon}$ is the heat kernel with time $\epsilon$ and the symbol $\star$ denotes the convolution. In the symmetry-preserving SGM, we consider the symmetrized measure $\pi^{N,\epsilon}_{G}$, defined as []: $\frac{\mathop{}\!\mathrm{d}{\pi^{N,\epsilon}_{G}}}{\mathop{}\!\mathrm{d}{x}}=
\sum_{l=0}^{\infty}\exp(-\epsilon\lambda_{l})\mu_{l}\phi_{l}$, where $\mathop{}\!\mathrm{d}{x}$ indicates the uniform measure of $\Omega$, and $(\lambda_{l},\phi_{l})$ is the pair of the eigenvalues and eigenfunctions of the Laplace-Beltrami operator of $\Omega$, $\mu_{l}\coloneq\frac{1}{N}\sum_{i=1}^{N}\mathbf{1}_{G}(l)\phi_{l}(X_{i})$, and $\mathbf{1}_{G}(l)=1$ if and only if $\phi_{l}$ is $G$-invariant. In particular, we have $\pi^{N}_{G}\coloneq\pi^{N,0}=S^{G}[\pi^{N}]$. It is evident that $\pi^{N,\epsilon}_{G}=S^{G}[\pi^{N}]\star\Gamma_{\epsilon}$.

In summary, in the context of SGMs, $\pi^{N,\epsilon}=\pi^{N}\star\Gamma_{\epsilon}$ corresponds to early stopping; $\pi^{N}_{G}=S^{G}[\pi^{N}]$ refers to data augmentations; $\pi^{N,\epsilon}_{G}=S^{G}[\pi^{N}]\star\Gamma_{\epsilon}$ is the early stopping version of the data-augmented empirical distribution.

Here, we extend the $\mathbf{d}_{1}$ generalization bound as presented in [] to the case when the target distribution is $G$-invariant.

Let $\eta^{N,\epsilon}_{G}:\Omega\times[0,T]\to[0,\infty)$ be the solution to

$$ $\begin{cases}\partial_{t}\rho-\Delta\rho=0\,\,\text{in}\,\,\Omega\times(0,T], \\ \rho(0)=\pi^{N,\epsilon}_{G}\,\,\text{in}\,\,\Omega,\end{cases}$ (12) $$

We first prove the finite-sample generalization bound for $\mathbf{d}_{1}(\pi,m(T))$.

###### Theorem 1 .

Assume $\mathcal{J}_{D}(\eta^{N,\epsilon}_{G},\mathbf{s}_{\theta})\leq e_{nn}$. Then for $\epsilon<1$ and up to a dimensional constant $C=C(d)>0$,

$$ $\mathbf{d}_{1}(\pi,m(T))\lesssim\sqrt{\epsilon}+R^{3/2}(1+\sqrt{\left\lVert \nabla\mathbf{s}_{\theta}\right\rVert_{\infty}})\left(Re^{-\frac{wT}{R^{2}}} \mathbf{d}_{1}(\pi,\frac{1}{\text{vol}(R\mathbb{T}^{d})})+\sqrt{e_{nn}^{\prime }}\right),$ $$

where

$$ $e_{nn}^{\prime}\lesssim e_{nn}+\left(1-\frac{\log(\epsilon)}{\sqrt{\epsilon}}+ \frac{1}{\sqrt{T}}+T\left\lVert\mathbf{s}_{\theta}\right\rVert^{2}_{C^{2}( \Omega\times[0,T])}\right)\mathbf{d}_{1}(\pi^{N}_{G},\pi),$ $$

and $\pi^{N}_{G}$ is the symmetrization of non-symmetric empirical distribution $\pi^{N}$; i.e., $\pi^{N}_{G}=S^{G}[\pi^{N}]$.

###### Remark 1 .

The assumption that $\mathcal{J}_{D}(\eta^{N,\epsilon}_{G},\mathbf{s}_{\theta})\leq e_{nn}$ implies that the score approximation is trained via DSM with augmented samples. This suggests that equivariant SGMs can be implemented through data augmentations. As we shall see in Sections 5 and 6, a better implementation of equivariant SGMs should rely on equivariant parametrizations of the score function.

Similar to [], we derive the following averaged generalization bound by taking the expectation with respect to the empirical distributions and subsequently applying Jensen’s inequality. However, the $G$-invariance of the target distribution $\pi$ provides a significant improvement in the data efficiency in the bounds.

###### Theorem 2 (Average bound) .

Let $e_{nn},A>0$ and assume that for each empirical measure $\pi^{N}$ consisting of $N$ samples from $\pi$ there exists $\mathbf{s}_{\theta}$ such that

$$ $\mathcal{J}_{D}(\eta^{N,\epsilon}_{G},\mathbf{s}_{\theta})\leq e_{nn},$ $$

with

$$ $\left\lVert\mathbf{s}_{\theta}\right\rVert_{C^{2}(\Omega\times[0,T])}\leq A.$ $$

Let $m(T)$ be the generated distribution. Then for sufficiently large $T$, up to a dimensional constant $C$ that only depends on $R$ and $d$ and is independent of random samples or $N$, we have

$$ $\mathbb{E}\left[\mathbf{d}_{1}(\pi,m(T))\right]\lesssim\sqrt{\epsilon}+R^{3/2} (1+\sqrt{A})\left(Re^{-\frac{wT}{R^{2}}}\mathbf{d}_{1}(\pi,\frac{1}{\text{vol} (R\mathbb{T}^{d})})+\sqrt{e_{nn}^{\prime}}\right),$ $$

where

$$ $e_{nn}^{\prime}\lesssim e_{nn}+\left(1-\frac{\log(\epsilon)}{\sqrt{\epsilon}}+ \frac{1}{\sqrt{T}}+TA^{2}\right)\mathbb{E}\left[\mathbf{d}_{1}(\pi^{N}_{G},\pi )\right].$ $$

###### Theorem 1 .

###### Remark 1 .

###### Theorem 2 (Average bound) .

#### On the importance of 𝐝 1 subscript 𝐝 1 \mathbf{d}_{1} bold_d start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT .

The use of $\mathbf{d}_{1}$ distance on both sides of our generalization bounds has two key implications:

(1) We can take advantage of the $G$-invariance of $\pi$ and improve data efficiency since $\mathbf{d}_{1}$ allows gains on $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]$. First, it is shown in [] that on bounded domains of $\mathbb{R}^{d}$, we have

$$ $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]\lesssim\begin{cases}\left(\frac{1} {\left|G\right|N}\right)^{1/d}&\text{if }d\geq 3,\\ \left(\frac{1}{\left|G\right|N}\right)^{1/2}\log N&\text{if }d=2,\\ \frac{\text{diam}(\Omega/G)}{N^{1/2}}&\text{if }d=1,\end{cases}$ (13) $$

if $G$ is finite. Later, [] extends it to closed Riemannian manifolds with possibly infinite $G$ such that $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]\lesssim\left(\frac{\text{vol}(
\Omega/G)}{N}\right)^{1/d^{*}}$, where $\text{vol}(\Omega/G)$ is the volume of the quotient space $\Omega/G$ and $d^{*}=\text{dim}(\Omega/G)\geq 3$. This sample complexity gain cannot be derived for the KL or other $f$-divergences without additional regularization.

(2) The $\mathbf{d}_{1}$ bounds in Theorem 1 and Theorem 2 remain well-defined and meaningful even when the target probability distribution does not have a density. In particular, Theorem 2 has the following corollary when the target distribution is supported on a smooth submanifold $\mathcal{M}\subset\Omega$.

###### Corollary 1 .

Follow the same assumption and quantities as in Theorem 2, and assume that $\pi$ is supported on a closed submanifold $\mathcal{M}\subset\Omega$, and $G$ admits a unitary representation in $\Omega$ as in 1. Then up to a dimensional constant $C>0$ that also depends on $\mathcal{M}$, such that

$$ $\mathbb{E}\left[\mathbf{d}_{1}(\pi,m(T))\right]\lesssim\sqrt{\epsilon}+R^{3/2} (1+\sqrt{A})\left(Re^{-\frac{wT}{R^{2}}}\mathbf{d}_{1}(\pi,\frac{1}{\text{vol} (R\mathbb{T}^{d})})+\sqrt{e_{nn}^{\prime}}\right),$ $$

where

$$ $e_{nn}^{\prime}\lesssim e_{nn}+\left(1-\frac{\log(\epsilon)}{\sqrt{\epsilon}}+ \frac{1}{\sqrt{T}}+TA^{2}\right)\left(\frac{\text{vol}(\mathcal{M}/G)}{N} \right)^{1/d^{*}},$ $$

where $\text{vol}(\mathcal{M}/G)$ is the volume of the quotient space $\mathcal{M}/G$ and $d^{*}=\text{dim}(\mathcal{M}/G)\geq 3$, and $\mathbf{d}_{1}$ here denotes the Wasserstein-1 distance on $\Omega$.

###### Corollary 1 .

## 4 Theoretical result (II): equivariant parametrizations restore intrinsic equivariance of SGMs

Theorem 1 and Theorem 2 do not explicitly convey the significance of equivariant vector fields in score matching. First, we illustrate the importance of equivariance from a Hamilton-Jacobi-Bellman (HJB) perspective in Section 4.1 by showing that SGMs are *intrinsically* equivariant. Second, we highlight the role of $G$-equivariant vector fields (typically parameterized by neural networks) in score matching, an aspect that has only been addressed experimentally in previous studies. Our rigorous results indicate that it is sufficient to perform score matching with $G$-equivariant vector fields in relation to an unsymmetrized distribution. This approach will be particularly beneficial when we only have a finite set of unaugmented samples (i.e., a non-symmetric empirical distribution drawn from an invariant distribution).
This latter aspect will be discussed in detail in Section 4.2, Section 5 and tested in Section 6.

### 4.1 HJB describes the inductive bias of equivariant SGMs

First, we use the connections between SGMs and PDE theory to provably show that score-based generative models are intrinsically equivariant under relatively mild assumptions. Score-based generative models have been shown to be well-posed through their connections with stochastic optimal control and mean-field games (MFGs) [, , ]. In [, ], it was shown that score-based generative models are solutions of a mean-field game, more specifically, one that corresponds with the Wasserstein proximal of the cross-entropy. The peculiar structure of cross-entropy is why SGMs can be trained by score-matching alone. The MFG is an infinite-dimensional optimization problem

$$ $\displaystyle\min_{v,\rho}\left\{-\int_{\Omega}\log\pi(x)\rho(x,T)dx+\int_{0}^ {T}\int_{\Omega}\left[\frac{1}{2}\|v\|^{2}-\nabla\cdot f\right]\rho(x,t)dxdt\right\}$ (14) s.t. $\displaystyle\partial_{t}\rho+\nabla\cdot((f+\sigma v)\rho)=\frac{\sigma^{2}}{ 2}\Delta\rho,\,\rho(x,0)=\eta(x,T).$ $$

The density of particles evolve according to the controlled Fokker-Planck equation. The terminal cost is equivalent to the cross entropy of $\pi$ with respect to the terminal density $\rho(x,T)$. The running cost is, via the Benamou-Brenier formulation of optimal transport, the Wasserstein-2 distance with a state cost $-\nabla\cdot f$.

The solution of the MFG optimization problem is characterized by its optimality conditions, which are a pair of nonlinear partial differential equations.

$$ $\displaystyle\begin{dcases}-\partial_{t}U-f^{\top}\nabla U+\frac{1}{2}|\sigma \nabla U|^{2}+\nabla\cdot f=\frac{\sigma^{2}}{2}\Delta U\\ \partial_{t}\rho+\nabla\cdot(\rho(f-\sigma^{2}\nabla U))=\frac{\sigma^{2}}{2} \Delta\rho\\ U(x,T)=-\log\pi(x),\,\rho(x,0)=e^{-U(x,0)}.\end{dcases}$ (15) $$

This first equation is a Hamilton-Jacobi-Bellman equation, which determines the optimal velocity field $v^{*}(x,t)=-\sigma\nabla U$ for the second equation, a controlled Fokker-Planck.
By a Hopf-Cole (logarithmic) transformation, this pair of PDEs is equivalent to the noising-denoising SDE system. Let $U(x,t)=-\log\eta(x,T-t)$, then for $s=T-t$, we have

$$ $\displaystyle\begin{dcases}\frac{\partial\eta}{\partial s}=-\nabla\cdot(f\eta) +\frac{\sigma^{2}}{2}\Delta\eta\\ \frac{\partial\rho}{\partial t}=-\nabla\cdot(\rho(f+\sigma^{2}\nabla\log\eta(x ,T-t)))+\frac{\sigma^{2}}{2}\Delta\rho\\ \eta(x,0)=\pi(x),\,\,\,\,\rho(x,0)=\eta(x,T).\end{dcases}$ $$

We can then see that the optimal velocity field has the form $v^{\star}(x,t)=-\sigma(t)\nabla U(x,t)=\sigma(T-t)\nabla\log\eta(x,T-t)$, which is precisely related linearly with respect to the score function of the forward noising process.

Assuming that the target distribution $\pi$ is $G$-invariant and the drift function $f$ is $G$-equivariant, the resulting optimal velocity field, and therefore, the score function is necessarily $G$-equivariant.

###### Theorem 3 .

Consider the score-based generative model given by the equivalent MFG Eq. 14 and let $U$ be the solution to the HJB equation in Eq. 15. Assume the target data distribution $\pi$ is $G$-invariant and that the drift in the noising dynamics is $G$-equivariant. Then we have that the corresponding score function is $G$-equivariant, namely

$$ $\mathbf{s^{*}}(x,t)=-\nabla U(x,t)=\operatorname*{arg\,min}_{\mathbf{s}\in \Omega\times[0,T]\to\mathbb{R}^{d}}\mathcal{J}_{E}(\rho,\mathbf{s})\in V_{G}\,,$ (16) $$

where we denote by $V_{G}\subset\Omega\times[0,T]\to\mathbb{R}^{d}$, the subspace of $G$-equivariant vector fields.

The MFG perspective is useful as the proof for this theorem immediately follows from basic uniqueness results from PDE theory. This theorem states that, mathematically, SGMs are symmetry-preserving for invariant target measures when the drift function also preserves the same symmetry. The most trivial case is when $f=0$.

###### Remark 2 (Equivariant inductive bias) .

In the SGM algorithm the optimal vector field $\mathbf{s^{*}}(x,t)$ is the score and is learned as part of the algorithm. Therefore, this theorem shows that the corresponding neural network for the approximation of $\mathbf{s^{*}}(x,t)$ should be parameterized in a way that is also $G$-equivariant, thus incorporating an induced, equivariant (structural) inductive bias.

In practice, however, finite samples and approximation can break this *theoretical* symmetry. Therefore, the symmetry must be eventually *restored* by enforcing the inductive bias of the model. Benefits of this latter algorithmic perspective are explored and discussed in the following subsection and Section 6.
This set of challenges is in direct analogy with similar fundamental issues in molecular dynamics and Hamiltonian systems: for instance, “symplectic numerical integration” methods have been developed to transfer the theoretical structure preserving properties of Hamiltonian systems to the numerical schemes for Hamiltonian systems. For example, see [] and [].

###### Theorem 3 .

###### Remark 2 (Equivariant inductive bias) .

### 4.2 Properties of score-matching with equivariant vector fields

First, we show that for any distribution $\rho$, the ISM objective when restricted to $G$-equivariant vector fields, is equivalent to the ISM objective with respect to its symmetrized counterpart. Second, we prove that using equivariant vector fields can reduce the DSM error for $G$-invariant distributions.

###### Theorem 4 .

Consider the ISM problem in Eq. 9, in which $\rho$ is not necessarily $G$-invariant. Then for any $G$-equivariant vector field $\mathbf{s}$, we have

$$ $\mathcal{J}_{I}(\rho,\mathbf{s})=\mathcal{J}_{I}(S^{G}[\rho],\mathbf{s}).$ $$

###### Remark 3 .

Theorem 4 is important for practical implementations, in the sense that the optimal equivariant vector field can be obtained by score-matching for raw data without data augmentation. We will demonstrate this point in our numerical simulations in Section 6.

Moreover, for the ESM (or equivalently, the DSM) problem of a generic probability measure, the $G$-equivariant minimizer is exactly the score of the symmetrized probability measure, namely:

###### Proposition 1 .

Consider the ESM problem in Eq. 8, in which $\rho$ is not necessarily $G$-invariant. Denote by $V_{G}\subset\Omega\times[0,T]\to\mathbb{R}^{d}$, the subspace of $G$-equivariant vector fields. Then we have

$$ $\operatorname*{arg\,min}_{\mathbf{s}\in V_{G}}\mathcal{J}_{E}(\rho,\mathbf{s}) =\nabla_{x}\left[\log\left(S^{G}[\rho]\right)\right].$ $$

We propose the following definition as an error quantification for the non-equivariance of a vector field with respect to a $G$-invariant measure $\rho\in\mathcal{P}_{G}(\Omega)\times[0,T]$.

###### Definition 1 (Deviation from equivariance) .

The deviation from equivariance (DFE) of a vector field $\mathbf{s}$ with respect to $\rho\in\mathcal{P}_{G}(\Omega)\times[0,T]$ is defined as

$$ $\text{DFE}(\rho,\mathbf{s})\coloneq\int_{0}^{T}\int_{\Omega}\left|\mathbf{s}-S _{G}^{E}[\mathbf{s}]\right|^{2}\mathop{}\!\mathrm{d}{\rho(s)}\mathop{}\! \mathrm{d}{s}.$ (17) $$

It is evident that $DFE(\rho,\mathbf{s})=0$ if $\mathbf{s}$ is $G$-equivariant. Given this definition, we obtain the following decomposition of the ESM and DSM objectives.

###### Theorem 5 .

For any $\rho\in\mathcal{P}_{G}(\Omega)\times[0,T]$ and any vector field $\mathbf{s}$, we have

$$ $\mathcal{J}_{E}(\rho,\mathbf{s})=\text{DFE}(\rho,\mathbf{s})+\mathcal{J}_{E}( \rho,S^{E}_{G}[\mathbf{s}]).$ (18) $$

As DSM and ESM are equivalent objectives, we readily have

$$ $\mathcal{J}_{D}(\rho,\mathbf{s})=\text{DFE}(\rho,\mathbf{s})+\mathcal{J}_{D}( \rho,S^{E}_{G}[\mathbf{s}])\,,\,\,\text{for any}\,\,\rho\in\mathcal{P}_{G}( \Omega)\times[0,T]\,.$ (19) $$

Finally, the following proposition indicates that for any learned distribution $\eta$, its symmetrized counterpart $S^{G}[\eta]$ is always closer to the $G$-invariant target distribution $\pi$ in the $\mathbf{d}_{1}$ sense. The $G$-invariance of the generated distribution is guaranteed by the $G$-equivariant vector field $\mathbf{s}_{\theta}$ (see Corollary 2).

###### Proposition 2 .

For any $\eta,\pi\in\mathcal{P}(\Omega)$, and $\pi$ is $G$-invariant, we have

$$ $\mathbf{d}_{1}(\eta,\pi)\geq\mathbf{d}_{1}(S^{G}[\eta],\pi).$ $$

###### Theorem 4 .

###### Remark 3 .

###### Proposition 1 .

###### Definition 1 (Deviation from equivariance) .

###### Theorem 5 .

###### Proposition 2 .

## 5 The significance of equivariant vector fields in SGMs

With the theoretical results established in Section 3 and Section 4,
we can now focus on providing quantitative comparisons between equivariant vector fields and data augmentations.
Our strategy relies on making the generalization bound in Theorem 2 as small as possible. In particular, we take a closer look at the terms $e_{nn}$ and $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]$, which can be improved by selecting an appropriate structure for the vector field or by implementing data augmentations.

The assumption $\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\mathbf{s}_{\theta})\leq e_{nn}$ in Theorem 2 refers to the error of DSM with augmented data. Technically, this assumption ensures the same generalization bounds derived in Theorem 1 and Theorem 2, regardless of whether the vector field $\mathbf{s}_{\theta}$ is $G$-equivariant or not. Note also that the gain in $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]$ (see the paragraph after Theorem 2 for the sample complexity gain) is not affected no matter whether we use equivariant vector fields. However, $\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\mathbf{s}_{\theta})$ or $e_{nn}$ does depend on the structure of vector fields and can be improved accordingly as we see next.

- •
Data augmentation without equivariant structure:
If we perform data augmentations without using equivariant vector fields, then we have to pay the cost of data augmentations. Moreover, by Theorem 5,
$e_{nn}=\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\theta)=\text{DFE}(\eta_{G}^{N,
\epsilon},\mathbf{s}_{\theta})+\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},S_{G}^{E}
[\mathbf{s}_{\theta}])\,,$
(20)
therefore $e_{nn}$ has a lower bound of $\text{DFE}(\eta_{G}^{N,\epsilon},\mathbf{s}_{\theta})$ that measures the distortion of vector fields from equivariance, which can be large if the vector fields are highly “non-equivariant”.
- •
Equivariant structure without data augmentation:
On the contrary, if we simply use equivariant vector fields without data augmentations, by Theorem 4, we can automatically obtain the score approximations of $\eta_{G}^{N,\epsilon}$ by simply solving the ISM objective of unaugmented samples $\eta^{N,\epsilon}$. Thus, the assumption $\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\theta)\leq e_{nn}$ is valid in practice. The main difference with the simple data augmentation case discussed above is that here, due to restricting the SGM on equivariant vector fields, we have $\text{DFE}(\eta_{G}^{N,\epsilon},\mathbf{s}_{\theta})=0$. Therefore, the term $e_{nn}$ in the generalization bounds can be made as small as possible.

To summarize, the generalization bound in Theorem 2 can be re-written as

$$ $\mathbb{E}\left[\mathbf{d}_{1}(\pi,m(T))\right]\lesssim\text{DFE}(\eta_{G}^{N, \epsilon},\mathbf{s}_{\theta})+\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},S_{G}^{E} [\mathbf{s}_{\theta}])+\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]+C(\epsilon, T),$ (21) $$

where $C(\epsilon,T)$ accounts for the error from early stopping and time horizon, and is independent of the equivariance structure or data augmentations we are studying. This suggests that while data augmentations can provide gains in $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]$, in order to further minimize the generalization error, one should make $\text{DFE}(\eta_{G}^{N,\epsilon},\mathbf{s}_{\theta})=0$; that is, applying $G$-equivariant vector fields.

## 6 Numerical example

We provide a simple numerical experiment to validate the basic results of our theory. The primary purpose is to emphasize minimizing the score-matching objective with respect to a non-symmetric sample of $G$-invariant distribution $\pi$ within a class of $G$-equivariant vector fields is better than just augmenting the data through group actions, as is indicated by our analysis encapsulated in the generalization bound Eq. 21.

We consider a mixture of 4 Gaussians centered at $[\pm 5,\pm 5]$ in $\mathbb{R}^{2}$. The group is generated by the action of moving from one center to the next. We report the $\mathbf{d}_{1}$ distance between the generated distribution and the target distribution. We consider four experimental setups: the first case (Equivariant, not augmented) is where the score network is parametrized to be $G$-equivariant by parametrizing it as

$$ $\displaystyle\mathbf{s}^{G}_{\theta}(x,t)=\frac{1}{|G|}\sum_{g\in G}A_{g}^{ \top}\mathbf{s}_{\theta}(A_{g}x,t),$ (22) $$

where $|G|=4$ is the order of the group. The score is trained on $N_{training}$ samples that are not augmented. The second case (Equivariant, augmented) is where the score network is parametrized as in Eq. 22, and is trained on data that is augmented by applying each group action on each training sample (hence effectively $4\times N_{training}$ samples). The third case (Non-equivariant, augmented) is where the network $\mathbf{s}_{\theta}$ is trained directly but on augmented training data. The fourth case (Non-equivariant, not augmented) is where the network $\mathbf{s}_{\theta}$ is trained directly and the training data is not augmented. For each case, the function $\mathbf{s}_{\theta}$ is parametrized via a fully-connected neural network with 3 hidden layers and 32 nodes per layer. It is trained over 10000 iterations via stochastic gradient descent, where the batch size is $N_{batch}=32$. For $N_{training}=10$, we sample with replacement in the SGD.

The Wasserstein-1 distance is computed using its dual form $\mathbf{d}_{1}(\eta,\pi)=\sup\left\{\mathbb{E}_{\eta}[\psi]-\mathbb{E}_{\pi}[
\psi]:\psi\in\text{Lip}_{1}(\Omega)\right\}$. The function $\psi$ is parametrized by a fully-connected neural network with two hidden layers with 64 nodes per layer. Spectral normalization [] is applied to enforce the Lipschitzness of $\psi$.

For each value of $N_{training}$ we perform 25 runs of each method. The mean and standard deviation of the 25 runs are reported in Table 1 and in Figure 1. Notice that the equivariant case consistently performs better than the data-augmented case, which corroborates our theoretical analysis. Moreover, the results suggest that training a non-equivariant score network on augmented data may not necessarily produce a superior model to the case when the data is not augmented.

In Figure 2, we show the generated samples of each case when $N_{training}=40$. Observe that the only way to consistently produce an invariant generated distribution is to have use an equivariant score approximation. Moreover, note that the reduction of $\mathbf{d}_{1}$ becomes marginal for large $N_{training}$ as other errors in the Theorem 2 are independent of the number of training samples.

**Table 1: $\mathbf{d}_{1}$ value for a 2d Gaussian mixture**
| $N_{training}$ | Equivariant, | Equivariant, | Non-equivariant, | Non-equivariant, |
| --- | --- | --- | --- | --- |
| augmented | not augmented | augmented | not augmented |  |
| $10$ | $1.36\pm 0.06$ | $1.82\pm 0.08$ | $1.93\pm 0.49$ | $2.64\pm 0.65$ |
| $100$ | $0.70\pm 0.09$ | $0.88\pm 0.10$ | $1.26\pm 0.45$ | $1.43\pm 0.35$ |
| $1000$ | $0.51\pm 0.12$ | $0.70\pm 0.11$ | $1.14\pm 0.32$ | $1.04\pm 0.33$ |
| $10000$ | $0.52\pm 0.10$ | $0.57\pm 0.12$ | $1.02\pm 0.20$ | $1.02\pm 0.23$ |

Figure: Figure 1: Wasserstein distance as a function of training sample size
Refer to caption: extracted/5890974/figure/wasserstein_distance_plot._with_equivariant_augmentationpng.png

Figure: (a) Truth
Refer to caption: x1.png

## 7 Proofs

### 7.1 Proof of Theorem 1

We also define the $G$-regularized Wasserstein-1 distance $(\mathbf{d}_{1}^{G})$ as:

$$ $\mathbf{d}_{1}^{G}(\pi_{1},\pi_{2})=\sup_{\gamma\in\Gamma_{G}^{inv}}\left\{ \mathbb{E}_{\pi_{1}}[\gamma]-\mathbb{E}_{\pi_{2}}[\gamma]\right\},$ (23) $$

where $\Gamma_{G}^{inv}$ is the subset of $\Gamma$ that consists of all $G$-invariant 1-Lipschitz functions.

The following theorem is adapted from Theorem 3.1 in []. Here we prove a version with group symmetry. The main difference is that the test function is now restricted to the class of $G$-invariant 1-Lipschitz functions, which is guaranteed by the equivariance of $b^{1}$.

###### Theorem 6 (Wasserstein Uncertainty Propagation) .

Let $\Omega=R\mathbb{T}^{d}$. Let $G$-equivariant vector fields $b^{1},b^{2}:\Omega\times[0,T]\to\mathbb{R}^{d}$ be given with $\left\lVert\nabla b^{1}\right\rVert_{\infty}<\infty$ and $m_{1},m_{2}\in\mathcal{P}_{G}(\Omega)$. If $m^{i}$ for $i=1,2$ are given by

$$ $\partial_{t}m^{i}-\Delta m^{i}-\operatorname{div}(m^{i}b^{i})=0,\,\,m^{i}(0)=m _{i}.$ (24) $$

Then up to a universal constant $C>0$, we have

$$ $\mathbf{d}_{1}^{G}(m^{2}(T),m^{1}(T))=\mathbf{d}_{1}(m^{2}(T),m^{1}(T))\leq CR ^{\frac{3}{2}}(1+\sqrt{\left\lVert\nabla b^{1}\right\rVert_{\infty}})(\mathbf{ d}_{1}^{G}(m_{2},m_{1})+\epsilon_{1}),$ $$

if

$$ $\left\lVert b^{2}-b^{1}\right\rVert_{L^{2}(m^{2})}\coloneq\left(\int_{0}^{T} \int_{\Omega}\left|(b^{2}-b^{1})(x,t)\right|^{2}m^{2}(t,x)\mathop{}\!\mathrm{d }{x}\mathop{}\!\mathrm{d}{t}\right)^{\frac{1}{2}}\leq\epsilon_{1}.$ $$

###### Proof.

The measure $\lambda=m^{1}-m^{2}$ satisfies the PDE

$$ $\partial_{t}\lambda-\Delta\lambda-\operatorname{div}(\lambda b^{1}+m^{2}(b^{1} -b^{2}))=0\,\,\text{in}\,\,\Omega\times(0,T),\quad\lambda(0)=m_{2}-m_{1}\,\, \text{in}\,\,\Omega.$ (25) $$

Let $\phi:\Omega\times[0,T]\to\mathbb{R}$ be a test function in space and time. We integrate in space and time against the PDE Eq. 25 and apply integration by parts to obtain

$$ $\displaystyle\int_{\Omega}\lambda(x,T)\phi(x,T)-\lambda(x,0)\phi(x,0)\mathop{} \!\mathrm{d}{x}$ $\displaystyle+\int_{0}^{T}\int_{\Omega}\lambda(-\partial_{t}\phi-\Delta\phi+b^ {1}\cdot\nabla\phi)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{t}$ (26) $\displaystyle+\int_{0}^{T}\int_{\Omega}m^{2}\nabla\phi\cdot(b^{1}-b^{2}) \mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{t}=0$ $$

Notice that if we choose the test function $\phi$ to satisfy the Kolmogorov backward equation (KBE)

$$ $-\partial_{t}\phi-\Delta\phi+b^{1}\cdot\nabla\phi=0\,\,\text{in}\,\,\Omega \times[0,T),\quad\phi(x,T)=\psi(x)\,\,\text{in}\,\,\Omega$ (27) $$

with terminal condition $\psi\in\mathcal{F}$, then from Eq. 26, we have

$$ $\int_{\Omega}\lambda(x,T)\psi(x)\mathop{}\!\mathrm{d}{x}=\int_{\Omega}\lambda( x,0)\phi(x,0)\mathop{}\!\mathrm{d}{x}+\int_{0}^{T}\int_{\Omega}m^{2}(t)\nabla \phi(x,t)\cdot(b^{2}-b^{1})(t)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{t}.$ (28) $$

Let $\mathcal{F}$ be the set of $G$-invariant 1-Lipschitz functions on $\Omega$. Taking the supremum over $\mathcal{F}$ we have

$$ $\mathbf{d}_{1}^{G}(m^{2}(T),m^{1}(T))\leq\sup_{\psi\in\mathcal{F}}\left|\int_{ \Omega}\lambda(x,0)\phi(x,0)\mathop{}\!\mathrm{d}{x}\right|+\sup_{\psi\in \mathcal{F}}\left|\int_{0}^{T}\int_{\Omega}m^{2}\nabla\phi\cdot(b^{2}-b^{1}) \mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{t}\right|.$ (29) $$

Also recall that $\phi$ is related to $\psi$ via the KBE Eq. 27.
We first show that $\phi(x,t)$ is always $G$-invariant for any $t\in[0,T)$ as long as $\psi$ is $G$-invariant. Indeed, if we perform a Hopf-Cole transform $u=-2\log\phi$, then Eq. 27 is equivalent to the Hamilton-Jacobi-Bellman (HJB) equation for $u$

$$ $-\partial_{t}u-\Delta u+\frac{1}{2}\left|\nabla u\right|^{2}+V\cdot\nabla\phi= 0,\,\,u(x,T)=-2\log(\psi(x)).$ (30) $$

On the other hand, it can easily be verified that $h(x,t)=u(gx,t)$ also satisfies Eq. 30 for any $g\in G$ since $A_{g}$ is unitary and $b^{1}$ is $G$-equivariant. The existence and uniqueness of the solution to Eq. 30 [] guarantees that $h(x,t)=u(x,t)$ is $G$-invariant for any $t\in[0,T)$ and therefore we have $\phi(x,t)=\phi(gx,t)$ for any $g\in G$ and $t\in[0,T)$. The rest of the proof, i.e., the gradient estimate of $\phi$ is exactly the same as that of Theorem 3.1 in [] since any $\psi\in\mathcal{F}$ is 1-Lipschitz.
∎

###### Corollary 2 .

Suppose a probability measure $m(x,t)$ evolves according to the KBE Eq. 27. That is,

$$ $-\partial_{t}m-\Delta m+V\cdot\nabla m=0\,\,\text{in}\,\,\Omega\times[0,T), \quad m(x,T)=m_{0}\,\,\text{in}\,\,\Omega$ (31) $$

where the vector field $V$ is $G$-equivariant and the terminal measure $m_{0}$ is $G$-invariant. Then $m(x,t)$ is $G$-invariant for all $t\in[0,T)$.

###### Proof.

By a change of variable $t\mapsto-t$ in the KBE Eq. 27, the statement follows the proof after Eq. 30.
∎

The following proposition shows that for empirical measures, the action of diffusion and symmetrization are commutable.

###### Proposition 3 .

$S^{G}[\pi^{N,\epsilon}]=S^{G}[\pi^{N}]\star\Gamma_{\epsilon}$.

###### Proof.

For any $\gamma\in\mathcal{M}_{b}(\Omega)$, we have

$$ $\displaystyle\mathbb{E}_{S^{G}[\pi^{N,\epsilon}]}\gamma$ $\displaystyle=\mathbb{E}_{\pi^{N,\epsilon}}S_{G}[\gamma]$ $\displaystyle=\int_{\Omega}\pi^{N}\star\Gamma_{\epsilon}S_{G}[\gamma]\mathop{} \!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}\int_{\Omega}\pi^{N}(y)\Gamma_{\epsilon}(x- y)\mathop{}\!\mathrm{d}{y}\gamma(gx)\mu_{G}(\mathop{}\!\mathrm{d}{g})\mathop{} \!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}\int_{\Omega}\pi^{N}(y)\Gamma_{\epsilon}(g^ {-1}x-y)\mathop{}\!\mathrm{d}{y}\gamma(x)\mu_{G}(\mathop{}\!\mathrm{d}{g}) \mathop{}\!\mathrm{d}{x}\quad\text{(since the Jacobian of $g$ is unitary)}$ $\displaystyle=\int_{\Omega}\int_{G}\int_{\Omega}\pi^{N}(g^{-1}y)\Gamma_{ \epsilon}(g^{-1}x-g^{-1}y)\mathop{}\!\mathrm{d}{y}\gamma(x)\mu_{G}(\mathop{}\! \mathrm{d}{g})\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}\int_{\Omega}\pi^{N}(g^{-1}y)\Gamma_{ \epsilon}(x-y)\mathop{}\!\mathrm{d}{y}\gamma(x)\mu_{G}(\mathop{}\!\mathrm{d}{g })\mathop{}\!\mathrm{d}{x}\quad\text{(due to the property of the heat kernel)}$ $\displaystyle=\int_{\Omega}\int_{\Omega}\int_{G}\pi^{N}(g^{-1}y)\mu_{G}( \mathop{}\!\mathrm{d}{g})\Gamma_{\epsilon}(x-y)\mathop{}\!\mathrm{d}{y}\gamma( x)\mathop{}\!\mathrm{d}{x}\quad$ $\displaystyle=\mathbb{E}_{S^{G}[\pi^{N}]\star\Gamma_{\epsilon}}\gamma.$ $$

∎

We decompose $\mathbf{d}_{1}(\pi,m(T))$ as follows

$$ $\mathbf{d}_{1}(\pi,m(T))\leq\mathbf{d}_{1}(\pi,\pi^{\epsilon})+\mathbf{d}_{1}( \pi^{\epsilon},m(T)).$ (32) $$

For the early stopping error, by the proof of Theorem 3.3 in [], we have $\mathbf{d}_{1}(\pi,\pi^{\epsilon})\leq C\sqrt{\epsilon}$, where $C$ only depends on the dimension $d$. To bound the second term in Eq. 32, we define $\eta^{\pi,\epsilon}:[0,T]\times R\mathbb{T}^{d}\to\mathbb{R}$ given by

$$ $\begin{cases}\partial_{t}\eta^{\pi,\epsilon}-\Delta\eta^{\pi,\epsilon}=0\,\, \text{in}\,\,R\mathbb{T}^{d}\times(0,T),\\ \eta^{\pi,\epsilon}(0)=\pi^{\epsilon}\,\,\text{in}\,\,R\mathbb{T}^{d}.\end{cases}$ (33) $$

Moreover, we define the drift

$$ $\mathbf{b}^{\pi,\epsilon}(x,t)\coloneq\nabla\log(\eta^{\pi,\epsilon})(x,T-t)$ $$

and let $m^{\epsilon}(x,t)=\eta^{\pi,\epsilon}(x,T-t)$ which satisfies

$$ $\begin{cases}\partial_{t}m^{\epsilon}=\Delta m^{\epsilon}+2\operatorname{div}( m^{\epsilon}\mathbf{b}^{\pi,\epsilon}),\\ m^{\epsilon}(0)=\eta^{\pi,\epsilon}(T).\end{cases}$ (34) $$

Then by applying Theorem 6, we have

$$ $\displaystyle\mathbf{d}_{1}(\pi^{\epsilon},m(T))$ $\displaystyle=\mathbf{d}_{1}(m^{\epsilon}(T),m(T))$ $\displaystyle\lesssim R^{\frac{3}{2}}(1+\sqrt{\left\lVert\mathbf{b}_{\theta} \right\rVert_{\infty}})\left(\mathbf{d}_{1}(m^{\epsilon}(0),\frac{1}{\text{vol }(R\mathbb{T}^{d})})+\left\lVert\mathbf{b}^{\pi,\epsilon}-\mathbf{b}_{\theta} \right\rVert_{L^{2}(m^{\epsilon})}\right).$ $$

By proposition A.3 in [], we have

$$ $\displaystyle\mathbf{d}_{1}(m^{\epsilon}(0),\frac{1}{\text{vol}(R\mathbb{T}^{d })})=\mathbf{d}_{1}(\eta^{\pi,\epsilon}(T),\frac{1}{\text{vol}(R\mathbb{T}^{d} )})\leq CRe^{-\frac{wT}{R^{2}}}\mathbf{d}_{1}(\pi^{\epsilon},\frac{1}{\text{ vol}(R\mathbb{T}^{d})}).$ $$

It remains to show the following bound

$$ $\left\lVert\mathbf{b}^{\pi,\epsilon}-\mathbf{b}_{\theta}\right\rVert_{L^{2}(m^ {\epsilon})}^{2}=\mathcal{J}_{D}(\eta^{\pi,\epsilon},\theta)\leq e_{nn}^{ \prime}=e_{nn}+C\left(1-\frac{\log\epsilon}{\sqrt{\epsilon}}+\frac{1}{\sqrt{T} }+T\left\lVert\mathbf{s}_{\theta}\right\rVert^{2}_{C^{2}(\Omega\times[0,T])} \right)\mathbf{d}_{1}(\pi^{N}_{G},\pi).$ (35) $$

In the rest part of this section, we prove Eq. 35. The proof is based on the structure of Section 8 in [].

We denote by $\rho^{m_{0}}:\Omega\times[0,T]\to[0,\infty)$ the solution to

$$ $\begin{cases}\partial_{t}\rho^{m_{0}}-\Delta\rho^{m_{0}}=0\,\,\text{in}\,\, \Omega\times(0,T],\\ \rho^{m_{0}}(0)=m_{0}\,\,\text{in}\,\,\Omega.\end{cases}$ (36) $$

###### Lemma 1 (Proposition 8.1 in [ 23 ] ) .

Let $m_{0}$ be a probability density in $\Omega$, such that $m_{0}\log(m_{0})\in L^{1}(\Omega)$ and $\rho:\Omega\times[0,T]\to\mathbb{R}$ be given by Eq. 36. Then we have

$$ $4\left\lVert\nabla\sqrt{\rho}\right\rVert_{2}^{2}=\int_{\Omega}m_{0}\log(m_{0} )-\rho(T)\log(\rho(T))\mathop{}\!\mathrm{d}{x}.$ $$

###### Lemma 2 (Proposition 8.2 in [ 23 ] ) .

Let $\pi^{i}$ ($i=1,2$) denote two probability measures in $\Omega$ such that $\left\lVert\pi^{i}\log(\pi^{i})\right\rVert_{1}<\infty$ and $\rho^{i}$ the corresponding solutions to Eq. 36. Then there exists a dimensional constant $C>0$ such that

$$ $\displaystyle\left|\mathcal{J}_{I}(\rho^{2},\theta)-\mathcal{J}_{I}(\rho^{1}, \theta)\right|\leq CT\sup_{t\in[0,T]}\mathbf{d}_{1}(\rho^{1}(t),\rho^{2}(t)) \left\lVert\mathbf{s}_{\theta}\right\rVert^{2}_{C^{2}(\Omega\times[0,T])}\leq CT \mathbf{d}_{1}(\pi^{1},\pi^{2})\left\lVert\mathbf{s}_{\theta}\right\rVert^{2}_ {C^{2}(\Omega\times[0,T])}.$ $$

###### Lemma 3 (Lemma 8.3 in [ 23 ] ) .

Let $\pi^{\epsilon}=\pi\star\Gamma_{\epsilon}$, and $\pi^{N,\epsilon}_{G}$ be as in Theorem 1 with $\epsilon<1$. There exists a dimensional constant $C=C(d)>0$ such that

$$ $\mathbf{d}_{1}(\pi^{N,\epsilon}_{G},\pi^{\epsilon})\leq\mathbf{d}_{1}(\pi^{N}_ {G},\pi),$ (37) $$

$$ $\left\lVert\pi^{N,\epsilon}_{G}-\pi^{\epsilon}\right\rVert_{1}\leq C\frac{ \mathbf{d}_{1}(\pi^{N}_{G},\pi)}{\sqrt{\epsilon}},$ (38) $$

and

$$ $\left\lVert\pi^{\epsilon}\log(\pi^{\epsilon})-\pi^{N,\epsilon}_{G}\log(\pi^{N, \epsilon}_{G})\right\rVert_{1}\leq C\left(1-\frac{d}{2}\log(\epsilon)\right) \frac{\mathbf{d}_{1}(\pi^{N}_{G},\pi)}{\sqrt{\epsilon}}.$ (39) $$

Moreover, let $\eta_{G}^{N,\epsilon}$ and $\eta^{\epsilon}$ be solutions to Eq. 36 with initial conditions $\pi_{G}^{N,\epsilon}$ and $\pi^{\epsilon}$ respectively. Then for large enough $T$ that depends on $R$ and the dimension $d$ but is independent of random samples or $N$, we have

$$ $\int_{\Omega}\log(\eta^{N,\epsilon}_{G}(T))\eta^{N,\epsilon}_{G}(T)-\log(\eta^ {\pi,\epsilon}(T))\eta^{\pi,\epsilon}(T)\mathop{}\!\mathrm{d}{x}\leq\frac{C}{ \sqrt{T}}\mathbf{d}_{1}(\pi,\pi^{N}_{G}).$ (40) $$

###### Proof.

Inequalities $(\ref{eq:8.3-1})-(\ref{eq:8.3-3})$ follow directly from the proof of Lemma 8.3 in []. For the bound in Eq. 40, by the convexity of the function $f(x)=x\log x$, we have

$$ $\displaystyle\int\log(\eta^{N,\epsilon}_{G}(T))\eta^{N,\epsilon}_{G}(T)-\eta^{ \pi,\epsilon}(T)\log(\eta^{\pi,\epsilon}(T))\mathop{}\!\mathrm{d}{x}$ $\displaystyle\leq\int\left(1+\log(\eta^{N,\epsilon}_{G}(T))\right)\mathop{}\! \mathrm{d}{(\eta^{N,\epsilon}_{G}(T)-\eta^{\pi,\epsilon}(T))}$ $\displaystyle\leq\left\lVert 1+\log(\eta^{N,\epsilon}_{G}(T))\right\rVert_{ \infty}\left\lVert\eta^{N,\epsilon}_{G}(T)-\eta^{\pi,\epsilon}(T)\right\rVert_ {1}.$ $$

From the proof of Lemma 8.3 in [], we have

$$ $\left\lVert\eta^{N,\epsilon}_{G}(T)-\eta^{\pi,\epsilon}(T)\right\rVert_{1}\leq \frac{C}{\sqrt{T}}\mathbf{d}_{1}(\pi^{N,\epsilon}_{G},\pi^{\epsilon})\leq\frac {C}{\sqrt{T}}\mathbf{d}_{1}(\pi^{N}_{G},\pi),$ $$

where $C>0$ is a dimensional constant. It remains to bound $\left\lVert 1+\log(\eta^{N,\epsilon}_{G}(T))\right\rVert_{\infty}$. Indeed, by the property of the heat kernel on $R\mathbb{T}^{d}$, $\eta^{N,\epsilon}(t)\lesssim_{d,R}1+(\epsilon+T)^{-d/2}$, and it is lower bounded by $\eta^{N,\epsilon}(t)\gtrsim_{d,R}(\epsilon+T)^{-d/2}$. By Proposition 3, we have $\inf_{x\in\Omega}\eta^{N,\epsilon}(x,t)\leq\eta^{N,\epsilon}_{G}(x,t)\leq\sup_
{x\in\Omega}\eta^{N,\epsilon}(x,t)$ for any $t$. This finishes the proof.
∎

###### Proof of Eq. 35 .

Note that $\mathcal{J}_{D}(\eta^{\pi,\epsilon},\theta)=\mathcal{J}_{I}(\eta^{\pi,\epsilon
},\theta)+4\left\lVert\nabla\sqrt{\eta^{\pi,\epsilon}}\right\rVert_{2}^{2}$. We have

$$ $\displaystyle\mathcal{J}_{D}(\eta^{\pi,\epsilon},\theta)$ $\displaystyle=\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\theta)+4\left(\left\lVert \nabla\sqrt{\eta^{\pi,\epsilon}}\right\rVert_{2}^{2}-\left\lVert\nabla\sqrt{ \eta_{G}^{N,\epsilon}}\right\rVert_{2}^{2}\right)+\left(\mathcal{J}_{I}(\eta^{ \pi,\epsilon},\theta)-\mathcal{J}_{I}(\eta_{G}^{N,\epsilon},\theta)\right).$ $$

By assumption we have $\mathcal{J}_{D}(\eta_{G}^{N,\epsilon},\theta)\leq e_{nn}$. By Lemma 1, we have

$$ $\displaystyle\left\lVert\nabla\sqrt{\eta^{\pi,\epsilon}}\right\rVert_{2}^{2}- \left\lVert\nabla\sqrt{\eta_{G}^{N,\epsilon}}\right\rVert_{2}^{2}$ $\displaystyle=\int_{\Omega}\pi^{\epsilon}\log(\pi^{\epsilon})-\pi^{N,\epsilon} _{G}\log(\pi^{N,\epsilon}_{G})\mathop{}\!\mathrm{d}{x}+\int_{\Omega}\eta_{G}^{ N,\epsilon}(T)\log(\eta_{G}^{N,\epsilon}(T))-\eta^{\pi,\epsilon}(T)\log(\eta^{ \pi,\epsilon}(T))\mathop{}\!\mathrm{d}{x}.$ $$

From Eq. 39 in Lemma 3, we can bound the first integral; while the second integral can be bound by Eq. 40. Combining with Lemma 2, we finish the proof.
∎

###### Proof of Corollary 1 .

Note that $\mathcal{M}$ is compact and can be covered by finitely many charts, where the map in each chart is Lipschitz (though with possibly different Lipschitz constant within each chart), so $\mathcal{M}$ has a Riemannian metric that is equivalent to the Euclidean metric in the ambient space. Hence we can apply the result in [] to $\mathbb{E}[\mathbf{d}_{1}(\pi^{N}_{G},\pi)]$.
∎

###### Theorem 6 (Wasserstein Uncertainty Propagation) .

###### Proof.

###### Corollary 2 .

###### Proof.

###### Proposition 3 .

###### Proof.

###### Lemma 1 (Proposition 8.1 in [ 23 ] ) .

###### Lemma 2 (Proposition 8.2 in [ 23 ] ) .

###### Lemma 3 (Lemma 8.3 in [ 23 ] ) .

###### Proof.

###### Proof of Eq. 35 .

###### Proof of Corollary 1 .

### 7.2 Proof that score-based generative models are intrinsically equivariant

###### Proof of Theorem 3 .

From [], it is known that score-based generative models are the solution of a mean-field game

$$ $\displaystyle\begin{dcases}\partial_{t}\rho+\nabla\cdot(\rho(f-\sigma^{2} \nabla U))=\frac{\sigma^{2}}{2}\Delta\rho\\ -\partial_{t}U-f^{\top}\nabla U+\frac{1}{2}|\sigma\nabla U|^{2}+\nabla\cdot f= \frac{\sigma^{2}}{2}\Delta U\\ U(x,T)=-\log\pi(x),\,\rho(x,0)=e^{-U(x,0)}.\end{dcases}$ (41) $$

Let $G$ be some group, $g\in G$ be an element of the group, and $A_{g}$ be the group action corresponding with $g$. Assume data distribution $\pi$ is $G$-invariant Then it is clear that $U(x,T)$ is also $G$-invariant as

$$ $\displaystyle U(gx,T)=-\log\pi(gx,T)=-\log\pi(x,t)=U(x,T).$ (42) $$

Furthermore, since $f$ is assumed to be $G$-equivariant, the corresponding Hamilton-Jacobi-Bellman equations are identical for all $g\in G$. Therefore, by the uniqueness
of the solution to the Hamilton-Jacobi-Bellman equation, $U(gx,t)=U(x,t)$ for all $t\in[0,T]$. For the existence and uniqueness of smooth solutions of the HJB equation and their properties we refer to [] (Section 1.7 and references therein), see also []. Therefore, the solution of the HJB equation $U(x,t)$ is invariant, and therefore the score function $\mathbf{s}=-\nabla U$ must be $G$-equivariant. Moreover, it is shown in [] that the minimizer of the implicit score matching objective, and therefore the ESM, is equivalent to the solution of 41. Therefore, this shows that the neural net must be parameterized in a way that is $G$-equivariant, thus incorporating an induced, equivariant (structural) inductive bias.
∎

###### Proof of Theorem 3 .

### 7.3 Proof of propositions of vector fields

#### G 𝐺 G italic_G -equivariance of S G E ⁢ [ 𝐬 ] superscript subscript 𝑆 𝐺 𝐸 delimited-[] 𝐬 S_{G}^{E}[\mathbf{s}] italic_S start_POSTSUBSCRIPT italic_G end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_E end_POSTSUPERSCRIPT [ bold_s ] .

For any $\bar{g}\in G$, we have

$$ $\displaystyle S_{G}^{E}[\mathbf{s}](\bar{g}x,t)$ $\displaystyle=\int_{G}A_{g}^{\top}\cdot\mathbf{s}(g\bar{g}x,t)\mu_{G}(\mathop{ }\!\mathrm{d}{g})$ $\displaystyle=\int_{G}A_{\bar{g}}A_{\bar{g}}^{\top}A_{g}^{\top}\cdot\mathbf{s} (g\bar{g}x,t)\mu_{G}(\mathop{}\!\mathrm{d}{g})$ $\displaystyle=\int_{G}A_{\bar{g}}A_{g\circ\bar{g}}^{\top}\cdot\mathbf{s}(g\bar {g}x,t)\mu_{G}(\mathop{}\!\mathrm{d}{g})$ $\displaystyle=A_{\bar{g}}S_{G}^{E}[\mathbf{s}](x,t).$ $$

###### Proof of Theorem 4 .

It is sufficient to look at the integration of $x$ over $\Omega$. We have

$$ $\displaystyle\int_{\Omega}\left(\left|\mathbf{s}\right|^{2}+2\nabla\cdot \mathbf{s}\right)S^{G}[\rho](x)\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}S_{G}\left[\left|\mathbf{s}\right|^{2}+2\nabla\cdot \mathbf{s}\right]\rho(x)\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\left|\mathbf{s}\right|^{2}\rho(x)\mathop{}\! \mathrm{d}{x}+2\int_{\Omega}S_{G}\left[\nabla\cdot\mathbf{s}\right]\rho(x) \mathop{}\!\mathrm{d}{x},$ $$

where the last equality is due to that the module $\left|\mathbf{s}\right|$ is $G$-invariant since $\mathbf{s}$ is $G$-equivariant.
For the second integral, we have

$$ $\displaystyle\int_{\Omega}S_{G}\left[\nabla\cdot\mathbf{s}\right]\rho(x) \mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}\sum_{i=1}^{d}\frac{\partial(\mathbf{s}_{i} )}{\partial x_{i}}(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}\rho(x)\mathop{}\! \mathrm{d}{x}$ $\displaystyle=\int_{G}\int_{\Omega}\sum_{i=1}^{d}\frac{\partial(\mathbf{s}_{i} )}{\partial x_{i}}(gx)\rho(x)\mathop{}\!\mathrm{d}{(x)}\mathop{}\!\mathrm{d}{ \mu_{G}(g)}$ $\displaystyle=\int_{G}\int_{\Omega}\sum_{i=1}^{d}\frac{\partial(\mathbf{s}_{i} )}{\partial x_{i}}(x)\rho(g^{-1}x)\mathop{}\!\mathrm{d}{(g^{-1}x)}\mathop{}\! \mathrm{d}{\mu_{G}(g)}$ $\displaystyle=-\int_{G}\int_{\Omega}\mathbf{s}(x)^{\top}(A_{g}\nabla\rho|_{g^{ -1}x})\mathop{}\!\mathrm{d}{(g^{-1}x)}\mathop{}\!\mathrm{d}{\mu_{G}(g)}\quad \text{(use integration by parts)}$ $\displaystyle=-\int_{G}\int_{\Omega}(A_{g}^{\top}\mathbf{s}(x))^{\top}(\nabla \rho|_{g^{-1}x})\mathop{}\!\mathrm{d}{(g^{-1}x)}\mathop{}\!\mathrm{d}{\mu_{G}( g)}$ $\displaystyle=-\int_{G}\int_{\Omega}(\mathbf{s}(g^{-1}x))^{\top}(\nabla\rho|_{ g^{-1}x})\mathop{}\!\mathrm{d}{(g^{-1}x)}\mathop{}\!\mathrm{d}{\mu_{G}(g)} \quad\text{(by the equivariance of $\mathbf{s}$)}$ $\displaystyle=-\int_{G}\int_{\Omega}(\mathbf{s}(x))^{\top}(\nabla\rho(x)) \mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}$ $\displaystyle=\int_{G}\int_{\Omega}(\nabla\cdot\mathbf{s})(x)\rho(x)\mathop{} \!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}$ $\displaystyle=\int_{\Omega}(\nabla\cdot\mathbf{s})(x)\rho(x)\mathop{}\!\mathrm {d}{x}.$ $$

Therefore, we have

$$ $\displaystyle\int_{\Omega}\left(\left|\mathbf{s}\right|^{2}+2\nabla\cdot \mathbf{s}\right)S^{G}[\rho](x)\mathop{}\!\mathrm{d}{x}=\int_{\Omega}\left( \left|\mathbf{s}\right|^{2}+2\nabla\cdot\mathbf{s}\right)\rho(x)\mathop{}\! \mathrm{d}{x}.$ $$

∎

To prove Proposition 1, we need the following lemma.

###### Lemma 4 .

For a generic $\rho\in\mathcal{P}(\Omega)$, which may not be $G$-invariant, the score formula of its symmetrized measure $S^{G}[\rho]$, is given by

$$ $\nabla_{x}\left[\log\left(S^{G}[\rho]\right)\right](x)=\frac{\int_{G}A_{g}^{ \top}\cdot(\nabla_{x}\rho)|_{gx}\mathop{}\!\mathrm{d}{\mu_{G}(g)}}{\int_{G} \rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}},$ $$

where $(\nabla_{x}\rho)|_{gx}$ is the gradient of $\rho$ evaluated at $gx$.

###### Proof of Lemma 4 .

$$ $\displaystyle\nabla_{x}\left[\log\left(S^{G}[\rho]\right)\right](x)$ $\displaystyle=\nabla_{x}\left[\log\left(\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{ \mu_{G}(g)}\right)\right]$ $\displaystyle=\frac{\nabla_{x}\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g) }}{\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}}$ $\displaystyle=\frac{\int_{G}\nabla_{x}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g) }}{\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}}$ $\displaystyle=\frac{\int_{G}A_{g}^{\top}\cdot(\nabla_{x}\rho)|_{gx}\mathop{}\! \mathrm{d}{\mu_{G}(g)}}{\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}}.$ $$

∎

###### Proof of Proposition 1 .

It suffices to prove the result for each time $t$, so we omit the time parameter. Let $\Omega/G$ be the quotient space of $\Omega$ by $G$. By the definition in Eq. 8, denoting by $\nabla\log\rho|_{gx}$ the score $\nabla\log\rho$ evaluated at $gx$, up to a multiplicative constant $C_{G}$ the depends on $G$ ($C_{G}=1$ if $\text{dim}(\Omega/G)<d$ and $C_{G}=\left|G\right|$ if $G$ is finite), we have

$$ $\displaystyle\mathcal{J}_{E}(\rho,\mathbf{s})$ $\displaystyle=C_{G}\int_{\Omega/G}\int_{G}\left|\mathbf{s}(gx)-\nabla\log\rho| _{gx}\right|^{2}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}\mathop{}\!\mathrm{d} {x}$ $\displaystyle=C_{G}\int_{\Omega/G}\int_{G}\left|A_{g}\cdot\mathbf{s}(x)-\nabla \log\rho|_{gx}\right|^{2}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}\mathop{}\! \mathrm{d}{x}$ $\displaystyle=C_{G}\int_{\Omega/G}\int_{G}\left|\mathbf{s}(x)-A_{g}^{\top} \cdot\nabla\log\rho|_{gx}\right|^{2}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)} \mathop{}\!\mathrm{d}{x},$ $$

where the last equality is due to the group actions in $G$ are isometries. For each $x\in\Omega/G$, regardless of $C_{G}$, we have

$$ $\displaystyle\nabla_{\mathbf{s}}\left[\int_{G}\left|\mathbf{s}(x)-A_{g}^{\top} \cdot\nabla\log\rho|_{gx}\right|^{2}\rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}\right]$ $\displaystyle=2\int_{G}\mathbf{s}(x)-A_{g}^{\top}\cdot(\nabla\log\rho|_{gx}) \rho(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}.$ $$

Then the stationary point of the above equation is given by

$$ $\mathbf{s}^{*}(x)=\frac{\int_{G}A_{g}^{\top}\cdot(\nabla\log\rho|_{gx})\rho(gx )\mathop{}\!\mathrm{d}{\mu_{G}(g)}}{\int_{G}\rho(gx)\mathop{}\!\mathrm{d}{\mu_ {G}(g)}}.$ $$

Note that $\nabla\log\rho|_{gx}=\frac{(\nabla_{x}\rho)|_{gx}}{\rho(gx)}$. This combined with Lemma 4 proves the claim.
∎

###### Proof of Theorem 5 .

It suffices to prove the equality for each time $t$, thus we will omit the time parameter. Expanding the square, it is equivalent to show that

$$ $\int_{\Omega}(\mathbf{s}^{\top}\nabla\log\rho)\rho(x)\mathop{}\!\mathrm{d}{x}= \int_{\Omega}\left(\mathbf{s}^{\top}S^{E}_{G}[\mathbf{s}]-\left|S^{E}_{G}[ \mathbf{s}]\right|^{2}+S^{E}_{G}[\mathbf{s}]^{\top}\nabla\log\rho\right)\rho(x )\mathop{}\!\mathrm{d}{x}.$ $$

First, we show that $\int\mathbf{s}^{\top}S^{E}_{G}[\mathbf{s}]\rho(x)\mathop{}\!\mathrm{d}{x}=\int
\left|S^{E}_{G}[\mathbf{s}]\right|^{2}\rho(x)\mathop{}\!\mathrm{d}{x}$. We have

$$ $\displaystyle\text{LHS}=\int_{\Omega}\int_{G}\mathbf{s}(x)^{\top}\cdot A_{g}^{ \top}\mathbf{s}(gx)\mathop{}\!\mathrm{d}{\mu_{G}(g)}\rho(x)\mathop{}\!\mathrm{ d}{x}$ $$

by the definition of the operator $S^{E}_{G}$; while

$$ RHS $\displaystyle=\int_{\Omega}\int_{G}\int_{G}\mathbf{s}(g_{1}x)^{\top}A_{g_{1}}A _{g_{2}}^{\top}\mathbf{s}(g_{2}x)\mathop{}\!\mathrm{d}{\mu_{G}(g_{1})}\mathop{ }\!\mathrm{d}{\mu_{G}(g_{2})}\rho(x)\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}\int_{G}\mathbf{s}(g_{1}x)^{\top}A_{g_{2} \circ g_{1}^{-1}}^{\top}\mathbf{s}(g_{2}x)\mathop{}\!\mathrm{d}{\mu_{G}(g_{1}) }\mathop{}\!\mathrm{d}{\mu_{G}(g_{2})}\rho(x)\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{G}\int_{G}\int_{\Omega}\mathbf{s}(g_{1}x)^{\top}A_{g_{2} \circ g_{1}^{-1}}^{\top}\mathbf{s}(g_{2}x)\rho(x)\mathop{}\!\mathrm{d}{x} \mathop{}\!\mathrm{d}{\mu_{G}(g_{1})}\mathop{}\!\mathrm{d}{\mu_{G}(g_{2})}$ $\displaystyle=\int_{G}\int_{G}\int_{\Omega}\mathbf{s}(x)^{\top}A_{g_{2}\circ g _{1}^{-1}}^{\top}\mathbf{s}(g_{2}\circ g_{1}^{-1}x)\rho(x)\mathop{}\!\mathrm{d }{x}\mathop{}\!\mathrm{d}{\mu_{G}(g_{1})}\mathop{}\!\mathrm{d}{\mu_{G}(g_{2})}$ $\displaystyle=\int_{G}\int_{G}\int_{\Omega}\mathbf{s}(x)^{\top}A_{g}^{\top} \mathbf{s}(gx)\rho(x)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)} \mathop{}\!\mathrm{d}{\mu_{G}(g_{2})}$ $\displaystyle=\int_{G}\int_{\Omega}\mathbf{s}(x)^{\top}A_{g}^{\top}\mathbf{s}( gx)\rho(x)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}=\text{LHS}$ $$

where the fourth equality is due to the $G$-invariance of $\rho$ and $A_{g}$ is unitary for any $g\in G$, and the fifth equality is due to that $G$ is unimodular so the Haar measure $\mathop{}\!\mathrm{d}_{\mu_{G}}$ is left-, right- and inverse-invariant.

Then it remains to show that $\int(\mathbf{s}^{\top}\nabla\log\rho)\rho(x)\mathop{}\!\mathrm{d}{x}=\int(S^{E
}_{G}[\mathbf{s}]^{\top}\nabla\log\rho)\rho(x)\mathop{}\!\mathrm{d}{x}$. Indeed, we have

$$ $\displaystyle\int_{\Omega}(S^{E}_{G}[\mathbf{s}]^{\top}\nabla\log\rho)\rho(x) \mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{\Omega}\int_{G}(A_{g}^{\top}\mathbf{s}(gx))^{\top}\mathop{ }\!\mathrm{d}{\mu_{G}(g)}(\nabla\log\rho(x))\rho(x)\mathop{}\!\mathrm{d}{x}$ $\displaystyle=\int_{G}\int_{\Omega}\mathbf{s}(gx)^{\top}A_{g}(\nabla\log\rho(x ))\rho(x)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}$ $\displaystyle=\int_{G}\int_{\Omega}\mathbf{s}(gx)^{\top}(\nabla\log\rho|_{gx}) \rho(x)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}$ $\displaystyle=\int_{G}\int_{\Omega}\mathbf{s}(x)^{\top}(\nabla\log\rho(x))\rho (x)\mathop{}\!\mathrm{d}{x}\mathop{}\!\mathrm{d}{\mu_{G}(g)}$ $\displaystyle=\int_{\Omega}\mathbf{s}(x)^{\top}(\nabla\log\rho(x))\rho(x) \mathop{}\!\mathrm{d}{x},$ $$

where the 3-rd equality is due to that $\nabla\log\rho$ is $G$-equivariant, and the 4-th equality is by a change of variable and $\rho$ is $G$-invariant.
∎

###### Proof of Proposition 2 .

Let $\Gamma=\text{Lip}_{1}(\Omega)$, and $\Gamma_{G}^{inv}$ be the subspace of $\Gamma$ that consists of $G$-invariant functions. By 1, actions in $G$ are 1-Lipschitz. Thus, $S_{G}[\Gamma]\subset\Gamma$.
First note that $S^{G}[\pi]=\pi$ since $\pi$ is $G$-invariant. Then we have

$$ $\displaystyle\mathbf{d}_{1}(S^{G}[\eta],\pi)$ $\displaystyle=\mathbf{d}_{1}(S^{G}[\eta],S^{G}[\pi])$ $\displaystyle=\sup_{\gamma\in\Gamma}\left\{\mathbb{E}_{S^{G}[\eta]}[\gamma]- \mathbb{E}_{S^{G}[\pi]}[\gamma]\right\}$ $\displaystyle=\sup_{\gamma\in\Gamma_{G}^{inv}}\left\{\mathbb{E}_{\eta}[\gamma] -\mathbb{E}_{\pi}[\gamma]\right\}$ $\displaystyle\leq\sup_{\gamma\in\Gamma}\left\{\mathbb{E}_{\eta}[\gamma]- \mathbb{E}_{\pi}[\gamma]\right\}=\mathbf{d}_{1}(\eta,\pi),$ $$

where the second equality is by the definition of $\mathbf{d}_{1}$ metric, and the third equality is due to Theorem 4.6 in [].
∎

###### Proof of Theorem 4 .

###### Lemma 4 .

###### Proof of Lemma 4 .

###### Proof of Proposition 1 .

###### Proof of Theorem 5 .

###### Proof of Proposition 2 .

## 8 Conclusion and future work

We rigorously show that SGMs can learn distributions with symmetries efficiently with equivariant score approximations. Compared to data augmentations, using equivariant vector fields for score-matching has the additional gain of reducing the score approximation error without the need to augment the dataset. Numerical experiments further verify this theoretical result. Certain directions are still unexplored in the present work.
For instance, it would be valuable to explore the architecture of equivariant neural networks to ensure they possess sufficient expressive power while maintaining a manageable number of parameters, as in the group equivariant convolutional neural networks proposed in [].
Furthermore, our analysis does not account for the time discretization of SGMs, and it could be interesting to incorporate this aspect or explore symmetry-preserving numerical integrators within the theoretical framework.
Another extension of our work would be to consider the domain as $\mathbb{R}^{d}$, with the forward process being, for instance, an Ornstein–Uhlenbeck process or other nonlinear processes [, ]. This would also require extending the empirical estimation of the Wasserstein-1 distance with group symmetry from bounded to unbounded domains.

## 9 Acknowledgement

Z. Chen, M. Katsoulakis, B. Zhang are partially funded by AFOSR grant FA9550-21-1-0354. M.K. is partially funded by NSF DMS-2307115 and NSF TRIPODS CISE-1934846.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 3
presentation: 3
questions: * Line 140: Clasically the projection $S_G$ is referred to as the Reynolds operator. 
* In Line (3) - doesnt $\nabla \text{log} \rho$ transform as $A^{\top}\nabla \text{log} \rho(Ax)$?
* Line 296 - it seems that the main point is the improved scaling of $1/d^*$, though not emphasised.
* Theorem 3 appears to be almost completely irrelevant? Could you help me understand why the result is either not obvious or expand on why exactly it is relevant? 
* Line 370 onwards I am not sure how the analogy with symplectic integrators fits into the picture there
* I beleive that remark 3 once again is entirely obvious without the theorems from the objective function? 
* Line 430 - this seems to completely miss the point of non-optimisability. Oftentimes, optimizing equivariant functions is harder/more expensive. And while $e_nn$ can be made as small as possible in theory, practice will be different.
* Figure 1 - I am confused why blue and orange are not the same exact curves - the loss itself must be identical? Even more surprisingly - why are red and green so close?
* I believe that the approach that you use for approximating the Wasserstein distance is not correct. It is fairly well known in the literature that WGANs significantly overestimate the wasserstein distances: https://arxiv.org/pdf/2103.01678
* Line 491 - I believe that the provided results are not extensive enough to claim neither that training a non-equivariant with augmentation nor equivariant without are superior. What is more, training an equivariant network on augmented data does not appear to make much sense, as the parameter gradient are the same in both cases?
soundness: 3
strengths: The paper is easy to read and provides a new theoretical anlysis for equivariant score-based generative models for learning symmetric distributions, providing valuable insights into their improved generalization and sample efficiency. The theoretical results are able to provide practical insights on comparison of equivariance with data augmentation. 

The paper presents a first-time quantitative comparison between data augmentation and the use of equivariant score parametrizations, demonstrating the advantages of the latter in terms of generalization error bounds.
summary: This paper presents the first theoretical analysis of score-based generative models (SGMs) for learning distributions with symmetries. It proves that incorporating symmetries through equivariant score approximations leads to improved generalization bounds and sampling efficiency. The study demonstrates that equivariant SGMs can effectively learn symmetric distributions without needing data augmentation and quantifies the impact of non-equivariant score parametrization on generalization error. The theoretical results are investigated with simple numerical experiments.
weaknesses: There are multiple weakness of this paper: 
Major: 
The numerical experiments, are conducted on extremely simple examples and the conclusions reported do not appear to follow directly from the example considered. Score based models are often used in much more complex, higher-dimensional datasets. I expand further on this in the questions sections. Furthermore, the main result of Section 4.1 - Theorem 3 seems completely irrelevant to the rest of the paper (or at least not sufficiently connected to it - happy to be proven wrong). 

Minor: 
The analysis focuses on symmetries representable by unitary matrices, significantly limiting its applicability. While the paper advocates for equivariant SGMs, it doesn't fully address the potential increase in computational complexity associated with designing and training equivariant neural networks.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. I understand that G is a given group when it comes to the theoretical results, but what exactly is the G group when it comes to the numerical results? For instance, in Section 6, is the group G the collection of four Gaussians or some properties of these four Gaussians (or other things)? How the elements of the example correspond to the elements of the theoretical parts need to be addressed in more detail. Specifically, I wonder if the authors can explicitly define the G group for the numerical example and provide a clear mapping between the theoretical elements and their counterparts in the experiment. 

2. In the Appendix, at lines 748-752, I do not understand how applying Theorem 6 yields the approximate inequality. If you already have the exact inequality in Theorem 6, why don't we just use the exact inequality (instead of some approximation)? I wonder if the authors can provide a step-by-step derivation of how they arrive at the approximate inequality from Theorem 6, or explain why an approximation is used if an exact inequality is available.
soundness: 3
strengths: The paper provides many theoretical results in a rigorous mathematical setting, when the target distribution has some symmetries. Furthermore, many of the results directly translate to practical use cases (e.g., Theorem 4).
summary: For learning distributions with certain symmetries using score-based generative models, this paper performs the first theoretical analysis and provides some guarantees. In particular, the authors show an improved error bound when the data distribution is group-invariant, use Hamilton-Jacobi-Bellman (HJB) theory to describe the inductive bias, and provide numerical tests to reinforce their claims.
weaknesses: While the (only) numerical example well demonstrates the theoretical results, it is a very simple one-dimensional experiment. It would have been much more convincing to see the theory validated with a higher-dimensional learning problem. For example, there are several high-dimensional experiments in Birrell et al. (2022) in ICML (cited by the present paper) and I wonder if the authors can validate their theoretical results in any of them.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - **Limited Scope of Numerical Experiments:** The numerical experiments are minimal and fail to comprehensively illustrate the necessity of incorporating equivariance into the SGM architecture. The four-Gaussian example appears overly simplistic, raising questions about the generality of the findings. Moreover, even in this simple example the benefit of using equivariance adjustments does not seem to clear. Could a more complex real-world example better justify the proposed framework? Also, is the bound evaluated and, if so, is it tight?
- **Choice of Architecture:** Why did the authors choose the current architecture over U-Net, which has been crucial to the success of SGMs due to its computational efficiency and architectural adaptability? Clarifying the rationale could provide insights into the model’s design choices.
- **Inconsistencies Between Theory and Experiments:** The authors constrain the diffusion process to a compact set but then use a Gaussian distribution in experiments, which does not align with this assumption. How should the bound be applied or evaluated in this case? Addressing this discrepancy would clarify the applicability of the results across different setups.
soundness: 3
strengths: The article demonstrates several notable strengths that advance the understanding and effectiveness of score-based generative models by focusing on equivariance:
- **Efficient Generalization of Existing Literature:** Building on the recent work by *Mimikos-Stamatopoulos et al. (2024)*, this article rigorously shows that $G$-invariance (group symmetry) is a key property that strengthens the structure and efficiency of SGMs. By extending these theoretical foundations, the authors highlight equivariance as crucial for improving model convergence and generalization.
- **Insightful Theoretical Framework Linking Key Concepts:** The paper effectively utilizes connections between MFG optimization, the HJB equation, and SGMs to clarify the structural properties of these models. This theoretical linkage reveals how SGMs can leverage inherent symmetries, providing a clearer perspective on the convergence bounds and making the role of symmetry more explicit.
- **Detailed Separation of Symmetry Components in Convergence Bounds:** The authors efficiently disentangle the components of symmetry within convergence bounds, offering a nuanced view of how these factors contribute to model stability and performance. This separation facilitates a better understanding of the role of symmetry in error decomposition and contributes to more accurate theoretical predictions.
- **Guidance for Model Enhancement:** The article provides practical suggestions for leveraging $G$-invariance directly within model architectures. This approach minimizes the need for data augmentation and improves computational efficiency while supporting more robust and generalizable models, laying groundwork for future research that optimally exploits symmetry.
summary: > This paper addresses how generative models can benefit from leveraging symmetries in data distributions to improve sampling efficiency and generalization on score-based generative models (SGMs). The authors focus on the empirical success of SGMs in scenarios where data has intrinsic symmetries (e.g., physics, molecular simulations). They argue that symmetries embedded in data distributions, known as group symmetries, can lead to improved model performance.

> First, they analyze how equivariant score-matching (with vector fields preserving symmetry) is an optimal method for training SGMs with symmetric data, bypassing the need for extensive data augmentation. Building on *Mimikos-Stamatopoulos et al. (2024)*, the authors generalize their bound errors in SGMs, breaking it down into errors stemming from non-equivariance, score-matching, sample complexity, and early stopping. They use the Wasserstein-1 distance (*d1*) as a metric to show how errors in these areas impact model performance, relying on HJB techniques and the mFG interpretation of SGMs.
The authors provides also numerical simulations aiming to corroborate their theoretical claims and to illustrate that embedding equivariant structures directly into SGMs consistently outperforms data-augmented models in generating high-quality samples of symmetric distributions, with lower Wasserstein-1 distance errors.

> This paper advances the theoretical understanding of symmetry in generative models, offering both conceptual and practical insights into how SGMs can be designed to better handle symmetric data distributions, efficiently using equivariant inductive biases instead of data augmentation.
weaknesses: - **Bounds Inefficiency:** The generalization bounds presented appear to lack efficiency, especially as they depend on the radius $R$ of the compact domain in which the data is confined. This leads to a bound that increases with the data’s spatial constraints. Additionally, the time horizon $T$ dependency in the bound further amplifies this issue, potentially limiting the practical applicability of the results.
- **Limited Novelty in Symmetry Component Analysis:** While the authors successfully separate the symmetry component in the bounds from *Mimikos-Stamatopoulos et al. (2024)*, the originality of this analysis seems limited. Revisiting classical results without substantial improvements in simulation outcomes weakens the impact of this contribution on advancing symmetry-related insights.
- **Insufficient Context for Symmetry in Examples:** Although the paper addresses how the absence of symmetry could affect convergence bounds, it lacks clear examples that naturally demonstrate symmetry or its absence. A more thorough description of practical cases where symmetry is essential would strengthen the authors’ arguments and provide clearer motivation.

Meta-review: This paper presents a theoretical analysis of score-based generative models (SGMs) for learning distributions with symmetries. The authors claim to provide theoretical guarantees for SGMs in this context, offering generalization bounds for group-invariant data distributions. They use Hamilton-Jacobi-Bellman theory to describe the inductive bias of equivariant SGMs and argue that equivariant vector fields can learn symmetrized distributions without data augmentation.

The paper attempts to provide a theoretical framework for understanding symmetry in SGMs and offers mathematical analysis of equivariant structures in generative models. However, it suffers from limited numerical experiments, with only a simple one-dimensional example. The theoretical results may have limited practical applicability, and there is a lack of comprehensive comparison with existing methods in high-dimensional settings. The relevance of some theoretical results to the main claims remains unclear.

The paper's contribution appears incremental, primarily extending existing work on SGMs and equivariance. The limited numerical experiments do not convincingly demonstrate the advantages of the proposed approach in real-world scenarios. While the theoretical analysis is detailed, its practical implications and relevance to state-of-the-art SGMs are not sufficiently demonstrated. Consequently, the paper falls short of providing a substantial advancement in the field of symmetry-aware generative models, which is expected for acceptance at ICLR.

Decision: Reject