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

Weaknesses

1. We consider the domain to be bounded, with $R$ fixed, and it should not be taken to $\infty$. As noted in our Conclusion and Future Work, applying the Ornstein–Uhlenbeck (OU) process to unbounded domains would be a more appropriate approach in such cases. Furthermore, in practice, $T$ never approaches $\infty$. Instead, it functions as a multiplicative factor in the sample complexity and does not constrain practical applicability.

2. We want to remind the reviewer that our contributions to Symmetry Component Analysis in Section 4 are new and independent of Mimikos-Stamatopoulos et al. (2024). Furthermore, Mimikos-Stamatopoulos et al. (2024) is not ``classical" but rather recent. We provide, for the first time, the quantification of errors in score-based generative modeling due to the approximation of equivariant score functions with non-equivariant vector fields.

3. We have provided abundant references in the Introduction and Related Works for the importance the motivation of encoding symmetries into generative models.

Questions

1. We refer the reviewer to the recent paper (Lu et al., 2024) cited in our paper for extensive experiments and implementations for equivariant SGMs, which also directly inspired our current theoretical study. We emphasize, however, that our primary contribution is to explain the existing widely observed empirical results and to provide theoretical justification for the use of equivariant models in contrast to data augmentation.

2. The design or the architecture of the model is not the focus of our work, and the parametrization of equivariant neural networks is not unique. The main contribution of our work is to quantitatively compare the role of equivariant models with that of data augmentation, and our simple simulation using fully-connected ReLU networks also provides a fair comparison. We also refer the reviewer to (Lu et al., 2024) for a few experiments of equivariant SGMs using U-Nets.

3. We assume the domain to be bounded for technical simplicity to provide the first quantitative comparison between model equivariance and data augmentation. Although the Gaussian example has unbounded support, it supports that our conclusion to advocate model equivariance is not limited to bounded domains. But for the theoretical part, we should use the OU process instead as mentioned in our future work and conclusion section.

Finally, the review has some vague language and inconsistencies, making it difficult to properly address them. We provide you with some examples.

• "Insightful Theoretical Framework Linking Key Concepts: 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.'' Our HJB theorem (Theorem 3) is different from the result for the convergence bounds (Theorem 1,2).

• "Bounds Inefficiency'' seems to us vague and appears to ignore our assumptions on bounded domains. As we assume the domain to be bounded, the bound should depend on $R$. Therefore, it does not make sense to take $R$ to infinity under our assumption: one should use the OU process instead, as mentioned in our future work section.

• "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." We do not use anywhere image-based data sets hence there is no need to consider U-nets or other specialized architectures. The problem we study is not about whether we should use U-Net or not, but rather whether we should use equivariant score approximations.

• "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." The paper Mimikos-Stamatopoulos et al. (2024) was just accepted at NeurIPS 2024 and is not quite yet a ``classic".

• The section "Detailed Separation of Symmetry Components in Convergence Bounds" in the Strengths seems to be somewhat contradictory to the section called ``Limited Novelty in Symmetry Component Analysis'' in the Weaknesses.

• The other Weakness, namely "Insufficient Context for Symmetry in Examples" appears to be rather vague and does not provide actionable items for us to address, given the already extensive discussion of the bibliography on related questions in the manuscript. We also refer the reviewer to our Global Response.

We thank the reviewer for the feedback and for carefully reading our manuscript. Please see our responses below.

Weaknesses

We did not conduct extensive experiments, as numerous related experiments have already been presented in (Lu et al., 2024) cited in our paper. Lu et al., 2024 has several similar experiments as in Birrell et al. (2022) in ICML, such as rotated MNIST, LYSTO, and ANHIR. Similar observations have also been made in earlier works on GANs and normalizing flows, where the generator/discriminator in GANs is equivariant/invariant, and the velocity field in CNFs is equivariant. However, common in all these approach is that they lack a theory to explain why model equivariance is more essential than data augmentation. One of our main contributions is to provide the first quantitative explanation for why model equivariance is more critical than simply performing data augmentation, in the context of SGMs.

Questions

1. The group used in the numerical experiments is $C_4$, the rotational group, represented by rotation matrices of $0^\circ$, $90^\circ$, $180^\circ$, and $270^\circ$. Clearly, the mixture of four Gaussians is invariant for any combination of these rotations. To be specific, the representations of the group actions are

$

A_0 = \begin{bmatrix}
   1 & 0 \\\ 0 & 1
\end{bmatrix},  A_{90} = \begin{bmatrix}
   0 & -1 \\\ 1 & 0
\end{bmatrix}, A_{180} = \begin{bmatrix}
   -1 & 0 \\\ 0 & -1
\end{bmatrix}, A_{270} = \begin{bmatrix}
   0 & 1 \\\ -1 & 0
\end{bmatrix}

$

2. We assume the reviewer refers to the symbol '$\lesssim$' at Line 751. If that is the case, it is a notation, and we use it to mean up to a dimensional constant $C=C(d)>0$ in Theorem 6, and it is equivalent to rewriting it as $\leq C$ by putting $C$ back into Line 751, essentially a direct application of Theorem 6. We will add a comment about that in the proof to clarify the notation.

