Improving Equivariant Networks with Probabilistic Symmetry Breaking

We thank the reviewer for their time, and for their valuable feedback on our work. We are happy to see that they find the motivation and relevance of the paper clear. We also appreciate that they have found our theoretical results solid and insightful as well as our experiments detailed.

We will respond to each point individually below, and hope to resolve some of the reviewer’s concerns regarding presentation.

To begin, we hope to clarify the reviewer’s questions on probability and notation.

Question 1: Are distributions discrete or continuous?

In our paper, they can be either! For example, $X$ and $Y$ can be token sequences from a finite vocabulary (discrete), or graphs with real-valued node features (continuous). Happily, our theoretical results really do not depend on which of the two is the case. The appropriate notion of conditional distribution which encapsulates both cases is called regular conditional probability, and it gives the standard formulas for both conditional probability mass functions and conditional densities (or just “conditional probability distributions”).

Question 2: What are the input and output random variables, X and Y? What are their input and output spaces?

As mentioned at the beginning of the preliminaries section (lines 87 and 88), the input and output spaces are simply two sets on which we have probability distributions, and on which the action of $G$ satisfies some basic conditions (see Footnote 1) — so, technically, they should be measurable spaces, which we have added to the paper. In machine learning applications, one generally encodes inputs and outputs as vectors, which would make $\mathcal{X}$ and $\mathcal{Y}$ vector spaces. Again, a realistic example might be the vector space of point clouds in 3D.

Question 5: In Equation 2, line 137, Y is a random variable, but your function $f$ maps elements into $\mathcal{Y}$. [...] it’s unclear whether the probability is over $X$ or over $\epsilon$.

As a reminder, Equation 2 was: $Y =_{a.s.} f(X,\epsilon)$. Here, “a.s.” denotes that the equality holds almost surely, i.e. with probability 1. What does this mean? $X$, $Y$ and $\epsilon$ are all random variables, taking values in the sets $\mathcal{X}$, $\mathcal{Y}$ and $[0,1]$ respectively, as described in the paper. (See the preliminaries paragraph (line 95), where we specified this for $X$ and $Y$.) The fact that $\epsilon$ is a random variable is implied by the fact that we specify its sampling distribution, but we have now stated this fact explicitly to make sure it’s clear. The equation means that, when you sample each of $X$, $Y$, and $\epsilon$, the probability that the random variables satisfy $Y = f(X,\epsilon)$ is equal to 1. One could say the probability is thus “taken over” $X$, $Y$, and $\epsilon$ (but we emphasize that no integration/expectation is being taken, as the term “taken over” might suggest). We have added an explanation of the “a.s.” notation to the preliminaries.

Question 6: In Theorem 3.4, equation 3, it’s not clear what the probability is being taken over.

Equation (3) is $Y =_{a.s.} f(X,\tilde{g},\epsilon)$. The explanation is analogous to the one we gave above, except that $\tilde{g}$ is now an additional random variable sampled from the inversion kernel. We emphasize again that “almost surely” means with probability 1, with respect to the joint distribution of all the random variables.

Overall, we have aimed to use notations and concepts that are standard in probabilistic machine learning. These are presented for example in textbooks such as Durrett’s “Probability: Theory and Examples”, and used in prior papers in the area such as [4]. However, we have now additionally made sure to explicitly define the concepts above.

Dear reviewer,

Please engage in the discussion as soon as possible. Specifically, please acknowledge that you have thoroughly reviewed the authors' rebuttal and indicate whether your concerns have been adequately addressed. Your input during this critical phase is essential—not only for the authors but also for your fellow reviewers and the Area Chair—to ensure a fair evaluation. Best wishes, AC

We sincerely thank the reviewer for their thoughtful review and appreciation of our work. We are happy that they felt that our paper was easy to read, that it addresses an important and challenging problem, and that we propose a well-motivated method backed by experiments and theory.

We thank the reviewer for their positive feedback and valuable suggestions. We address specific questions below.

Increased computational cost due to canonicalization

This is a valid concern. However, to summarize, we have found that this is not an issue in practice.

First, we reiterate why the reviewer’s point is valid. Depending on the group and data type, properly canonicalizing may not even be solvable in polynomial time -- for example, in the case of $S_n$ acting on graphs, it would solve graph isomorphism (since, to determine if two graphs are isomorphic, one need merely canonicalize them both). However, Proposition 5.1 is our saving grace. It states that actually, we do not need a strict canonicalization function to represent any equivariant distribution; instead, any equivariant variable $Z$ will suffice. (The point of using a canonicalization, as discussed in Section 5 and Appendix E, is for learning efficiency -- we wish to inject as little noise as possible into the network.) Therefore, even if we are not exactly canonicalizing, such as using an invariant graph network to output permutations in the graph diffusion experiment (Section 8), we can still break symmetries perfectly.

Moreover, outside of cases like the aforementioned graphs, canonicalization via fixed analytical procedures is a widespread preprocessing step in many fields of science (as just one example, it is standard practice to canonicalize small molecules via SMILES strings using computational chemistry software like rdkit). Therefore, there are often heuristics readily available for canonicalization, or at least for outputting an equivariant $Z$ as described above. For example, at the end of Appendix C, we suggest a simple and efficient heuristic for point clouds based on randomly choosing two points of maximal radius. Our work allows one to harness such procedures for more general (learnable) symmetry-breaking.

