Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance

Summary

The paper proposes a probabilistic version of the canonization method to construct equivariant architectures from generic non-equivariant backbones. This approach is inspired by the symmetrization solution, which involves averaging a non-equivariant function to obtain an equivariant one, but replaces the uniform averaging distribution with an input-conditional one to achieve a lower sample complexity.

Strengths

The paper is clearly written and well motivated. The proposed idea overcomes some important limitations of both the symmetrisation and the canonicalisation methods. I liked the idea of using a general-purpose transformer architecture, which makes the proposed method very flexible and easy to apply in more generic contexts. The surprising performance of pre-trained vision models on non-vision tasks is also a very interesting result and I wonder if it can tell us something more general about these architectures. Despite the generality and flexibility of the approach, the experiments show it can also achieve competitive performance.

Weaknesses

With respect to the canonicalization method from [26], the proposed approach essentially averages over $N>1$ samples and, therefore, is $N$ times more expensive than canonicalization. Since this work is partially motivated by computational efficiency, I think this aspect has not been sufficiently discussed in the manuscript. Moreover, I imagine the number of samples chosen can play an important role on 1) the model's final performance, 2) the training dynamic, 3) the variance of the learned distributions and 4) the training time. I think this point also requires some additional discussion. I would encourage the authors to include a short experiment to study the effect of the sample size on (at least some of) these aspects.

Questions

Why is it necessary to assume the group is compact? line 96-97: "2) learn to produce low-variance samples $g \sim p_\omega(g|x)$ that can offer more stable gradients for the base function $f_\theta$ during training." I think this claim requires some additional comments to justify it. Why should the network learn to produce low-variance samples? On a related note, do you train the models with some low-variance regularization? what prevents the model from outputting always a uniform distribution? lines 114-117: this seems more a problem of parameterizing the group G itself, not its representations. The representation $\rho(g)$ can be computed in a differentiable way from $g$ when $G$ is a Lie group at least. Eq. 8: it is not really clear how you output a permutation matrix. You use the notation $P_g \approx \hat{P}_g$ but you don't explain what it means. The matrix $\hat{P}_g$ is a real valued matrix and doesn't belong to $(0, 1)^{n \times n}$. Do you apply rounding on it? If so, isn't it possible to just use the matrix $\hat{P}_g$? Is there any benefit in using the boolean $P_g$ instead? You compare your probabilistic method with the deterministic canonicalization method. It is not clear to me if the benefit is coming from the fact the backbone is averaged over $N>1$ samples or if the stochasticity of $p_\omega$ plays a role. Could you try a version of your model which outputs $N$ samples but deterministically (e.g. sample the noise variables only once before training and keep them fixed)? While this is already mentioned by the authors among the limitations as a future work, it would be interesting to understand if this kind of architectures enjoys the same data-efficiency property of other equivariant networks. Do the authors have any insights about the effect of the training set size on the ability of the model to learn the symmetrising distribution? The paper is missing a related works section. Moreover, I think an important part of the literature about equivariant networks using group-convolution is not cited.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

Some limitations are discussed in the supplementary material. I encourage the authors to move the main points mentioned there in the main paper. See also my comments under Weaknesses.

Summary

The paper presents a method to learn to symmetrize a neural network using data. The method considers a learnable probability distribution over the group and uses group averaging to enforce (relaxed) equivariance w.r.t. the learned distribution.

Strengths

Over recent years there has been a growing interest in the community to learn relaxed symmetries (invariance and equivariance) from data. The paper addresses this important problem and provides a practical method that scales to interesting model classes and data problems (e.g. the use of Transformers, particle dynamics, and graphs is noteworthy).

Weaknesses

