SymmetricDiffusers: Learning Discrete Diffusion on Finite Symmetric Groups

Thank you for the insightful and constructive comments. We appreciate your positive feedback and address the questions below.

Q1: You later say that you can merge steps of the forward process if each individual step does not induce enough mixing and so the stated weakness that these styles of shuffle have slow mixing seems moot.

A1: We’d like to clarify that we can only merge steps in the reverse process and not in the forward process. While merging steps in the reverse process makes loss computation more efficient, we still need to run the whole forward process for random transpositions and random insertions. And since the mixing time for random transpositions and random insertions are much slower than riffle shuffles, riffle shuffles should be the preferred shuffling method.

Q2: I think it should also be discussed that $q(X_{t-1}|X_t,X_0)$ is also unavailable since this distribution is used in standard diffusion models to re-write the variational bound in a lower variance form...

A2: Yes, the posterior $q(X_{t-1}|X_t,X_0)$ is unavailable for most shuffling methods, so we cannot re-write the ELBO in a lower variance form like in common diffusion models. It is worth noting that even though $q(X_{t-1}|X_t,X_0)$ is available for riffle shuffles, the common parameterization presented in diffusion models is still not helpful. Please refer to the General Replies for a detailed discussion.

Q3: However, the ablations with these different styles of shuffle are in the appendix and I think it should be in the main …

A3: We appreciate this suggestion and will reorganize the paper's structure to include the ablation.

Q4: I find it difficult to follow the description of the inverse transposition parametrization, there is no intuition given for the functions $\phi$ and $\psi$ nor the functional form of $p_{\text{IT}}(\sigma)$...

A4: We apologize for the confusion regarding the inverse transposition parameterization. Due to the space limit, we indeed shortened the description. We will reorganize the sections and include the following intuitions.

The support of the IT distribution is the set of all transpositions plus the identity permutation. Let us treat the identity permutation $\mathrm{Id}$ separately, and we use a parameter $\tau$ to assign $p_{\mathrm{IT}}(\mathrm{Id}) = 1 - \mathrm{sigmoid}(\tau)$. Then we assign probabilities to the transpositions. A transposition is essentially an unordered pair of distinct indices, so we use $n$ parameters $\mathbf{s}=(s_1,\ldots,s_n)$ to represent the logits of each index getting picked. Thus, for indices $i\neq j$, it is natural to let

for a transposition $\sigma=\begin{pmatrix}i&j\end{pmatrix}\neq\mathrm{Id}$, where $\mathrm{sigmoid}(\tau)$ is the probability of not choosing $\mathrm{Id}$. The expression in the parentheses is the probability of choosing the unordered pair $i$ and $j$, which is equal to the probability of choosing $i$ and then $j$, plus the probability of choosing $j$ and then $i$.

Q5: Isn't there an alternative where although you need to run the full forward process since $q\left(X_t | X_0\right)$ is intractable, you could still only compute the loss on some subset of the reverse transitions…

A5: Our framework allows computing the loss on some subset of the trajectory to be more memory efficient. For example, we could randomly sample one timestep from a denoising schedule $[t_0,\ldots,t_k]$ and compute the loss as

omitting constant terms with respect to $\theta$. In fact, for riffle shuffles, we could sample $X_{t_{i-1}}$ directly for arbitrary timestep $t_{i-1}$ from $X_0$ [1]. For other shuffling methods, we would have to run the entire forward process. It is also worth noting that computing the loss on a subset of the trajectory could potentially introduce more variance during training. So, there is a tradeoff, and the exact design choice should depend on the problem setup.

In our current implementations of the experiments, we make $T$ (parallelized) calls to the neural network to compute the expectation w.r.t. time in the ELBO exactly. This is because, as discussed in our paper, the forward trajectories are usually short for riffle shuffles, so the exact computing of the variational bound is feasible and can help reduce variance.

[1] Dave Bayer and Persi Diaconis. Trailing the dovetail shuffle to its lair. The Annals of Applied Probability, 2(2):294 – 313, 1992.

Thank you for the insightful and constructive comments. We appreciate your positive feedback and address the questions below.

Q1: The model proposes to directly learn the reverse transition densities $p_{\theta}(X_{t-1}|X_t)$. The issue with doing this for standard diffusion models is that this seems to hurt model training since it increases the variance of training (as, in particular, one must sample the two $X_{t-1}$, $X_t$ for training instead of just one $X_t$). As such, most works use the (ultimately equivalent) mean/score-parameterizations [1, 2]. I would want to hear a bit more about if this would be applicable in the $S_n$ case (and training with this parameterization might improve the model) or if this is not possible.