We appreciate the reviewer’s recognition of the strengths of our paper, particularly in acknowledging that 1) it provides the first theoretical analysis of score-based generative models (SGMs) for learning distributions with symmetries, and 2) it offers a first-time quantitative comparison between data augmentation and the use of equivariant score parameterizations.

However, we respectfully disagree with most of the concerns raised, as we elaborate next. A key contribution of our paper lies in the quantitative comparison of building equivariance directly into the model, through equivariant score approximations, versus relying solely on data augmentation. We hope our clarifications below address the reviewer’s concerns and further highlight the significance of our contributions.

Weaknesses

1. We did not conduct extensive experiments, as numerous related experiments have already been presented in (Lu et al., 2024) cited in our paper. Similar observations have also been made in earlier works on GANs and normalizing flows, where the generator/discriminator in GANs is equivariant/invariant, and the velocity field in CNFs is equivariant. Our primary focus is twofold: 1) to provide the first rigorous explanation of why model equivariance is fundamentally more impactful than merely applying data augmentation, specifically in the context of SGMs, and 2) to quantitatively analyze the resulting improvements in the generalization bound. The numerical example included in our work serves to validate the theoretical findings.

2. Theorem 3 rigorously identifies the intrinsic inductive bias of score-based generative models, which justifies the importance and necessity of using $G$-equivariant vector fields using the fundamental characterization of SGMs as a mean-field game. This result is new and not provided in previous symmetry-related works. The main focus of our work is to rigorously justify the necessity and importance of building symmetry into SGMs, and therefore Theorem 3 is a key result in providing a rigorous foundation for this goal.

3. The characterization that our analysis only applies to groups represented by unitary matrices is categorically false. We emphasize, again, that Theorem 3 is valid for any group $G$ as it does not rely on a unitary representation, and shows that the inductive bias of SGMs states that a $G$-invariant distribution demands a $G$-equivariant vector field. This fact is true for any group $G$. Furthermore, unitary representations are not a limitation, as they encompass important transformations such as rotations and reflections. For the subsequent results, we can assume $G$ is compact, which ensures the existence of a unitary representation by leveraging a change of variables through the Peter–Weyl Theorem.

4. The computational complexity associated with designing and training equivariant neural networks is beyond the scope of our work. However, we refer the reviewer to Group Equivariant Convolutional Networks (G-CNNs) [Cohen & Welling, ICML 2016] and Steerable CNNs [Cohen & Welling, ICLR 2017] to address the concerns regarding the efficiency of equivariant neural networks.

Questions

1. Thank you for the related notion. We will mention this in the revision.

2. Equation (3) is the definition of being $G$-equivariant.

3. Thank you, we will emphasize the improvement in the convergence rate that depends on the intrinsic dimension. This important improvement is only possible because we are using the $\mathbf{d}_1$ metric in our generalization bounds.

4. Theorem 3 justifies the importance of using $G$-equivariant vector fields from the mean-field perspective. This result is new and not provided in previous symmetry-related works. As the main focus of our work is to rigorously justify the necessity and importance of building symmetry into the model, we do not believe that Theorem 3 is irrelevant. On the contrary, it provides a rigorous explanation of why the equivariance of the vector field is an intrinsic inductive bias, and therefore should be included in the algorithms. Theorem 3 is applicable for any group $G$ without any restrictions.

5. In practice we only have access to finite samples drawn from the unknown distribution, and those finite samples do not preserve the symmetry even if the underlying distribution does have symmetry. Therefore, the practical numerical schemes should incorporate symmetry so that the outcome would respect the symmetry. ``Symplectic numerical integration'' is one of the well-known examples in numerical analysis that has a similar idea that the numerical scheme should have the structure-preserving properties inherited from the continuous problem, such as Hamiltonian systems and energy conservation.

6. This is an important property of score-matching with equivariant vector fields, that you essentially gain data augmentation for free by simply using equivariant vector fields. Therefore, it is not clear to us why Remark 3 is mathematically obvious or known without the theorems from the objective function.

7. We want to clarify that we do not assume $e_{nn}=0$. In fact, quite the opposite, we rigorously show in Theorem 5 that the score approximation error has a lower bound, quantified by the DFE when using a non-equivariant score approximation (NN), even assuming perfect training and optimization. The comment of the reviewer ``optimizing equivariant functions is harder/more expensive'' is widely understood and addressed in the existing literature. In particular, while this is not the focus of our work, this statement contradicts many existing works related to equivariant neural networks, which are more efficient and easy to implement, such as G-CNNs. Also, the paper (Lu et al., 2024) cited in our paper provides numerous practical implementations.

8. The blue and orange curves are not expected to overlap: one represents results with data augmentation on the training samples, while the other does not, and both are plotted using 25 replicates. Since the samples are drawn randomly, some variation is naturally expected. As for the observation, “Even more surprisingly—why are the red and green curves so close?”, this outcome is entirely consistent with our quantitative discussion in Section 5, which builds on our theoretical results. Specifically, data augmentation offers limited benefits when the network is non-equivariant, as it cannot reduce the DFE error in such cases. This behavior is independent of data augmentation itself. Overall, Figure 1 supports our theoretical conclusion that leveraging equivariant score approximations is fundamentally more impactful than relying solely on data augmentation.