* Missed prior work There have been several approaches that aim to learn symmetrizing distributions (invariance or equivariance). Placing a probability distribution on the group over which a function is averaged and learning this distribution from data is not new. For instance, see approaches listed in survey Sec. 6 of [2]. * Proposed method already exists The paper misses this literature and it seems that the method has been published before (Equivariant Augerino, Sec. 3.1 of [1]). An argument could be made that although the Augerino paper describes the method for general groups, it only considers very simple affine groups in practice. This paper does consider more interesting group structures and domains (e.g. graphs). However, if the same method has been published before and is applied to new domains/data, this should at least be credited and discussed. This would make it more clear what the contributions of the work are. * Objective function A well-known difficulty with learning symmetry distribution is not necessarily describing a probability distribution on the group but rather the objective to learn the distribution over the group. Although the method claims to be probabilistic, it is not clear what objective is used to train the model (or whether inference is being performed). If regular maximum likelihood loss is used, it seems there is no encouragement in the objective that prevents the collapse of the symmetrization distribution into a delta peak at the identity. In such a case, it would result in not learning the symmetrization that generalises well. Prior works often consider more sophisticated losses, such as regularization [1], lower bounds or model selection, to prevent this. The paper seems to skip over this issue. [1] Benton, Gregory, et al. "Learning invariances in neural networks from training data." Advances in neural information processing systems 33 (2020): 17605-17616. [2] Rath, Matthias, and Alexandru Paul Condurache. "Boosting deep neural networks with geometrical prior knowledge: A survey." arXiv preprint arXiv:2006.16867 (2020).

Questions

a) How does the method differ from equivariant Augerino (Sec. 3.1 of [1])? b) What objective is used, or how is inference over the probability distribution being performed? c) If a maximum likelihood objective is used, what mechanism prevents the symmetrization from collapsing into the identity? Did authors notice such a collapse in experiments, or is this mitigated somehow? [1] Benton, Gregory, et al. "Learning invariances in neural networks from training data." Advances in neural information processing systems 33 (2020): 17605-17616.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

- The paper seems to miss prior works that consider learning similar parameterizations over group structures. - The proposed method seems to be proposed already (equivariant Augerino, Sec. 3.1 of [1]) - It is not clear what objective is being used to train the model and what prevents collapse of the symmetrization distribution.

Summary

The paper suggests a probabilistic approach to symmetrization, where an input conditional distribution is used to replace the untractable haar measure distribution of infinite groups. In turn, the paper identifies what are the conditions on the conditional distribution under which the symmetrization yields an equivariant function. The method is evaluated on several benchmarks involving permutation and rotation symmetries.

Strengths

Overall the paper is well-written and easy to follow. The formulation of the method seems to be adequate. The proposed idea is simple, and a natural extension of previous methods.

Weaknesses

Relation to previous works The paper in the introduction states that existing approaches focus on either manually deriving smaller subsets [1], or implementing a relaxed version of equivariance [2]. It is not clear what manually means in this context. In fact, [1] identifies a similar condition to eq(5) when the function p(g|x) is a delta function. I noticed section 2.4 relates to these previous works more thoroughly. More importantly, I feel [2] and the suggested approach are very similar, as [2] does not only suggest an approximate equivariant but rather suggests a learnable frame function (group subset), satisfying the equivariance constraint (similar to eq (5)). Thus, I feel this work should convince us of the necessity of the probabilistic model for learnable frames. To this end, I would have expected an ablation study testing “apples-to-apples” architectures where the only difference is the addition of the input noise term. It would also be beneficial to separate the test of averaging over S_n to O(3) in the experiment in table (2) and use symmetrization only to the O(3) symmetries with S_n equivariant networks. Missing discussion In fact, the network prediction in eq(4) is only approximated. Thus I assume, the loss terms used for training only estimate the true loss (gradients). How stable are these estimates? How well do they approximate the true loss?

Questions

I would appreciate any response from the authors regarding the weakness stated above.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Could not find a discussion on limitations

Summary