A1: We agree that learning $p_{\theta}(X_{t-1}|X_t)$ directly increases the variance of training. However, the reason why previous works can use mean/score-parameterizations is that they use Gaussian kernels on continuous data, which admits analytical forms for $q(X_{t-1}|X_t,X_0)$ and $q(X_t|X_0)$.

However, such parameterizations are not applicable to $S_n$ since the transitions are not Gaussian. In particular, the posterior $q(X_{t-1}|X_t,X_0)$ is unavailable for most shuffling methods. We cannot rewrite the ELBO in a lower variance form like in standard diffusion models. It is worth noting that even though $q(X_{t-1}|X_t,X_0)$ is available for riffle shuffles, the common parameterization presented in diffusion models is still not helpful.

Please refer to the General Replies for a detailed discussion.

Q2: In particular, works like [3, 4, 5] have established a working theory for discrete diffusion in continuous time, so it would be beneficial to discuss how the proposed framework might fit into the established theory.

A2: Thanks for bringing up this point and mentioning the interesting work. Since the shuffling methods are inherently discrete-time Markov chains, it seems non-trivial to convert them into continuous-time Markov processes. In other words, the commonly used linear ODE parameterization in [3, 4] does not match the shuffling dynamics considered in our work. On the other hand, the concrete score and the score entropy idea in [5] seem to be compatible, i.e., we could possibly change the parameterization to develop the discrete-time counterpart in our context. We will explore this topic more and add the discussion in the later version.

Q3: The experiments, while showing good results, do not show that the method is particularly scalable… While some discussion is made here that talks about transformer layers, I think large values of $n$ aren't that big of an issue in transformers due to systems like Flash Attention. So, it should be made more clear if this is a fundamental problem with the existing method, or a larger scale example (even toy) should be presented.

A3: We agree that techniques like Flash Attention could improve our model’s scalability. A large value of $n$ is indeed not a big issue since our backbone network is also a Transformer.

During submission, since some metrics (e.g., accuracy) drop significantly as $n$ increases and other competitive methods do not scale up well, we only report the results with $n\le100$ for a better comparison.

During the rebuttal, we further investigated our model for the task of sorting 4-digit MNIST numbers with $n=200$ and show the results below.

This result is already better than the performance of sorting network based models on the easier task with $n=100$ (Diffsort/Error-free Diffsort: Kendall-Tau=0.166/0.140, Correct (%)=23.2/20.1). We will include more large-scale experiments in the later version.

Q1: In Robin Winter et al.'s paper, the authors have studied the symmetric group and my understanding is disentangling invariant features and equivariant groups enables the design of flexible diffusion models since you only need to perform diffusion modeling in the invariant latent space.

A1: Thanks for introducing Winter’s interesting work [1]. Specifically, [1] employs a graph-invariant encoder, $\eta$ , ensuring that regardless of the group action $\rho_X$ applied from group $G$ on a discrete state $x \in X$ , the encoder's output remains constant ($\eta(\rho_X(g) x) = \eta(x), \forall x \in X, \forall g \in G$ ). This invariant characteristic is particularly beneficial for molecule generation [2], where the goal is to learn a probability distribution over molecules that is $\mathrm{SE}(3)$-invariant.

However, we’d like to clarify that performing diffusion modeling in the invariant latent space can not solve the problem in our paper. In particular, our goal is to learn a distribution over a finite symmetric group $S_n$, where it is critical that representations of different elements within the group are distinct. For instance, since distinct permutations result in different images in Jigsaw Puzzles, an invariant encoder can not differentiate between incorrect and correct permutations based on the latent representations. Therefore, a group-invariant encoder does not solve our problem.

On the other hand, one could use the group-equivariant encoder $\psi$ proposed in [1], which predicts group actions in one shot instead of performing diffusion. However, the resulting model would be essentially similar to the one in Gumbel-Sinkhorn [3], which has been shown to be less effective than ours in several tasks.

Therefore, the approach in [1] does not align with the objectives of our study. We will discuss this work in a future version.

[1] Winter, Robin, et al. "Unsupervised learning of group invariant and equivariant representations." Advances in Neural Information Processing Systems 35 (2022): 31942-31956.

[2] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." International Conference on Machine Learning. PMLR, 2023.

[3] Mena, Gonzalo, et al. "Learning latent permutations with gumbel-sinkhorn networks." arXiv preprint arXiv:1802.08665 (2018).

Q2: My main concern is the proposed method is not the state-of-the-art… I want to see the comparison of the approach proposed in Robin Winter et al's paper. My understanding is that Robin Winter's approach enables building state-of-the-art diffusion models for molecular generation.