Alternatively, if one uses an equivariant network to learn a canonicalization, as introduced by [3], the canonicalization function can be made much smaller than the neural network used for the prediction, resulting in a small additional overhead. Evidence for this is provided in the main text for graph autoencoder experiment and in the appendix for the other experiments: in the case of a learned canonicalization, using SymPE results in a negligible increase in the number of parameters compared to using only an equivariant network. The canonicalization networks are themselves cheap. We have provided explicit evidence for this in Tables 3 and 4 of the revised manuscript by including inference time for the spin experiment, and training time for the graph diffusion model. As shown, the additional computational overhead is relatively low (28h vs 34h for training the graph diffusion models, and 0.49s vs 0.52s for the forward time on the spin experiment).

Comparison with other positional encodings

That is a great comment. We have added some discussion of this work in the revised manuscript, along with other graph positional encodings such as RFP (Random Feature Propagation) [1] and RRWP (Relative Random Walk Position) [2], since our results have important implications in this regard. The core idea is that the learned positional encodings in Dwivedi et al. are initialized from Random Walk Positional Encodings. Although they can improve the expressivity of a GNN, they are still strictly equivariant. Therefore, as highlighted by our theoretical results, they also cannot break symmetries (i.e. nodes in the graph that are equivalent under automorphism will have the same embedding) and do not allow one to represent arbitrary equivariant distributions. An important contribution of our work is the proof that for a positional encoding to break graph symmetries, it must be sampled from a distribution that satisfies the conditions of Proposition 5.1.

Please let us know if we have satisfyingly addressed your concerns, and if you have any other questions. Note that the changes in the revised manuscript have been highlighted in blue for clarity.

[1] M. Eliasof, F. Frasca, B. Bevilacqua, E. Treister, G. Chechik, and H. Maron. Graph positional encoding via random feature propagation. In International Conference on Machine Learning, pages 9202–9223. PMLR, 2023. [2] L. Ma, C. Lin, D. Lim, A. Romero-Soriano, P. K. Dokania, M. Coates, P. Torr, and S.-N. Lim. Graph inductive biases in transformers without message passing. In International Conference on Machine Learning, pages 23321–23337. PMLR, 2023. [3] S.-O. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonicalization functions. ICML 2023.

We sincerely thank the reviewer for their thoughtful and detailed feedback. We are glad that they have found our paper clear, the theoretical results helpful, and the proposed method effective. We respond individually to their points below.

Positional encoding in Transformers versus our method

We agree that there are some subtleties to this suggestive choice of terminology, and thank the reviewer for insightfully highlighting them.

Breaking model vs input symmetries

First, the reviewer is correct that the primary goal of positional encodings (PEs) in transformers is not to break permutation symmetries of the input, but in the model architecture itself. This is part of why we call ours “symmetry-breaking” positional encodings, rather than ordinary positional encodings. Nonetheless, in both cases (as the reviewer points out), we are using some kind of positional encoding to turn an equivariant learning method into a less equivariant learning method. However, the analogy runs deeper than this -- we’ll dive into this in the subsequent paragraphs.

Relative positional encodings

Work on relative positional encodings [1] has argued that some degree of inductive bias towards equivariance, such as sequence translations, is desirable. Imagine then that we want our language model to be equivariant to translations (in an idealized case of zero-padded sequences). One could no longer use ordinary positional encodings, which fully break permutational symmetry (because they are constant with respect to input permutation, indeed they are independent of the input). Instead, one must use input-dependent positional encodings, which are equivariant with respect to input translation. (Then, even if we translate the input sequence, the positional encodings will translate too.) We would therefore apply SymPE with respect to the translation group, not the symmetric group. This would make the full model (SymPE + Transformer) equivariant to only translations, achieving the correct level of equivariance.

The two meanings of “position” then coincide

Seen in this way (above), the two meanings of position as "position of the tokens in the sentence" and "position of the data sample in the orbit" coincide. This is because specifying a single element of the translation group (via sampling from the inversion kernel), specifies the position of all the tokens --- since the relative ordering is preserved by translations. We think that the analogy is therefore still correct in this setting.

However, given the subtleties discussed above, the reviewer makes a fair point that this sequence model example could be more intuitive.

Graph positional encodings

The example of PEs in graphs is more intuitive. In this setting, the main use of PEs is to break symmetries’(automorphisms) and increase model expressivity. Like for sequences, we also have that the group element (permutation) obtained at a sample (graph) level completely specifies the position of individual components (nodes). , Relatedly, the difficulty of designing PEs for graphs has often [2,3] been attributed to the lack of a canonical ordering for this data structure. In SymPE, this problem is tackled directly by learning and sampling a canonical ordering from an inversion kernel. (Or, more realistically, a more general equivariant distribution as in Proposition 5.1.)

Based on the reviewer’s feedback, we have updated our motivation for the term “symmetry-breaking positional encodings”, at the bottom of page 5, in terms of graph positional encodings.

Encoding of the group element

It is true that things could be framed in the language of group representation for this part. However, since most of the results of the paper do not assume that we deal with representations and apply to more general actions, we have preferred to keep the formalism uniform.

The reviewer’s other suggestion to remove Prop 4.1 from the main text is a good idea, and we have implemented it (moving Prop 4.1 to the appendix).

Notation for encoding vector

This is a good suggestion. We have implemented it and changed the notation for the positional encoding from $e$ to $\tilde{v}$.

Eq (1): should it be this way or $P(Y|X=g^{−1}x)=P(gY|X=x)$?

By $P(gY|X=x)$, what we mean is $P(gY=y|X=x)$, i.e. $P(gY=y|X=x)=P(Y=y|X=gx)$ for all $y$ (interpreted as probability density functions).

Different results from Satorras et al. 2021

Thanks to the reviewer for pointing this out. Indeed, although we used the codebase of Satorras et al, the experimental setup was slightly different (we use the “erdosrenyinodes_0.25_none” dataset, the “AE” architecture, and train for fewer epochs) -- but the comparisons within our experiments should be fair. We will clarify these distinctions in the revision.