This paper proposes a new symmetrization method for achieving equivariance for any base function. It absorbs previous methods with the same objective as special cases, including group averaging, frame averaging, and canonicalization. The method uses a learnable, equivariant map to generate the group representation necessary for achieving equivariance from an invariant, external noise variable, followed by necessary post-processing to ensure the validity of the group representation. The authors empirically demonstrate that the learned stochasticity in symmetrization leads to improved model performance than the previous methods, and showcase the benefits of general pre-training for equivariant models.

Strengths

* The proposed method is general, providing a unified perspective for previous methods and may lay the foundation for further theoretical exploration of the "optimal" conditional group distribution. * The presentation is clear and easy to read. * The finding of the benefit of non-symemtric pretraining for equivariant models seems interesting and new.

Weaknesses

* My primary concern is that the supriorty of the proposed method compared to the canonicalization method is not clear enough. - The advantages of the former over the latter seem to be more evident only in the first experiment, as the canonicalization method appears to have difficulties during optimization (although it is unknown whether this issue can be mitigated through some warm-up techniques). In the other two experiments, the performance of both methods is quite similar. - On the other hand, compared to the canonicalization method, the proposed method incurs more computational cost in practice, including the additional parameters and hyperparameters in the learnable module, as well as the sampling cost (the normalization method only requires a single "sample", whereas here we need 10-20 samples). + This paper lacks thorough empirical and theoretical analyses for the two claimed superiorities over previous methods. Specifically, - "Learn to collaborate with base function to maximize task performance": No sensitivity analysis is provided to demonstrate how this method adapts better to different base functions. - "learn to produce low-variance samples that can offer more stable gradients for the base function": No analysis of gradient magnitudes is provided to indicate its stability. * The limitation is not discusses, including the extra computational cost. * The post-process seems interesting but also laborious, as it requires different designs for each different (compact) groups.

Questions

* What's the computational cost of this method, including training and testing, compared to the canonicalization method? * Analyzing the sampling efficiency for estimating expectations would be beneficial, especially in comparison to the uniform distribution, as it is a major claim of the effectiveness of this method.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors have not discussed the limitations.

Summary

This work presents a generic approach to symmetrize wide range of base models. Instead of relying on uniform average sampling, the approach introduces a trainable transformation to model the group equivariant distribution. The framework theoretically encompasses existing approaches such as group averaging, frame averaging, and canonical function approaches. Furthermore, the author provides empirical evidence demonstrating competitive performance.

Strengths

1. The proposed method demonstrates the ability to symmetrize architectures in a group-agnostic manner for general purposes, supported by sound theory. 2. The theoretical analysis presented in this work showcases the ability of the proposed approach to encompass interesting literature that assigns distributions on the compact group $G$. 3. The paper is written in a reader-friendly manner, making it easy to follow.

Weaknesses

1. A thorough discussion of the empirical and theoretical advantages and disadvantages of literatures, specifically group average, frame average, and canonicalization, would greatly enhance our understanding, given the shared problem setup with the proposed method. 2. I have reservations about the assertion made in Line 3 that >we use an arbitrary base model (such as an MLP or a transformer)... While the argument of this work suggests that the base model $f_{\theta}$ can be arbitrary, in the experiments, only MLP and transformer architectures were explored and evaluated. 3. Typically [1,2], showcasing improved performance on image classification problems is one of the common applications used to demonstrate the effectiveness of a newly proposed equivariant network. It is encouraged to include such experiments in order to comprehensively validate and demonstrate the efficacy of the proposed equivariant network. [1] S. Basu, P. Sattigeri, K. N. Ramamurthy, V. Chenthamarakshan, K. R. Varshney, L. R. Varshney, and P. Das. Equi-tuning: Group equivariant fine-tuning of pretrained models [2] S. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonicalization functions

Questions