9. We want to correct the reviewer that our approximation of Wasserstein-1 distance $\mathbf{d}_1$ is different from WGANs. First, we do not train a generator. Second, we do not use mini-batches when computing the approximate $\mathbf{d}_1$ between generated and true samples. Moreover, we do not use gradient penalty but rather apply spectral normalization, which offers a type of ``neural net distance'' [R1] that approximates $\mathbf{d}_1$.

[R1] Arora, S., Ge, R., Liang, Y., Ma, T. and Zhang, Y., 2017, July. Generalization and equilibrium in generative adversarial nets (GANs). In International conference on machine learning (pp. 224-232). PMLR.

10. Line 491: our simulation results are not extensive, but they corroborate our theory. Training an equivariant network on augmented data is part of our ablation study that makes the comparisons in Figure 1 complete. More empirical results can be found in (Lu et al., 2024) cited in our paper. Training equivariant NNs on augmented data makes sense for the completeness of the experiments. We also want to remind the reviewer that ``the parameter gradient are the same in both cases" is not true in general, as it also depends on the loss function. Your statement is actually supported by our Theorem 4.

We thank the reviewer for the quick response and appreciate the discussion. Please see our responses to your two sets of comments below.

W1: We want to clarify that we do not claim "G-CNNs would be 'better' than normal CNNs" generically: for example, we should not expect a G-CNN with very simple architecture to perform better than a normal CNN with a much larger depth and number of channels and filters. What we claim through the logic of our theory is that: whenever you have a non-equivariant vector field, you should use its Reynolds averaging, meaning you should always use an equivariant vector field or NN. However, within the space of equivariant vector fields or the family of equivariant NNs, Reynolds averaging might not be the most efficient way to build the equivariant models, and we leave space for any design of equivariant NNs, e.g., G-CNNs are equivariant and are not designed by averaging. Besides, we also want to emphasize that the specific design and training of equivariant NNs are not the focus of our work. We think this will be a good clarification, and will add it to the revision.

With a change of variables, we should define all the forward-backward processes and the score-matching objectives in the new coordinate systems, in which the matrix representation of $G$ is unitary. Then all our analysis can be applied to the new coordinate system. Finally, you can always reverse the transformation by the change of variables on the generated samples.

W4: The training/expense of equivariant NNs is not the focus of our work and is a complementary task to our paper. However, we want to refer the reviewer to the empirical results in G-CNNs [Cohen & Welling, ICML 2016]: for example, in Section 8.1 Rotated MNIST, they empirically compare the normal CNN with G-CNNs with approximately the same number of parameters and use the Adam algorithm for optimization, and the result illustrates that G-CNNs achieve smaller errors (almost halves the error) than normal CNNs.

6. Yes, it is for all time $T$.

7. "Equivariant models with the same error in training would be better than their non-equivariant counterparts" -- Your conclusion is correct regarding Eq. (21), but it is not the only possible conclusion. Another interpretation is what we explained in the response to point W1 above.

We meant "optimizing equivariant functions is harder/more expensive" contradicts the empirical results in the existing literature. See also our response to W4. However, the training/expense of equivariant NNs is not the focus of our work.

8. "Thus, the only case in which they would not overlap - is due to stochasticity in batch selection. If that were the case - the error bars would at least intersect. However, they do not - meaning that there must be another systematic - my guess would be the wasserstein distance approximation being incorrect." -- The stochasticity comes also from random samples drawn from the target distribution. The error bars of the blue and orange curves (both using equivariant NNs) indeed intersect a lot of times. To be specific, the blue and orange curves are the bottom two curves in the range between $10^3$ and $10^4$ for $N$. On the other hand, the error bars of the red and green curves (both using non-equivariant NNs) also intersect a lot of times.

"In regards to the non-equivariant models - your theory does not analyze the non data-augmented case, right?" -- Our analysis is built upon the non-equivariant case with no data augmentation (Mimikos-Stamatopoulos et al., 2024) cited in our paper.

"So i am not sure how you claim that it makes sense? The non-augmented model should not be able to generate a quarter of all possible samples, thus the support of augmented is extremely different from the support of the non-augmented one - how could it make sense that the distances agree?" -- To clarify any misunderstanding about data augmentation, here is our explanation. In our experiment, the target distribution is a mixture of 4 Gaussians. Even with no data augmentation, you can still draw samples from the four centers and it is unlikely your samples are all clustered at one center. As a result, even use a non-equivariant NN, you can still generate a distribution from which you can observe four centers. Regarding the theorem, our Theorem 4 shows that the optimal equivariant vector field can be obtained by score-matching for raw training data without data augmentation whenever you apply equivariant NNs. For example in our Figure 2, in which we use 40 training samples, the probability of missing one quarter is $4\times0.75^{40}\approx 10^{-5}$.

For responses to 9. and 10., see the next block.