A2: To the best of our knowledge, our method achieves state-of-the-art performances in the tasks of sorting 4-digit MNIST images and jigsaw puzzles. It is important to note that our focus is not on developing diffusion models for molecular generation.

I noticed this thread was developing, and I feel obligated to say something.

Reviewer UKH2's comments seem to highly mischaracterize the contributions of Winter et al. to the point of being largely unfair and invalid. In particular,

1. [1] does not build a diffusion model anywhere in the paper, counter to claims made in the original comment. In fact, [1] is not even a probabilistic model at all. The work is effectively a equivariant AE (that operates by instead mapping to $X/G \times G$ instead of $X$ in the mathematical framework). Note that, as a strict autoencoder, it is not even probabilistically meaningful, which means that we aren't learning a distribution on $X$ anyways.

2. [1] does not seem to be SOTA. In particular, the work is qualitative and shows the performance of the embeddings in, for example, visualizing data. The work is not quantitative and does not present any major numerical results.

3. [1] is primarily focused on building group equivariant models. In particular, I emphasize that, instead of modeling a distribution over the group $G$, one would model a distributions over a group $X$ that $G$ acts on (ie $X = \mathbb{R}^{n \times 3}$ and $G = SO(3) \times S_n$ would be a typical setup for molecules) if [1] was a proper VAE instead of an AE. This is extremely different from the proposed work, which models a distribution on $G$.

4. [2] is completely different from [1], is more relevant to the proposed work, and, even then, would not be applicable here. [2] learns a proper probabilistic model (ie a VAE + diffusion), which is completely different from the goal of [1]. [2] would not be applicable here either since we would be looking to model a distribution over $X$ that is equivariant with respect to $G$. If $X = S_n$, then there is no such $G$ besides the trivial group which would give the required flexibility. Such a framework would then have to continuously embed $|S_n| = n!$ number of elements into the latent space and decode there, which is computationally infeasible since even the most performant VQVAEs (ie discrete quantized representation VAEs, trained through some variant of straight through estimator) can have a codebook size of 1024.

[1] Winter, Robin, et al. "Unsupervised learning of group invariant and equivariant representations." Advances in Neural Information Processing Systems 35 (2022): 31942-31956.

[2] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." International Conference on Machine Learning. PMLR, 2023.

This is, unfortunately, the second time in this thread that I feel the need to step in.

To reviewer UKH2:

• Please clarify your prior position and argument. Both the authors and I have stated explicitly that the prior work that you reference does not work in the problem setting of the submitted paper. In particular, the prior work you reference can enable one to learn a $\mathbf{S_n}$-equivariant probabilistic model while the submitted work learns a probabilistic model on $\mathbf{S_n}$, which are two very different things and do not overlap in any meaningful way. For example, this means the submitted paper is not actually building a "symmetric diffusion". For clarity sake, I am asking you to directly address which part of this assertion you think is incorrect (e.g. if the statement does not accurately portray your references, the submitted work, or the differences/similarities between the two bolded ideas).

• Please make more concrete points. In particular, please do not reference course projects or intuitions about methods or anything else like this that we (authors, reviewers, ACs) are unable to directly access. For example, if UKH2 has played with the provided code in this paper and has tried a baseline method that they have mentioned (which could explain the motivation behind your comments), it would better support your point if you can provide us all with the code through an anonymous github link or something similar.

• Please review in good faith. Please do not accuse the authors (or me) of not reading your references. Please do not lower your score to spite the authors when they disagree with you. Please do not insult the submitted work in a dismissive way (e.g. "weak" and "not worth publishing" vs saying something like "does not adequately compare against relevant baselines", which also happens to be a more precise characterization of your critique).

I think my point on how to solve the problem by parameterizing the encoder and denoising process is clear, and the implementation is also clear - replacing E(n)-equivariant neural networks in GeoLDM with S(n)-equivariant neural networks used by Winter et al. The implementation is not challenging at all, and it is not my responsibility to provide code for this. As a reviewer, my role is to properly assess the paper based on my knowledge, though I cannot guarantee completely adequate.

I lowered the score because I think the paper has major flaws after discussing it with the authors.

Let us recall how you violate good faith: recall that I first raise my comments on another solution to solve the problem, and look at how insulting and lack of scientific merits in your first comment. Then I made my point clear on how to solve the problem in detail, then looked at your insulting language in your latest comment.

I think your position is dubious and lacks good faith. My points are more than clear, before the authors showed me reasonable evidence that the proposed approach is better than using the latent diffusion, I insist on rejecting this paper.

I am not going to respond to any of your further unscientific comments.