1. Is it feasible to explore the convergence of the probabilistic equivariant distribution $p_{\omega}$ and assess its dissimilarity to a uniform distribution? This analysis would provide insights into the nature of the mechanism. 2. What is the sensitivity of the proposed method to the architecture scales of MLP and transformer? Can the efficacy of the proposed method be maintained when the base models have a large number of parameters? 3. Is there a general framework or guideline for designing the $G$ equivariant neural network $q_{\omega}$? How significantly does the expressivity of $q_{\omega}$ impact the performance? 4. Why some experimental comparisons between [1-3] are missing? For instance, in ``3.1 Graph Isomorphism Learning with MLP'', what limits the comparison with [3]? [1] S. Basu, P. Sattigeri, K. N. Ramamurthy, V. Chenthamarakshan, K. R. Varshney, L. R. Varshney, and P. Das. Equi-tuning: Group equivariant fine-tuning of pretrained models [2] S. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonicalization functions [3] O. Puny, M. Atzmon, E. J. Smith, I. Misra, A. Grover, H. Ben-Hamu, and Y. Lipman. Frame averaging for invariant and equivariant network design

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have discussed limitations and potential societal impact associated with this work.

Official Review by Reviewer SgC2

The authors have addressed all my questions and concerns well. Thus, I increased my evaluation.

Additional In-Depth Analysis on Convergence and Sample Sizes (1/2)

Dear Reviewers, We sincerely appreciate insightful and constructive comments from all reviewers. Following common requests (SgC2, Ry8Q, wgQS, SdiQ), we have conducted an in-depth empirical analysis of the convergence and sample complexity of our method using the EXP-classify dataset (Sn invariance; Section 3.1). We want to share our main findings that might be of general interest. Since we cannot share the figures directly in the discussion period, we used the tables to present the results. We will update our draft to reflect our new findings with better presentations. Let us recap that our algorithm uses an equivariant distribution p(g|x) and a base function f(x) to obtain a symmetrized function ϕ(x) = E_g[ g⋅f(g-1⋅x) ].       ---       **Q1. Does p(g|x) learn to produce low-variance samples compared to Unif(G)? (SgC2, Ry8Q, SdiQ)** Experiment: For every few training epochs, we sampled $N = 100$ permutation matrices $\mathbf{P}\_{g} \sim$ p(g|x) for each validation data $\mathbf{x}$ and obtained their average $\bar{\mathbf{P}}\_{\mathbf{x}} = \sum \mathbf{P}\_{g} / N$. Then, we measured the mean of row- and column-wise entropy of $\bar{\mathbf{P}}\_{\mathbf{x}}$ for quantitative analysis. If the distribution p(g|x) has learned low-variance samples, the entropy of $\bar{\mathbf{P}}\_{\mathbf{x}}$ is expected to be low. The results are summarized in Table A. Note that a similar analysis is visually presented in the Figure 1 of the main text. **Table A.** Entropy of $\bar{\mathbf{P}}\_{\mathbf{x}}$ averaged over validation dataset $\\{\mathbf{x}\\}$. | Train epoch | 100 | 500 | 1000 | 1500 | 2000 | | --- | --- | --- | --- | --- | --- | | p(g\|x) | 6.25 | 6.04 | 5.93 | 5.84 | 5.71 | | Unif(G) | 7.03 | 7.03 | 7.03 | 7.03 | 7.03 | Observation: The entropy reduces throughout training and is consistently smaller than that of Unif(G). **A1. In EXP-classify, p(g|x) learns to produce lower-variance samples than Unif(G).**       ---       **Q2. Does p(g|x) offer more stable gradients for the base function than Unif(G)? (Ry8Q, SdiQ)** Experiment: For every few training epochs, we computed the full gradient of the loss over the entire training dataset with respect to the parameters of the base function f(x). This averages out the variance from individual data point and provides the net direction of the gradient on the base function offered by p(g|x) or Unif(G). For a controlled comparison, we fixed the initialization of the base function to be identical in both settings. **Table B.** Norm of the full gradient w.r.t. the parameters of the base function trained with p(g|x) or Unif(G). | Train epoch | 100 | 500 | 1000 | 1500 | 2000 | | --- | --- | --- | --- | --- | --- | | p(g\|x) | 1.13 | 1.12 | 0.34 | 2.00 | 3.40 | | Unif(G) | 1.35e-2 | 4.27e-5 | 1.74e-4 | 2.44e-4 | 1.65e-4 | Observation: When we use p(g|x) for symmetrization, the net gradient for the base function has a consistently larger magnitude in contrast to Unif(G) that leads to near-zero net gradients. This indicates, for Unif(G), the gradients from each training data instances are oriented in a largely divergent manner and therefore the training signal is collectively not very informative, while using p(g|x) leads to more consistent gradient across training data instances, *i.e.*, it offers a more stable training signal. **A2. In EXP-classify, p(g|x) offers more stable gradients for the base function.**       ---       **Q3. Does p(g|x) offer better sample efficiency (lower variance estimation) than Unif(G)? (Ry8Q, SdiQ)** Experiment: We fixed a randomly initialized base function f(x) and computed the variance of the estimated output ϕ(x) and loss L(ϕ(x), y) with 100 repeated measurements for a range of inference sample sizes. **Table C.** Variance of estimated output ϕ(x) under p(g|x) and Unif(G). Scaled by 1e7 for readability. Lower variance for each inference sample size is boldfaced. | Inference sample size | 1 | 2 | 5 | 10 | 20 | 50 | 100 | 200 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | p(g\|x) | **1.65** | **0.81** | **0.33** | **0.17** | **0.083** | **0.033** | **0.016** | **0.008** | | Unif(G) | 3.05 | 1.53 | 0.63 | 0.31 | 0.16 | 0.062 | 0.030 | 0.015 | **Table D.** Variance of estimated loss L(ϕ(x), y) under p(g|x) and Unif(G). Scaled by 1e7 for readability. Lower variance for each inference sample size is boldfaced. | Inference sample size | 1 | 2 | 5 | 10 | 20 | 50 | 100 | 200 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | p(g\|x) | **6.83** | **3.36** | **1.36** | **0.69** | **0.345** | **0.138** | **0.068** | **0.035** | | Unif(G) | 12.63 | 6.35 | 2.59 | 1.27 | 0.641 | 0.256 | 0.126 | 0.064 | Observation: For all measurements, p(g|x) consistently offers smaller variances of estimation than Unif(G). **A3. In EXP-classify, p(g|x) offers a better sample efficiency than Unif(G).**

Further Response

Dear Reviewer Ry8Q, Thank you again for the insightful and constructive comments. Following common requests, we have conducted an in-depth empirical analysis of the convergence and sample complexity of the algorithm using the EXP-classify dataset (Sn invariance; Section 3.1). The results can be found in the latest common response: https://openreview.net/forum?id=phnN1eu5AX&noteId=ci75ubQbgK. We think the findings might be of your interest, in particular the analysis on the stability of gradients (**Q2**) and sample complexity in comparison to the uniform distribution (**Q1 and Q3**). Sincerely, Authors of submission 12717

Further Response

Dear Reviewer wgQS, Thank you again for the insightful and constructive comments. Following common requests, we have conducted an in-depth empirical analysis of the convergence and sample complexity of the algorithm using the EXP-classify dataset (Sn invariance; Section 3.1). The results can be found in the latest common response: https://openreview.net/forum?id=phnN1eu5AX&noteId=ci75ubQbgK. We think the findings might be of your interest, in particular the analysis on how the sample size for estimation affects training (**Q4**). A particularly interesting finding is that smaller sample size for training regularizes the model towards low-variance estimation. This supplements our previous response on sampling-based loss, and provides some understanding of the side effect induced by optimizing the sampling-based upper-bound to the true objective. Sincerely, Authors of submission 12717