SymmetricDiffusers: Learning Discrete Diffusion on Finite Symmetric Groups
Authors: Yongxing Zhang
Donglin Yang
Renjie Liao
Venue: ICLR 2025 Oral (2025)
Keywords: Finite Symmetric Groups
Discrete Diffusion
Permutations
Riffle Shuffles
Plackett-Luce Distribution
Sorting
Jigsaw Puzzle
Primary area: generative models

Abstract: The group of permutations $S_n$, also known as the finite symmetric groups, are essential in fields such as combinatorics, physics, and chemistry. However, learning a probability distribution over $S_n$ poses significant challenges due to its intractable size and discrete nature. In this paper, we introduce *SymmetricDiffusers*, a novel discrete diffusion model that simplifies the task of learning a complicated distribution over $S_n$ by decomposing it into learning simpler transitions of the reverse diffusion using deep neural networks. We identify the riffle shuffle as an effective forward transition and provide empirical guidelines for selecting the diffusion length based on the theory of random walks on finite groups. Additionally, we propose a generalized Plackett-Luce (PL) distribution for the reverse transition, which is provably more expressive than the PL distribution. We further introduce a theoretically grounded "denoising schedule" to improve sampling and learning efficiency. Extensive experiments show that our model achieves state-of-the-art or comparable performance on solving tasks including sorting 4-digit MNIST images, jigsaw puzzles, and traveling salesman problems. Our code is released at <https://github.com/DSL-Lab/SymmetricDiffusers>.

Published as a conference paper at ICLR 2025
SYMMETRICDIFFUSERS LEARNING DISCRETE DIFFUSION ON FINITE SYMMETRIC GROUPS
Yongxing Zhang .3+, Donglin Yang Renjie Liao
1 University of Waterloo University of British Columbia Vector Institute nick. zhang@uwate ca, (ydlin, rjliao} &ece. ubc. ca
ABSTRACT
The group of permutations Sn, also known as the finite symmetric groups, are essential in fields such as combinatorics, physics, and chemistry. However, learning a probability distribution over Sn poses significant challenges due to its intractable size and discrete nature. In this paper, we introduce SymmetrricDiflusers,, a novel discrete diffusion model that simplifies the task of learning a complicated distri- bution over S, by decomposing it into learning simpler transitions of the reverse diffusion using deep neural networks. We identify the riffle shuffle as an effective forward transition and provide empirical guidelines for selecting the diffusion length based on the theory of random walks on finite groups. Additionally, we propose a generalized Plackett-Luce (PL. distribution for the reverse transition, which is provably more expressive than the PI. distribution. We further introduce theoretically grounded "denoising schedule" to improve sampling and learning efficiency. Extensive experiments show that our model achieves state-of-the-an comparable performance on solving tasks including sorting 4-digit MNIST images, jigsaw puzzles, and traveling salesman problems. Our code is released at attps://githut com DS///SSLLLLL Laapm/DDDLLLLaLSLL-- http:a/aaaaymmeera....feeeeeee
INTRODUCTION
As vital area of abstract algebra, finite groups provide a structured framework for analyzing symme- tries and transformations which are fundamental to a wide range of fields, including combinatories, physics, chemistry, and computer science. One of the most important finite groups is the finite symmetric group Sn. defined as the group whose elements are all the bijections (or permutations) from set of elements to itself, with the group operation being function composition.
Classic probabilistic models for finite symmetric groups Sn, such as the Plackett-Luce (PL) model Plackett 1975 Luce, 1959), the Mallows model (Mallows) 1957). and card shuffling methods (Diaconis) 1988 are crucial in analyzing preference data and understanding the convergence of random walks. Therefore, studying probabilistic models over S, through the lens of modern machine learning is both natural and beneficial. This problem is theoretically intriguing as bridges abstract algebra and machine learning. For instance, Cayley's Theorem, fundamental result in abstract algebra, states that every group is isomorphic to a subgroup of a symmetric group. This implies that learning probability distribution over finite symmetric groups could, in principle, yield distribution over any finite group. Moreover, exploring this problem could lead to the development of advanced models capable of addressing tasks such permutations in ranking problems, sequence alignment in bioinformatics, and sorting.
However, learning a probability distribution over finite symmetric groups S, poses significant challenges. First, the number of permutations of n objects grows factorially with n, making the inference and learning computationally expensive for large n. Second, the discrete nature of the data brings difficulties in designing expressive parameterizations and impedes the gradient-based learning.
In this work, we propose a novel discrete-time discrete (state space) diffusion model over finite symmetric groups, dubbed as Dyffetrrcciffusser. It overcomes the above challenges by decomposing the difficult problem of learning a complicated distribution over S, into a sequence of simpler
Work done while intern Vector Institute.

---

Published as a conference paper at ICLR 2025
problems, L.e., learning individual transitions of a reverse diffusion process using deep neural networks. Based on the theory of random walks on finite groups, we investigate various shuffling methods as the forward process and identify the riffle shuffle the most effective. We also provide empirical guidelines on choosing the diffusion length based on the mixing time of the riffle shuffle. Furthermore, examine potential transitions for the reverse diffusion, such as inverse shuffling methods and the PL distribution, and introduce novel generalized PL distribution. We prove that our generalized PL is more expressive than the PL distribution. Additionally, we propose theoretically grounded "denoising schedule" that merges reverse steps to improve the efficiency of sampling and learning. To validate the effectiveness of our SymmetricDiffusers we conduct extensive experiments on three tasks: sorting 4-Digit MNIST images, solving Jigsaw Puzzles on the Noisy MNIST and CIFAR- 10 datasets, and addressing traveling salesman problems (TSPs). Our model achieves the state-of-the-an or comparable performance across all tasks.
2 RELATED WORKS
Random Walks on Finite Groups. The field of random walks on finite groups, especially finite symmetric groups, have been extensively studied by previous mathematicians (Reeds) 1981 Gilbert 1955 Bayer & Diaconis 1992, Saloff-Coste)] 2004). Techniques from a variety of different fields, including probability, combinatorics, and representation theory, have been used to study random walks on finite groups (Sateff-Coste] 2004). In particular, random walks on finite symmetric groups are first studied the application of card shuffling, with many profound theoretical results of shuffling established. A famous result in the field shows that riffle shuffles are enough to mix up deck of 52 cards (Bayer & Diaconis 1992), where riffle shuffle is a mathematically precise model that simulates how people shuffle cards in real life. The idea of shuffling to mix up deck of cards aligns naturally with the idea of diffusion, and we seek to fuse the modern techniques of diffusion models with the classical theories of random walks on finite groups.
Diffusion Models, Diffusion models (Sohl- Dickkteiiktnn et al. 2015 Song & Ermon 2020 Ho et al. 2020 Song et al. )021) are powerful class of generative models that typically deals with continuous data. They consist of forward and reverse processes. The forward process is typically a discrete-time continuous- state Markov chain or continuous-tin time continuous-star Markov process that gradually adds noise to data, and the reverse process learn neural networks to denoise. Discrete (state space) diffusion models have also been proposed to handle discrete data like image, text (Austin et al. 2023). and graphs (Vignac et (003)) However, existing discrete diffusion models focused on cases where the state space is small has a special (e.g., dccomposable) structure and are unable to deal with intractable-sizer state spaces like the symmetric group. In particular, Austin et al. (2023) requires an explicit transition matrix, which has size n! x n! in the case of finite symmetric groups and has no simple representations or sparsifications. Finally, other recent advancement includes efficient discrete transitions for sequences (Varma et al. 2024), continuous-
discrete- state diffusion models (Campbell et al. 2022: Sun et al.) [023] Shi et al.) (024) and discrete score matching models (Meng) jet al. 2023, Lou al. 2024), but the nature of symmetric groups again makes non-trivial to adapt to these cxisting frameworks.
Differentiable Sorting and Learning Permutations. A popular paradigm to learn permutations is through differentiable sorting or matching algorithms. Various differentiable sorting algorithms have been proposed that uses continuous relaxations of permutation matrices (Grover et al.
2018 Cuturi
et al. 2019 Blondel et al. 2020), or uses differentiable swap functions (Petersen al. Kim et al. 2024). The Gumbel-Sinkhon latent permutations using the continuous Sinkhorn operator. Such methods often focus on finding the optimal permutation instead of learning a distribution over the finite symmetric group. Moreover, they tend be less effective as n grows larger due their high complexities.
202 2022
method (Mena et al., (018)) has also been proposed to learn
LEARNING DIFFUSION MODELS ON FINITE S SYMMETRIC GROUPS
We first introduce some notations. Fix N. Let [n] denote the set 21, ... n... A permutation
2
o(n) The identity permutation, denoted by Id, is the permutation given by Id(i) = for all i E [n]. Let S, be the set of all permutations (or bijections) from set of elements to itself, called the finite
on [n] is function from [n] to [n], and we usually write a as ((1) (2)
2

---

Published as a conference paper at ICLR 2025
9(X1
231
132
231
321
123
821
198
312
213
312
213
TT
0:
σ1
x(Xr|XX--1:0 or)
X(Xi X(11
Xxr
X,
Xi
X1
p(X
1 of)
p(Xo)XX1.0f)
Xr)
Po(o!1x,)
pa(oi](ii)
d2
7 #
Probability
Ferwand Step
X
3
9 3
4 7
Asconating
X3 1 S
Xa)
7 7
2 9
X2
[x), Xo, Xs}
Merged Revene Step
Figure I: This figure illustrates our discrete diffusion model on finite symmetric groups. The middle graphical model displays the forward and reverse diffusion processes. We demonstrate learning distributions over the symmetric group S3 via the task of sorting three MNIST 4-digit images. The top part of the figure shows the marginal distribution of ranked list of images X₃ at time t, while the bottom shows randomly drawn list of images.
symmetric group, whose group operation is the function composition. For a permutation a E Sn. the permutation matrix Q, Rax= associated with satisfies e, Q, (()) for all [n]. In this paper, we consider a set n distinctive objects X = {xI..... Xn where the i-th object is represented by a d-dimensional vector x; Therefore, a ranked list of objects can be represented as a matrix X = [xx..... Xn, E Raxd where the ordering of rows corresponds to the ordering of objects. We can permute X via permutation to obtain Q.
Our goal is to learn a distribution over Sn. We propose learning discrete (state space) diffusion models, which consist of a forward process and reverse process. In the forward process, starting from the unknown data distribution, we simulate a random walk until it reaches a known stationary "noise" distribution. In the reverse process, starting from the known noise distribution, we simulate another random walk, where the transition probability is computed using a neural network, until it recovers the data distribution. Learning transition distribution over S, is often more manageable than learning the original distribution because: (1) the support size (the number of states that can be reached in one transition) could be much smaller than n!. and (2) the distance between the initial and target distributions is smaller. By doing
we break down the hard problem (learning the original distribution) into a sequence of simpler subproblems (learning the transition distribution). The overall framework is illustrated in Fig. [ In the following, we will introduce the forward card shuffling process in Section 3.1 the reverse process in Section 3.2 the network architecture and training in Section 3.3 denoising schedule in Section 3.4 and reverse decoding methods in Section 3.5
3.1 FORWARD DIFFUSION PROCESS: CARD SHUFFLING
Suppose we observe a set of objects X' and their ranked list Xo- They are assumed to be generated from an unknown data distribution in an IID manner, i.e., Xo, iid Paata X, X). One can construct a bijection between ranked list of n objects and an ordered deck of cards. Therefore, permuting objects is equivalent to shuffling cards. In the forward diffusion process, we would like to add "random noise" to the rank list so that it reaches to some known stationary distribution like the uniform. Formally, we let S C S, be a set of permutations that are realizable by a given shuffling method in step. S does not change across steps in common shuffling methods. We will provide concrete examples later. We then define the forward process as a Markov chain,
II T
q(Xi.T) X) ((Xir)(Xo)
(1)
where
q( (X o2 (g())) and the first equality in Eq. ① holds since Xo implics X. In the forward process, although the set X' does not change, the rank list of objects X, changes. Here g(o.) has the support S and describes the permutation generated by the underlying
3

---

Published as a conference paper at ICLR 2025
shufiling method. Note that common shuffling methods are time- honnoeooogeeeuus Markov chains, i.e., (())) stays the same across time. qXx, oe) is a delta distribution ó (X, = Px-1) since the permuted objects X2 uniquely determined given the permutation o, and X,-1- 1- We denote the neighhouring states of X via one-step shuffling as (x(X)= = (Qa X2xX|o l E S}. Therefore, we have, if X E (x(
{ (g(())
q( xX
10
otherwise.
Note that X₂ E (N((X, is equivalent to E S and X, Qo, X1-1-
3.1.1 CARD SHUFFLING METHODS
We now consider several popular shuffling methods the forward transition, i.e., random transpo- sitions, random insertions, and riffle shuffles. Different shuffling methods provide different design choices of q( thus corresponding to different forward diffusion processes. Although all these forward diffusion processes share the same stationary distribution, i.e., the uniform, they differ in their mixing time. We will introduce stronger quantitative results on their mixing time later.
Random Transpositions. One natural way of shuffling is to swap pairs of objects. Formally, transposition or a swap is a permutation a E S, such that there exist + j E [n] with (ii) = j. a(i) = i, and (k) k for all k {i.j}, in which case we denote = (i j). We let S = {(i j): ;) E [n]] {Id]. For any time f, we define q( by choosing two indices from [n] uniformly and independently and swap the two indices. If the two chosen indices are the same, then this means that we have sampled the identity permutation. Specifically, g( = (i j)) = 2/1² when * j and g(0, Td) = =1/n.
Random Insertions. Another shuffling method is to insert the last piece to somewhere in the middle. Let insert, denote the permutation that inserts the last piece right before the 5th piece, and let S = {insert; : E [n]]. Note that insert, = Id. Specifically, we have (a; = insert.) = 1/n when * and q(o, Id)= = 1/n..
Riffle Shuffles. Finally, we introduce the riffle shuffle, a method similar to how serious card players shuffle cards. The process begins by roughly cutting the deck into two halves and then interleaving the two halves together. A formal mathematical model of the riffle shuffle, known as the GSR model, was introduced by Gilbert and Shannon (Gilbert 1955), and independently by Reeds (198T). The model is described as follows. A deck of n cards is cut into two piles according to binomial distribution, where the probability of having cards in the top pile is (2)/2" for < k S n. The top pile is held in the left hand and the bottom pile in the right hand. The two piles are then riffled together such that, if there are A cards left in the left hand and B cards in the right hand, the probability that the next card drops from the left is A/(A + B), and from right is B (A + B). We implement the riffle shuffles according to the GSR model. For simplicity, we will omit the term "GSR" when referring to riffle shuffles hereafter.
There exists an exact formula for the probability over Sn obtained through one-step riffle shuffle. Let Sn. A rising sequence of a is subsequence of constructed by finding a maximal subset of indices i, < i2 < < ij such that permuted values are contiguously increasing, i.e.,
(i2)) - (11)
(is) (12)
o(i,) a(iy-1)
For example, the permutation
2 3
1 4 5 3 has 2 rising sequences, iee,, 123 (red) and 45 (blue). Note that permutation
has rising sequence if and only if it is the identity permutation. Denoting by (Rs(() (a) the probability of obtaining a through one-step riffle shuffle, was shown by BBayer & Diaconis (1922) that if Id if a has two rising sequences otherwise,
n + (m+t2-)
(n + 1)/2"
(RS(O)
1/2"
(3)
2n
where r is the number rising sequences a. The support S thus the set all permutations with at most two rising sequences. We let the forward process be q(σ,) (Rs( for all t.
3.1.2 MIXING TIMES AND CUT-OFF PHENOMENON
All of the above shuffling methods have the uniform distribution as the stationary distribution. However, they have different mixing times (iee. the time until the Markov chain close to its
4

---

Published as a conference paper at ICLR 2025
stationary distribution measured by some distance), and there exist quantitative results their mixing times. Let (yRT, 9RI, gns), and for & E N, let q()) be the marginal distribution of the Markov chain after shuffles. We describe the mixing time in terms of the total variation (TV) distance between two probability distributions, i.e., Drv (g u), where u is the uniform distribution.
For all three shuffling methods, there exists a cut-off phenomenon, where Drv(g u) stays around 1 for initial steps and then abruptly drops to values that are close to O. The cut-off time is the time when the abrupt change happens. For the formal definition, we refer the readers to Definition 3.3 Saloff Coste (004). In Moloff-Costte/ee (2004) they also provided the cut-off time for random transposition. random insertion, and rifile shuffle, which are log n. n log n. and 3 log, n respectively. Observe that the riffle shuffle reaches the cut-off much faster than the other two methods, which means it has much faster mixing time. Therefore, we use the riffle shuffle in the forward process.
3.2 THE REVERSE DIFFUSION PROCESS
We now model the reverse process another Markov chain conditioned on the set of objects X' We denote the set of realizable reverse permutations as T. and the neighbours of X with respect to T as (x) = (QaX E 7). The conditional joint distribution given by
II pe (
po -(
where po (x(x-1(Xx) ET p(,, oo)oo To sample from p(Xr((X), one simply samples random permutation from the uniform distribution and then shuffle the objects accordingly to obtain Xr. p(Xt-1lX, of) is again delta distribution ó(X,-1 Qo) We have
{0o X., if E otherwise,
pa Pa(X.-1
(5)
where X,-1 E (X,) is equivalent of E T and X,-1 Qo; In the following, will introduce the specific design choices of the distribution po(o)(X,).
3.2.1 INVERSE CARD SHUFFLING
A natural choice is to the inverse operations of the aforementioned card shuffling operations in the forward process. Specifically, for the forward shuffling S, we introduce their inverse operations T {o- :o E S). from which we can parameterize X.)
Inverse Transposition. Since the inverse of a transposition is also transposition, we can let T S S ((i j): i E [n]} U [Id]. We define distribution of inverse transposition (IT) over T using n + real-valued parameters s = (81,..., .8) and T such that
1 1- x(T)
if Id,
pır(o)
pb())
exp(s:) exp((*)
(x,) exp((tt exp((s)
exp(s,)
exp (si) exp(sx)
if (i (b)).
(6) where (b.) is the sigmoid function. The intuition behind this parameterization is to first handle the identity permutation Id separately, where we use Φ(T) to denote the probability of not selecting ld. Afterwards, probabilities are assigned to the transpositions. A transposition is essentially an unordered pair of distinct indices, so we use n parameters s = (s1,..., 8n) to represent the logits of each index getting picked. The term in parentheses represents the probability of selecting the unordered pair i and j, which is equal to the probability of first picking i and then j. plus the probability of first picking j and then i.
Inverse Insertion. For the random insertion, the inverse operation is to insert some piece to the end. Let inverse inverse_insee insert, denote the permutation that moves the ,th component to the end, and let T = (inverse. i insert; i [n]]]. We define categorial distribution of inverse insertion (II) over T using parameters s = (s,,..., 8x) such that,
pn(a diveers_insertee = exp(st)/ (E/- exp(s;)).
5

---

Published as a conference paper at ICLR 2025
Inverse Riffle Shuffle. In the riffle shuffle, the deck of card is first cut into two piles, and the piles are riffled together. So to undo a riffle shuffle, we need to figure out which pile each card belongs to, i.., making a sequence of binary decisions. We define the Inverse Riffle Shuffle (IRS) distribution using parameters $ = (81,.... 8n) as follows. Starting from the last (the n² object, each object i has probability ((s;) of heing put the top of the left pile. Otherwise, it falls the top of the right pile. Finally, put the left pile on top of the right pile, which gives the shuffled result.
3.2.2 THE PLACKETT- DISTRIBUTION AND ITS GENERALIZATION
Other than specific inverse shuffling methods to parameterize the reverse process, we also consider general distributions pa (o(ol whose support are the whole S,,, i.e., T = Sn.
The PL. Distribution. popular distribution over S, is the Plackett-Luce (PL.) distribution Plackettt
1975
Luce 1959). which is constructed from scores ($1...., 8n) as follows,
PPL(O) = II[ exp (sa(a)/ (E}= exp (8o(5)))
(8)
for all Sn- Intuitively, ($1...., Sn)) represents the preference given to each index in [n]] To sample from pl,, we first sample (1) from Cat(n, max Then we remove a(1) from the list and sample (2) from the categorical distribution corresponding to the rest of the scores (logits). We continue in this manner until we have sampled a (1),.... o (n). By Cao et al. (2007) the mode of the PL distribution is the permutation that sorts s in descending order. However, the PL distribution is
not very expressive. In particular, we have the following result, and the proof is given in Appendix E
Proposition The PL. distribution cannot represent a delta distribution over Sn-
The Generalized PL (GPL) Distribution. We then propose a generalization of the PI. distribution, referred to as Generalized Plackett- Luce (GPL) Distribution. Unlike the PL distribution, which uses set of n scores, the GPL distribution uses n² scores {s1, s.,). where each s; (5,,1.... 8i,)) consists of n scores. The GPL. distribution is constructed as follows,
PG) == II exp (Si,(1)) / (E exp (St)6)))
(9)
Sampling of the GPL distribution begins with sampling (1) using n scores $1. For S n, we remove i - scores from s; that correspond to (I)), a 1) and sample (i) from a categorical distribution constructed from the remaining n i scores in sj. It is important to note that the family of PL distributions is strict subset of the GPL. family. Since the GPL distribution has more parameters than the PL distribution, is expected to be more expressive. In fact, we prove the following significant result, and the proof is given in Appendix E
Theorem 2. The reverse diffusion process parameterized using the GPL distribution in Eq. 9 can model any distribution over S..
3.3 NETWORK ARCHITECTURE AND TRAINING
We now briefly introduce how to use neural networks to parameterize the above distributions used in the reverse process. At any time t, given X, Igd we use a neural network with parameters 0 to construct po Xx)) In particular, we treat n rows of X, as tokens and use Transformer architecture along with the time embedding of and the positional encoding to predict the previously mentioned scores. For example, for the GPL distribution, to predict scores, we introduce n dummy tokens that correspond to the n permuted output positions. We then perform few layers masked self-attention (2n. 2n) to obtain the token embedding Zj E X corresponding to input tokens and Ze E 1n Xhundel corresponding to n dummy tokens. Finally, the GPL score matrix is obtained Sn Z,Z, R=Xx Since the aforementioned distributions have different numbers of scores, the specific architectures of the Transformer differ. We provide more details in Appendix B
To learn the diffusion model, we maximize the following variational lower bound:
1X Xx)
(Xo. log Pappee(ea (Xol 2 X0a (Xo.)))))
logp(Xrlx)) log x(x.
(10)
In practice, one can draw samples to obtain the Monte Carlo estimation of the lower bound. Due to the complexity of shuffling transition in the forward process, we not obtain g ((x,)x)) analytically,

---

Published as a conference paper at ICLR 2025
as is done in common diffusion models. Therefore, we have to the forward process collect samples. Fortunately, it is efficient as the forward process only involves shuffling integers. We include more training details in Appendix G
Note that most existing diffusion models, such as those proposed by Ho et al. (2020) and Austin et al. (023), use equivalent form of the above variational bound, which involves the analytical KL divergence between the posterior q(X,-1 xX,,,,X, Xo) and po ( (x,) for variance control. How- ever, this variational bound cannot be applied S, because the transitions are not Gaussian, and ((X-iiX,, Xo) is generally unavailable for most shuffling methods. Most existing diffusion models also sample a random timestep of the loss. While this technique is also available in our framework, introduces more variance and does not improve efficiency all the time. And for riffle shuffles, the trajectory is usually short enough that we can compute the loss on the whole trajectory. A detailed discussion can be found in Appendix C
3.4 DENOISING SCHEDULE VIA MERGING REVERSE STEPS
one merges some steps in the reverse process, sampling and learning would be faster and more memory efficient. The variance of the training loss could also be reduced. Specifically, at time I of the reverse process, instead of predicting po (X-11 ( we can predict po (Xxe X2) for any S < t. Given a sequence of timesteps 0 to < < tx T, we can now model the reverse process
pa (Xto: X x) = p ( II.
po .
(11)
To align with the literature of diffusion models, call the list [M,,-. ....... Lx) the denoising schedule. After incorporating the denoising schedule Eq. (1O). we obtain the loss function:
Po (o(X.__,
((0) = (Y)
E Ea log logg(Xr((X) - - log ((...)
(12)
Note that although we may not have the analytical form of q (x,-;(Xx;11) .). we can draw samples from it. Merging is feasible if the support of Pe(Xt__, (Xe,) is equal or larger than the support of (X,,, otherwise, the inverse of some forward permutations would be almost surely unrecoverable. Therefore, we can implement a non-trivial denoising schedule (i.e., k < T), when po(o) Xx) follows the PI. or GPI. distribution, as they have whole S, as their support. However, merging is not possible for inverse shuffling methods, as their support is smaller than that of the corresponding multi-step forward shuffling. To design a successful denoising schedule, we first describe the intuitive principles and then provide some theoretical insights. 1) The length of forward diffusion T should be minimal so long as the forward process approaches the uniform distribution. 2) If distributions of X, and X1+1 similar, should merge these two steps. Otherwise, we should not merge them, as it would make the learning problem harder.
To quantify the similarity between distributions shown in 1) and 2), the TV distance is commonly used in the literature. In particular, where (1) is the distribution at time in the forward process and is the uniform distribution. For first introduce the Eulerian Numbers An,r (OEIS Foundation Inc. 2024), i.e., the number of permutations S, that have exactly r rising sequences where 1 - < An, can be computed using the following recursive formula Aur = rAn-1,r + (n -r + 1) An-1r-11 where A1,1 = 1. We then have the following result. The proof is given in Appendix F
measure Drv(q ((')) for t + t' and Drv (q(9), u).
riffle shuffles, the total variation distance can be computed exactly. Specifically,
Proposition 3. Let + t' be positive integers. Then
n + 2
n 2
(1) Drv
IRS SRS
=
2
An.r
2
2tn
(13)
pml
and
DTV
(1) IRS.
Σ
A
1
+
2
1 n!
(14)
2
Note that Eq. [4) was originally given by Kanungo 2220). We restate it here for completeness. Once the Fulerian numbers precomputed, the TV distances can be computed in O(n) time instead
7

---

Published as conference paper at ICLR 2025
TV Distance to Unitrom, n-
0.8
5.8
0.6
22-
15
(a)
(b)
Figure 2: (a) DTV (Ins: u) computed using Eq. (4). We choose T 15 (red dot) based on (t)
threshold 0.005. (b) heatmap for DTv (Ins: Uns for n 100 and 1 S 1' S 15, computed
using Eq. (13). Rows are and columns are t'. We choose the denoising schedule [0, 8, 10, 15].
of O(n)). Through extensive experiments, we have following empirical observation. For principle 1), choosing 7' so that Drv (ARS u) ss 0.005 yields good results. For the principle 2),
(T')
denoising schedule [fa..... ... tu.] with Drv (9RS IRS
(t:) (t:+1) = 0.3 for most i works well. We show an
example on sorting n 100 four-digit MNIST images in Fig.
3.5 REVERSE PROCESS DECODING
We now discuss how to decode predictions from the reverse process at test time. In practice, one is often interested in the most probable state few states with high probabilities under po (") However, since we can only draw samples from Pa (Xoo via running the reverse process, exact decoding is intractable. The simplest approximated method is greedy search, ie., successively finding the mode or an approximated mode of (o(X,, x; I Another approach is beam search, which maintains a dynamic buffer of & candidates with highest probabilities. Nevertheless, for one-step reverse transitions like the GPL distribution, even finding the mode is intractable. To address this, we employ hierarchical beam search that performs an inner beam search within n² scores each step of the outer beam search. Further details are provided in Appendix
EXPERIMENTS
We now demonstrate general applicability and effectiveness of our model through solving a variety of tasks, including sorting 4-digit MNIST numbers, jigsaw puzzles, and traveling salesman problems (TSPs). Additional details, including an additional synthetic experiment that compares our method with other discrete diffusion models, are provided in Appendix G due to space constraints.
4.1 SORTING 4-DIGIT MNIST IMAGES
We first evaluate our SymmetricDiffusers on the four-digit MNIST sorting benchmark, well- established testbed for differentiable sorting (Blondel et al. 2020 Cuturi et al. 2019 Grover et al.)
2018 Kim et al.) 2024 Petersen
2021 2022). Each four-digit image this benchmark
obtained by concatenating 4 individual images from MNIST, and our task is to sort four-digit MNIST
numbers. For evaluation, we employ several metries compare methods, including Kendall Tau
coefficient (measuring the correlation between rankings), accuracy (percentage of images perfectly reassembled), and correctness (percentage of pieces that are correctly placed).
Ablation Study. We conduct an ablation study to verify our design choices for reverse transition and decoding strategies. As shown Table when using riffle shuffles as the forward process, combining PL with either beam search (BS) greedy search yields good results in terms of Kendall-Tar and correctness metrics. In contrast, the IRS (inverse riffle shuffle) method, along with greedy search, performs poorly across all metrics, showing the limitations of IRS in handling complicated sorting tasks. At the same time, combining GPL and BS achieves the best accuracy in correctly sorting the entire sequence of images. Finally, that random transpositions (RT) and random insertions
8

---

Published as a conference paper at ICLR 2025
Table 1: Results (averaged over 5 runs) on solving the jigsaw puzzle on Noisy MNIST and CIFARIO.
Noisy MNIST
CIFAR-I0
Method
Metrics
x
x 4
5 x 5
6 6
2 x 2
3 X
Gumhel-
Kendall-Tan f
0,9984
0.6908
0.3578
0.2430
0.1755
0,8378
0.5044
0.4016
Sinkhorn
Accuracy (%)
99.81
44,65
00.86
0.00
0.00
76.54
6.07
0.21
Correct (%)
99,91
80,20
49.51
26.94
14.91
86,10
43,59
25.31
Network
RMSE
0.0022
0,1704
0.4572
0,8915
1,0570
0.3749
0,9590
1.0960
Menn al
2018
MAE +
0.0003
0.0233
0.1005
0.3239
0.4515
0.1368
0.5320
0.6873
Kendall-Tau f
0.9931
0.3054
0.0374
0.0176
0.0095
0.6463
0,1460
0.0490
Accuracy (%)
99.02
5.56
0.00
0.00
0,00
59.18
0,96
0.00
Dift'Srtt
Currect (%)
99.50
42.25
10.77
6.39
75.48
27.87
12.27
Petersen
RMSE
0.0689
1,0746
1.3290
1,4883
1.5478
0.7389
1.2691
1.3876
ct al. 2022
MAE
0,0030
0,4283
0.6531
0.8204
0,8899
0,2800
0.8123
0.9737
Kendall-Tau
0,9899
0.2014
0.0100
0.0034
-0.0021
0.6604
0.1362
0.0318
Accuracy (%)
98.62
0.82
0.00
0.00
0.00
60.96
0.68
0.00
Error-free Sor (Kim
Correct (%)
99.28
32.65
7.40
4.39
2,50
75.99
26.75
10.33
RMSE 4
0.0814
1.1764
1,3579
1,5084
1.5606
0.7295
1.2820
1.4095
ct al. 2024
MAE +
0.0041
0.5124
0.6818
0.8424
0,9041
0,2731
0.8260
0.9990
Kendall- Tau f
0.9992
0.8126
0,4859
0.2853
0,1208
0,9023
0,8363
0,2518
Symmetric Diffusers
Accuracy (%)
99.88
57.38
1.38
0.00
0.00
90.15
70.94
0.64
Correct (%)
99.94
86.16
58.51
37.91
18.54
92.99
86.84
34.69
(Ours)
RMSF.
0.0026
0.0241
0.1002
0,2926
0,4350
0.3248
0.3892
0.8953
MAE
0,0001
0,0022
0.0130
0.0749
0.1587
0.0651
0.0977
0.5044
Table 2: Results (averaged over 5 rums) on the four-digit MNIST sorting benchmark. For 200,
due to efficiency reasons, we use PL for the reverse process, and we randomly sample a timestep when computing the loss (see Appendix C.2)
Sequence Length
Method
Metries
15
32
52
100
200
DiffSort
Kendall-Tau f
0.930
0.898
0.864
0.801
0.638
0.535
0.341
0,166
0.107
(Petersen ct al. 2022
Accuracy (%)
93.8
83.9
71.5
52.2
10.3
0.2
0.0
0.0
0.0
Correct (%)
95.8
92.9
90.1
85.2
82.3
61.8
42.8
23.2
15.3
Emrrr-ree
K f
0.974
0.967
0.962
0.952
0.938
0.879
0,170
0,140
0.002
DiffSort Kim
Accuracy (%)
97.7
95.3
92.9
89.6
83.1
57.1
0.0
0.0
0.0
al. 2024
Correct (%)
98.4
97.7
97.2
96.3
95.1
90.1
24.2
20.1
0.8
Symmetric
Kendall-Tau f
0.976
0.967
0,959
0.950
0.932
0.858
0.786
0.641
0.453
Diffusers
Accuracy (%)
98.0
95.5
92.9
90.0
82.6
55.1
27.4
4.5
0.1
(Ours)
Correct (%)
98.5
97.6
96.8
96.1
94.5
88.3
82.1
69.3
52.2
(RI) both out of memory for large instances due their long mixing time. Given that accuracy is the most challenging metric to improve, we select riffle shuffles, GPL and BS for all remaining experiments, unless otherwise specified. More ablation study (e.g.. denoising schedule) is provided in Appendix G.3
Full Results. From Table 2 we can see that Error-free DiffSort achieves the best performance in sorting sequences with lengths up to 32. However, its performance declines considerably with longer sequences (e.g., those exceeding 52 in length). Meanwhile, DiffSort performs the worst due to the error accumulation of its soft differentiable swap function (Kim et al. |024 [Petersen et al. (022)). In contrast, our method is on par with Error-free DiffSort sorting short sequences and significantly outperforms others on long sequences.
4.2 JIGSAW PUZZLE
We then explore image reassembly from segmented "jigsaw" puzzles (Mena et al. 2018 Noroozi & Favaro 2016 Santa Cruz et al. 2017). We evaluate the performance using the MNIST and the CIFAR 10 datasets, which comprises puzzles of up to 6 x 6 and 4 x 4 pieces respectively. We add slight noise to pieces from the MNIST dataset to ensure background pieces are distinetive. To evaluate models, we use Kendall- -au coefficient, accuracy, correctness, RMSE (root mean square error of reassembled images), and MAE (mean absolute error) as metrics.
Table ] presents results comparing our method with the Gumbel-Sinkhorn Network (Mena et al. 2018). Diffsort (Petersen et al. 2022). and Error-free Diffsort (Kim et al., 2024). DiftSort and free DiffSort primarily designed for sorting nigh-dimensional ordinal data which have clearly different patterns. Since jigsaw puzzles on MNIST and CIFAR10 contain pieces that are visually similar, these methods do not perform well. The Sunkeo-Sinkhorn performs better for tasks

---

Published as a conference paper at ICLR 2025
Table 3: Ablation study transitions of reverse diffusion and decoding strategies. Results averaged over three runs on sorting 52 four-digit MNIST images. GPL: generalized Plackett-Luce IRS: inverse riffle shuffle; RT: random transposition; IT: inverse transposition; RI: random insertion; II: inverse insertion.
Forward
Riftle Shuffles
Reverse
GPI
GPI. Greedy
Greedy
BS
IRS Greedy
Greedy
Greedy
Kendall-Tau |
0.786
0.799
0.799
0.797
0,390
Accuracy (%)
27.4
24.4
26.4
26.4
0.6
Out of Memory
Correct (%) 82.1
83.1
Table 4: Results on TSP-20 and TSP-50. We compare our method with OR solvers such as Concorde
(Applegate et al. 2006), LKH-3 (Helsgaun) (017), and 2-Opt (Lin & Kernighan 1973), ) as well as learning-based approaches including GCN (Joshi al. 2019) and DIFUSCO (Sun & Yang 2023) on 20-node and 50-node TSP instances. An asterisk (*) indicates that
heuristies were
removed to ensure fair comparison. Baselines for TSP-50 are taken from [un & Yang (2023)
TSP-20
TSP-50
Method
Tour Length
Optimality Gap (5)
Tour Length 4
Optimality Gap (%)
Concorde
3.84
0.00
5.69
0.00
OR Solvers
LKH-3
3.84
0.00
5.69
0.00
2-Opt
4.02
4.64
5.86
2.95
Learning-
GCN
3.85*
0.21
5.87
3.10
Based
DIFUSCO
3.88*
1.07
5,70
0.10
Models
Ours
3.85
0.18
5.71
0.41
involving fewer than 4 x pieces. In more challenging scenarios (e.g., 5 and 6 x 6), our method significantly outperforms all competitors.
4.3 THE TRAVELLING SALESMAN PROBLEM
At last, we explore the travelling salesman problem (TSP) to demonstrate the general applicability of our model. TSPs are classical Nocoopee combinatorial optimization problems which are solved using integer programming or heuristic solvers (Arora 1998, Gonzalez 2007). There exists vast
literature on learning-bast based models to solve TSPs (Kipf & Welling 2017 Kool et al. 2019 Joshi
al. 2019 2021 Bresson & Laurent 2021 K Kwon al.] 2021] Fu et al. [202]] Quu
2022 Kim
2023 Sun & Yang 2023] Min et al. 2024] Sanokowski et al. (2224). They often focus on the Euclidean TSPs, which are formulated as follows. Let V = (VI,.... Un be points in IR2 We need to find some E S, such that Σ ((ass)) (o(++t))l2 2 is minimized, where we let ( + 1) = :1))
ct al.
We compare with operations research (OR) solvers and other learning based approaches on TSP instances with 20 nodes and 50 nodes. The metrics are the total tour length and the optimality gap. Given the ground truth (GT) length produced by the best OR solver, the optimality gap given by (predicted length- (GT length)) /GGT length). As shown in Table SymmetricDiffusee achieves comparable results with both OR solvers and the state-of-the-art learning-basee methods. Further experiment details are provided in Appendix G
CONCLUSION
In this paper, we introduce novel discrete diffusion model over finite symmetric groups. We identify the riffle shuffle as an effective forward transition and provide empirical rules for selecting the diffusion length. Additionally, we propose a generalized PL distribution for the reverse transition, which provably more expressive than the PL. distribution. We further introduce theoretically grounded "denoising schedule" to improve sampling and learning efficiency. Extensive experiments verify the effectiveness of our proposed model. Despite significantly surpassing the performance of existing methods on large instances, our method still has limitations in larger scales. In the future, we would like to explore methods to improve scalability even more. We would also like to explore how can fit other modern techniques in diffusion models like concrete scores (Meng et al.) (023) and score entropy (Lou et al. 2024) into our shuffling dynamics. Finally, we are interested in generalizing model to general finite groups and exploring diffusion models Lie groups
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their helpful comments. This work was funded, in part, by the NSERC DG Grant (No. RGPIN-2022-04636, the Vector Institute for AI, and Canada CIFAR AI Chair. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through the Digital Research Alliance of Canada alliance. can ca and companies sponsoring the Vector Institute www. , vector insiinsittt a #partners and Advanced Research Computing at the University of British Columbia. Additional hardware support was provided by John R. Evans Leaders Fund CFI grant.
REFERENCES
David Applegate, Robert Bixby, Vasek Chvatal, and William Cook. Concorde tsp solver, 2006. URL https://wne /twp://wwwwww. math. uwaterloo http://......ppcooooeeeeee/e/e// http://occooennnneeooo /ineex html
Sanjeev Arora. Polynomial time approximation schemes for euclidean traveling salesman and other geometric problems. J. ACM, 45(5):553-782, Sep 1998. ISSN 0004-5411, doi: 10.1145/290179.
290180. URL/tttps: /tte:////oooie/ http:///00..rg/111111/////////// 7990111197102218 290180
Jacob Austin, Daniel Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2023.
Dave Bayer and Persi Diaconis. Trailing the dovetail shuffle to its lair. The Annals of Applied Probability, 2(2)2944 - 313, 1992. doi: http///aapp11770057000 URL hRtpstts://ddo org/ 10. 121 (t.12214/aaap/177005777
Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning, pp. 950-959. PMLR, 2020.
Xavier Bresson and Thomas Laurent. The transformer network for the traveling salesman problem, 2021.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models, 2022. URL
hitts:///arxlv wrg. ht.rg/abs/2203. 14987
Zhe Cao. Tao Qin, Tie- Yan Liu, Ming-Feng Tsai, and Hang Li. Learning to rank: from pairwise approach listwise approach. In Proceedings of the 24th International Conference on Machine Learning, ICML '07, pp. 129-136, New York, NY, USA, 2007. Association for Computing Machinery. ISBN 9781199937933.. doi: 2..1145//733496.2273311 Utt.https://doi. org//
10 10.1145/1273496.122
Marco Cuturi. Olivier Teboul, and Jean-Philippe Vert. Differentiable ranking and sorting using optimal transport. Advances in neural information processing systems, 32, 2019.
Persi Diaconis. Group representations in probability and statistics. Lecture notes monograah series, 11:i-192, 1988.
Zhang-Hua Fu, Kai-Bin Qiu, and Hongyuan Zha. Generalize a small pre-trained model to arbitrarily large tsp instances, 2021.
N. Gilbert. Theory of shuffling. Bell Telephone Laboratories Memorandum, 1955.
Teofilo F. Gonzalez. Handbook of Approximation Algorithms and Metaheuristics (Chapman & Hall/Cre Computer & Information Science Series). Chapman & Hall/CRC, 2007. ISBN 1584885505.
Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic optimization of sorting networks via continuous relaxations. In International Conference on Learning Representations, 2018
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL gurobi ccm
https: www
11

---

Published as a conference paper at ICLR 2025
Keld Helsgaun. An extension of the fin-kernighan-helsaee tsp solver for constrained traveling salesman and vehicle routing problems, Dec 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
Chaitanya K. Joshi, Thomas Laurent, and Xavier Bresson. An efficient graph convolutional network
technique for the travelling salesman problem, 2019.
Chaitanya a K Joshi, Quentin Cappart, Louis-Martin Rousscau, and Thomas Laurent. Learning tsp requires rethinking generalization. In International Conference on Principles and Practice of Constraint Programming. 2021.
Shihan Kanungo. Mixing time estimates for the riffle shuffle. Euler Circle, 2020.
Jungtack Kim, Jeongbeen Yoon, and Minsu Cho. Generalized neural sorting networks with error- differentiable swap functions. In International Conference on Learning Representations (ICLR), 2024.
Minsu Kim, Junyoung Park, and Jinkyoo Park. Sym-nco: Leveraging symmetricity for neural combinatorial optimization, 2023.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks, 2017.
Wouter Kool, Herke van Hoof, and Max Welling. Attention, learn to solve routing problems!, 2019.
Yeong-Dae Kwon, Jinho Choo, Byoungjip Kim, Iljoo Yoon, Youngjune Gwon, and Seungjai Min. Pomo: Policy optimization with multiple optima for reinforcement learning, 2021.
Shen Lin and Brian W Kernighan. An effective heuristic algorithm for the travelling-salesman problem. Operations research, 21(2):498-51) 1973.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution, 2024. URL
R. D. Luce. Individual Choice Behavior. John Wiley, 1959.
Colin L Mallows. Non-null ranking models. i. Biometrika, 44(1/2):114-330 1957.
Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permutations with gumbel-sinkhorn networks. In International Conference Learning Representations, 2018. 88
Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data, 2023. URL httpp:://axxvv.oor///aabtt
00802
Yimeng Min, Yiwei Bai, and Carla P. Gomes. Unsupervised learning for solving the travelling salesman problem, 2024.
Mchdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. In European conference on computer vision, pp. 69-84. Springer, 2016.
OEIS Foundation Inc. The culerian numbers, entry a008292 in the On-Line Encyclopedia of Integer Sequences, 2024. Published electronically at httpt:tt////e..// ort.tlosrrr/A0022900
Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Differentiable sorting networks for scalable sorting and ranking supervision. In International conference on machine learning, pp-
6555-8555. PMLR. 2021.
Felix Petersen, Christian Borgelt, Hilde Kuchne, and Oliver Deussen. Monotonic differentiable sorting networks. In International Conference Learning Representations (ICLR), 2022.
R. I.. Plackett. The analysis of permutations. Applied Statistics, 24(2):193 - 1975.
12

---

Published as a conference paper at ICLR 2025
Ruizhong Qiu, Zhiqing Sun, and Yiming Yang. Dimes: A differentiable meta solver for combinatorial optimization problems, 2022.
J. Reeds. Theory of shuffling. Unpublished Manuscript, 1981.
Laurent Saloff-Coste. Random Walks on Finite Groups, pp. 263-346. Springer Berlin Heidelberg, Berlin, Heidelbbrrg 2004. ISBN 978-3-662-004444444 09444-0. doi: 10.10079988.3-66--004444
URL
https://doi. 1ttp:///11110000---1111331 http://.....2.2200000/./.665566.
Sebastian Sanokowski, Sepp Hochreiter, and Sebastian Lehner. A diffusion model framework for abs
unsupervised neural combinatorial optimization, 2024. URL
2406 01661
Rodrigo Santa Cruz, Basura Fernando, Anoop Cherian, and Stephen Gould. Deeppermnet: Visual permutation learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3949-3957, 2017.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K. Titsias, Simplified and generalized masked diffusion for discrete data, 2024. URL
2406. 04329
Jascha Sohl- Dicksseii, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. 2256-22665 Lille, France, 07-09 Jul 2015. PMLR. URL https:// proceedings mlm.
www.einns.........mm..m........m html
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations, 2021.
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based tone discrete diffusion models, 2023. bttt::::: http://..wrraaaa//////.../
Zhiqing Sun and Yiming Yang. Difusco: Graph-based diffusion solvers for combinatorial optimiza- tion, 2023.
Harshit Varma, Dheeraj Nagaraj, and Karthikeyan Shanmugam. Glauber generative model: Discrete diffusion models via binary classification, 2024. URL https://azxiv.oqq/
17035
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need, 2023.
Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation, 2023.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3-4):229-256 1992.
13

---

Published as a conference paper at ICLR 2025
A ADDITIONAL DETAILS OF THE GSR RIFFLE SHUFFLE MODEL
There are many equivalent definitions of the GSR riffle shuffle. Here we also introduce the Geometric Description (Bayer & Diaconis T992), which is easy to implement (and how implement riffle shuffles in our experiments). We first sample points in the unit interval [0, [) uniformly and independently, and suppose the points labeled in order as ø1 x2 < .... < In- Then, the permutation that sorts the points {2x1} .......... (2z,,} follows the GSR distribution, where {o} := [ is the fractional part of æ.
B DETAILS OF OUR NETWORK ARCHITECTURE
We now discuss how to use neural networks to produce the parameters of the distributions discussed Section 3 and 3.2.2 time I, and suppose X; = (x 1.... .....,,.))) xi=)) E Raad Let encoder, be an object-specific encoder such that encodero E IR nxddda For example, encodery can be a CNN if X, an image. Let
Y == (X.. + time_
(y 1.... (11)
E IR Tedda
(15)
where t time_ _ereeeebb_____ is the sinusoidal time embedding. Then, we would like to feed the embeddings into Transformer encoder (Vaswani et al. 2023). Let t woensffrrer_eeccdae ry be the encoder part of the Transformer architecture. However, each of the distributions we discussed previously has different number of parameters, we will have discuss them separately.
Inverse Transposition. For Inverse Transposition, we have n + parameters. To obtain a + 1 tokens from transformer _ecordeen_c,,, append dummy token of 0's to Y. Then we input (t) ,...,,... yn .)) into gransformer_encon to obtain Z E [In+ Finally, we apply an MLP to obtain (81,..., E R
Inverse Insertion, Inverse Riffle Shuffle, PL Distribution. These three distributions all require exactly n parameters, so we can directly feed Y, into framssfrmer_eemm of transformer_encome be Z E Rn where we then apply MLP to obtain the scores Sy E R"
Let the output
The GPL Distribution. The GPL distribution requires n² parameters. We first append n dummy tokens of O's to Y,, with the intent that the jth dummy token would learn information about the jth column of the GPL parameter matrix, which represents where the jth component should be placed. When com-
We then pass (yl .....yn y(ll, ........)) R dn transformer_encaee
puting attention, we further apply 2n x 2n attention mask
x
-00
A B
x
M =
where A is an n x matrix of - oo, B
is x n.
The reason for having B as an upper triangular matrix of -oo is that information about the jth component should only require information from the previous components. Let
7 Zz
Srensfornerrenncoeee..
where Z2 E Raxdeas Finally, we obtain the GPL parameter matrix as Sp = 2, E Ruxn
For hyperparameters, refer the readers Appendix G.5
14

---

Published as a conference paper at ICLR 2025
C DISCUSSIONS ON OTHER FORMS OF THE LOss
C.I USING KL DIVERGENCE
Many diffusion models will rewrite the variational bound Eq. 10 the following equivalent form of
KL divergences to reduce the variance (Austin et al.) 2023 Ho et al. 2020):
Exmi(Xo X Xo) (ia(o(X,(Xo) ((Xrl))
Σ
|| po (X DK. f>1
po Po
(16)
However, we cannot use this objective for Sn in most cases. In particular, since
X() x()x(Xx-1)
(17)
we can only derive the analytical form of q(X,-1|X, Xo) if we know the form of q( (o). However, (o)) is unavailable for most shuffling methods used in the forward process except for the riffle shuffles. For riffle shuffle, (X,) is actually available and permits efficient sampling (Bayer Diaconis, 1992). However, DKI. (g(Xx-1|Xx, Xo) || po(Xx_1(Xx)) still does not have an analytical form, unlike in common diffusion models. As result, we cannot use mean/score parameterization (Ho et al. 20200] Song et al.) 202I) commonly employed in the continuous setting. Therefore, need to rewrite the KL term as follows and resort Monte Carlo (MC) estimation,
DKL (KL
po (X
((,))xo) log g( Po Po(X,_)) (x))
= ((x,
9(X)-1lXo) ((X,[[x-1)
log ((x,11 |x,, Xp)
= (xxx)))))
((Xx|{xo) po(Xi-1 (x,)
(g(x+1X,-1) log q(x,- Xn) x()
= [x) ByCX Xo)
(18)
Po(X.-1(X.)
Note that X, ~ ((,, Xp) and X,-1 ((X,-1(X))) are drawn independently. However, there is high chance that (X,1) = for the X, and X,-1 that are sampled. Consequently, we only draw a few MC samples, the resulting estimator will likely be zero with zero-valued gradients, impeding the optimization the training objective. Therefore, writing the loss the form of divergences does not help in the case of discrete diffusion on S..
C.2 SAMPLING RANDOM TIMESTEP
Another technique that many diffusion models use is to randomly sample timestep t and just compute the loss at time t. Our framework also allows for randomly sampling one timestep and compute the loss as
E Bys (No (X) Ey (xx [- log log (Xc-
(19)
omitting constant terms with respect to 0. With denoising schedule of [to..... .0, :t],, the loss is
E E Xo, E,,y |No) Extx, x. log Po (X.,...
(20)
again omitting constant terms with respect 8. It is also worth noting that computing the loss subset of the trajectory could potentially introduce more variance during training, which leads to a tradeoff. For riffle shuffles, although can sample X,-1 directly for arbitrary timestep from
15

---

Published as a conference paper at ICLR 2025
Xy previously mentioned in this section, the whole trajectory would be really short. For other shuffling methods, we would still have to run the entire forward process to sample from g(Xe-1|Xo) which does not solve the inefficiency problem of other shuffling methods. Therefore, we opt to use the loss in Eq.(O 10 in most cases, and we would resort to Eq. 20 for riffle shuffling really large instances.
D ADDITIONAL DETAILS OF DECODING
Greedy Search. At each timestep & in the denoising schedule, we can greedily obtain or approx- imate the mode of (X,, We then the (approximated) mode X for the next timestep pa( (Xe, 2lx,, )). Note that the final Xo obtained using such a greedy heuristic may not necessarily be the mode of (.).
beam search improve the greedy approach. The basic idea is that, at each timestep 1 in the denoising schedule, we compute or approximate the pop-k-most-probabbe results from po Xt, (Xx.)). For each of the top-k results, sample top-k from pe (X X t; 2)- Now we have k2 candidates for X and we only keep the top k of the k2 candidates.
Beam Search. We
However, it is not easy to obtain the rop-k-most-probabl results for some of the distributions. Here we provide algorithm approximate top-k of the PL and the GPL distribution. Since the PL distribution is strict subset of the GPL distribution, it suffices to only consider the GPL distribution with parameter matrix S. The algorithm for approximating top-k of the GPL distribution is another beam search. We first pick the & largest elements from the first row of S. For each of the k: largest elements, we pick & largest elements from the second row of S, excluding the corresponding element picked in the first row. We now have 2²2 candidates for the first two elements of permutation, and we only keep the top-k-most-probabbe candidates. We then continue in this manner.
E THE EXPRESSIVENESS OF PL AND GPL
In this section, prove the expressiveness results of the PL and GPL distribution. We first show that the PI. distribution has limited expressiveness.
Proposition The PL distribution cannot represent a delta distribution over Sn-
Proof. Assume for a contradiction that there exists some a E Sn and such that PL = dg. Then we have
n exp (So(i)
II
1.
Eド= exp (8o(i)
Since each the term in the product is less than or equal to 1, we must have
exp (so
(21)
Σj= exp (8o(3)
for all E [n]. In particular, have
exp (8a(1))
Σド= exp (sa(i))
which happens if and only if 8o()) 00 for all j 2 2. But this contradicts 21
We then prove that the reverse process using the GPL distribution can model any target distribution. We first introduce two lemmas.
Lemma 4. The GPL distribution can represent any delta distribution on S, if allowing - oo scores.
Proof. Fix o E Sn. For all E [n], we let (1)) = O and si.j = -0 for all j } o()). Then it is clear that GPI,,,,) 8z.
16

---

Published as a conference paper at ICLR 2025
1/4
1/2
2
3)
33
1)
1/8
1/4
8/4
3/4
Tops
1/2
Figure 3: A simple example for the GPL expressivees theorem on S3-
Lemma 5. Let E Sn and let == n) where (n-1 n) a transposition that swaps the last two indices and o is function composition. That is, π is obtained from by swapping the last two components. Let p be any probability distribution on Sn whose support subset of {a, x). Then there exists scores (8ig)sJE(m) (possibly -oo) such that GPI)).,)) P.
Proof. Note that we have
2
7- (n)
71
क = ((1) (2)
a onn 1)
For all < 2, we let (1))(() = O and 8i, = 0o for all j * (1).
- Let 8n-1,o(n-1) = Inp(o)
and let 8,-1,o(n) = In p(r). Finally, let sngj 0 for all j E [n]. It is then easy to verify that GPL(,.))) p.
We then state the main expressiveness theorem.
Theorem 6. Let Yo E Sn, be a random variable with arbitrary distribution q(Y)). Let be any distribution over Sa- Then there exists l E N and random variables Y1,---, Y& with GPL (allowing - oo scores) transition distributions q(Y; Y,11) for each i E [k] such that q(Y)) p.
Before proceeding to the proof, we first provide a small example illustrating the construction we are going to use. Consider Fig. Jn Sg. Suppose we start with g gY)) being the uniform distribution, and the target distribution p is listed at the top row of the diagram. We partition Sg into 3 pairs (a,, R;) indicated by their color in the diagram. The permutations within each pair differ by last two indices, so each pair is the pair considered in Lemma The first step is to concentrate all probability mass æi using GPL. and Lemma 4 Now note that only need 1/4 + 1/2 3/4 probability for the first pair (a1, TT), so there is 1/4 excess. We then use Lemma 5 to move the excess amount to σ1. Then we use Lemma to move the excess amount of pair T2 pair Finally, we use Lemma 5 distribute the correct mass to σj and πj from the 3/4 that π currently holds. We now present the formal construction.
swap of the
Proof. For σ E Sn. say that and æ are pair if (n n) a.. Note that can partition S, into n1/2 disjoint pairs. We write the pairs as (11, (an1/2, kn1/2).
We now give algorithm that explicitly constructs the transitions from Yo to the target distribution p. The intuition of the algorithm that we distribute the probability mass for one pair permutations time. The variable Pietover records how much mass have yet to distribute.
17

---

Published as a conference paper at ICLR 2025
Let Pieltover = Let g(1) Yo) 8x.-
2. Iterate through all pairs (0i, Ti) for i from 1 to n2/2:
(a) Let Pexcess Pleftoer p(o,) p(TT:).
(b) Define (Y 2) such that
Pexcess
0( oi Y3i = = ni)
Piatover
(Y = Mi 13i Tz)
Pexcess
Pietover
and q\Yii_ = Y3-- ns) 0 for all other & (o,, Let (Y
Yxi- =) = 6, for r $ {
(c) Define Ys--1) such that (Yss Ys--1 = o;) 8mzt and 9 Y3i-1
7) = 6, for all * di.
(d) Finally, define g(Y3(11 Y3.) such that
((33+1 = Yys = mT) p(o.)
Pexcess
((33+11 = n, Yy: = mT)) p(TT.)
Pexcess
and 31+1 Y3 n;)) 0 for all other & (o,, m. Let (Ya441 |Ys r) = 8r for $ { (e) Update Pletover Pexcess.
Return 4(Y Yo), (( Y3n1/2).
Note that q(1 = T = 1. Also note that by Lemma and 5 all transition distributions can be modeled by the GPL distribution. We claim that the end of iteration we must have
n2/ j=i+i p(o,)) p(TT));
(1) Pleftover
(2) g(Yzi41 = oi) = p(o;) and g(Y3i+1 = T.) = p(TT);;
(3) g(Ys;41 =T++1) Pietover
We proceed by induction on For is clear that Pleftover the first iteration. We also note that qY2 = σ1) Pexcess and (122 11) =1 =1-Pesees = p(
n!/ -i p(a,) + p(T,) at the end of
p(TT). After step 2(c) of the algorithm, we have g(Y33 = σ1) = 0, gY3 TI) =p(a)) + =p(T)) and
<(33 = T2) = Pexcess Finally, after step 2(d) and 2(e), we get (Y4 01) = =p(an), (Y4 = m)
p(T1), and qY₁ )2)
Plettover-
For the inductive step, let 2 2. We know by the inductive hypothesis that at the start of iteration 2t/2
i, we have Plettover = ΣØd p(a,) + p(TT). So 0 < Pexcess 1, and all transition distributions in iteration are well-defined. is easy to verify that:
After step 2(b), 4(Y3i-1 σ.) Pexcess and g(Ya)(( = n.) p(ox) p(m.))
After step 2(c), q(Ys) = ..) = 0, qYss = nT)) = p((.) + p(T.), and q(sss = T:t1) =
Pexcess
After step 2(d)(e), g(Ys441 = σi) = p(o.), 9(v = m)) = p(T)), and g(Ya441 = n2/2
7;-1) Pleftoer ()))t1((a,) p(T,).
18

---

Published as a conference paper at ICLR 2025
The Piettover the end of iteration m1/2 is Peacess p((ont/2) - Finally, we observe that after iteration i, :((, = oz)) and g(Y3 = mT) will never be changed for j 2 i. This finishes the proof.
m2/2)) = 0. This finishes the induction.
Finally, Theorem 2 which is stated in the main paper, follows immediately from Theorem by setting (Y)) to be the uniform distribution and p to be the target distribution over Sn.
Theorem 2. The reverse diffusion process parameterized using the GPL distribution in Eq. 9 can
model any distribution over S..
F RESULTS ON TV DISTANCES BETWEEN RIFFLE SHUFFLES
Proposition 3. Let t & t' be positive integers. Then
(0) (e))
1
12
+
n + 2
|
Drv ARS IRS
An.
(13)
2
2
2
r=1
and
2
Drv
Ins: v
(2)
=
An,
2
72
n!!
(14)
Proof. Let E Sn. It was shown in Bayer & Diaconis (1992) that
(8) Ins
1
n2 + 2
2
where is the number of rising sequences of Note that if two permutations have the same number of rising sequences, then they have equal probability. Hence, we have
(t) Drv (ARS 9RS
(t))
Σ
(e) IRS
(a)
9RS (a)
(t')
1
An.
9RS
(1)
9RS
(t')
(a)
2
T
Σ
An,
n2 2
r
1
+
2
2
n
2t'
n
as claimed. For 14 replace TRS (o) with u(()) = n!
1 in the above derivations.
G ADDITIONAL DETAILS ON EXPERIMENTS
G.1 ADDITIONAL SYNTHETIC EXPERIMENT: LEARNING A SINGLE PERMUTATION LIST
We present an additional synthetic experiment to illustrate the effectiveness of Dyffetrri A natural approach to modeling permutations is to represent them as sequences of numbers from {0, L.I....,,, n - 1} and apply sequence generation models. However, existing sequence generation models relax the permutation constraint, allowing numbers to repeat and modeling transitions in larger sequence space. To highlight the advantage, i.e., restricting the diffusion trajectory to S, and modeling transitions in a smaller space (O(n!) vs O(n")). )) of our approach, we conduct a synthetic experiment: learning a delta distribution over Sn, i.e., single permutation. particular, we compare our method with SEDD (Lou et al. 2024), one of the strongest discrete diffusion models.
In the experiment, we train models to learn the identity permutation (n 100) and a fixed arbitrary distribution (n 100 and 200). For details on hyperparameters and training, please refer to Appendix G.5. From Table 5 we see that our method reaches 100% accuracy in all settings. While SEDD also performs well, the gap between its accuracy and ours grows with sequence length. These experiments verify the advantage and the effectiveness of our method.
We would also like to highlight another key limitation of other discrete diffusion models (including SEDD). In fact, it is nearly impossible for other discrete diffusion models to solve the tasks introduced in our paper, including the jigsaw puzzle, sorting multi-digit MNIST numbers, and the TSP. The
19

---

Published as a conference paper at ICLR 2025
Table 5; Results on the toy experiment of learning a single permutation list.
Method
Identity
100
Arbitrary Permutation 100
Arbitrary Permutation 200)
Accurncy (%)
Correct (5%
Accuracy (9)
Correct (5)
Accuracy (%)
Correct (%)
SEDD
Lou al. 2024
95.47
99.95
93.24
99.93
88.75
99.94
100
100
100
reason is that all prior discrete diffusion models assume a fixed alphabet or vocabulary, and they model categorical distributions on the fixed alphabet. For example, in NLP tasks, the vocabulary is predefined and fixed. However, the alphabet is the set of all possible image patches for image tasks such as jigsaw puzzles and sorting MNIST numbers. It is impractical to gather the complete alphabet beforehand. We could potentially train VQVAFs to obtain quantized image embeddings. However, such an approach introduces the approximation error from the quantized alphabet. For the TSP, each node of the graph is point in the continuous space IR, it is also impossible gather the complete alphabet. In contrast, our method can be successfully applied to these tasks because we model a distribution the fixed alphabet Sn, and we treat each permutation function that can be applied to an ordered list of objects.
G.2 DATASETS USED IN THE FULL PAPER
Jigsaw Puzzle. We created the Noisy MNIST dataset by adding i.i.d. Gaussian noise with a mean of 0 and standard deviation of 0.01 to each pixel of the MNIST images. No noise added to the CIFAR-10 images. The noisy images are then saved as the Noisy MNIST dataset. During training, each image is divided into X n patches. A permutation is then sampled uniformly at random to shuffle these patches. The training set for Noisy MNIST comprises 60,000 images, while the CIFAR-10 training set contains 10,000 images. The Noisy MNIST test set, which is pre-shuffled, also includes 10,000 images. The CIFAR-10 test set, which shuffles images on the fly, contains 10,000 images well.
Sort 4-Digit MNIST Numbers. For each training epoch, we generate 60,000 sequences of 4-digit MNIST images, each of length n, constructed dynamically the fly. These 4-digit MNIST numbers are created by concatenating four MNIST images, each selected uniformly at random from the entire MNIST dataset, which consists of 60,000 images. For testing purposes, we similarly generate 10,000 sequences of n 4-gigit MNIST numbers on the fly.
TSP. We take the TSP-20 and TSP-50 dataset fromJoshi et al. (2021) [ The train set consists of 1,512,000 graphs, where each node is an i.i.d. sample from the unit square [0, 1]2. The labels are optimal TSP tours provided by the Concorde solver (Applegate al.) 2006). The test set consists of 1,280 graphs, with ground truth tour generated by the Concorde solver as well.
G.3 ABLATION STUDIES
Choices for Reverse Transition and Decoding Strategies. As demonstrated in Table 6 we have explored various combinations of forward and inverse shuffling methods across tasks involving different sequence lengths. Both GPL and PL consistently excel in all experimental scenarios, highlighting their robustness and effectiveness. is important note that strategies such random transposition and random insertion paired with their respective inverse operations, are less suitable for tasks with longer sequences. This limitation attributed to the prolonged mixing times required by these two shuffling methods, a challenge that is thoroughly discussed in Section 3.1.2
Denoising Schedule. We also conduct an ablation study on how we should merge reverse steps. As shown in Table 7 the choice of the denoising schedule can significantly affect the final performance. In particular, for n = 100 on the Sort 4-Digit MNIST Numbers task, the fact that [0, 15) has accuracy justifies our motivation to use diffusion to break down learning into smaller steps. The result we get also matches with our proposed heuristic Section 3.4
20

---

Published as a conference paper at ICLR 2025
Table More results on sorting the 4-digit MNIST dataset using different combinations of forward process methods and reverse process methods. Results averaged over 3 runs with different seeds. RS: riffle shuffle; GPL: generalized Plackett-Lucee IRS: inverse riffle shuffle; RT: random transposition; IT: inverse transposition; RI: random insertion; II: inverse insertion.
Sequence Length
32
52
Denoising Schedule Kendall-
30,3,5. 5, 9) 0,948
[0,5, 7, 12 0.857
[0, 5, 6, 10, 13) 0.779
RS (forward) + GPL. (reverse) + greedy
Accuracy (5)
89.4
54.8
24.4
Correct (%)
95.9
88.1
81.6
Denoising Schedule
30,3,5, 9) 0.953
[0,5, 7, 12] 0.867
[0, 5, 6, 7, 10, 13]
Kendall- -au
0.799
RS (forward) + PL (reverse) + greedy
Accuracy (%)
90.9
26.4
Correct (%)
96.4
Denoising, Schedule Kendall-Tau Accuracy (5)
3, 9] 0.955
[0, 5, 7, 12] 0.869
[0, 5, 6, 7, 10, 13]
0.797
RS (forward) + PL (reverse) beam search
91.1
57.2
26.4
Correct (%)
96.5
89.2
83.1
T
9
12
13
Kendall-
0.947
0.794
0.390
RS (forward) + (reverse) + greedy
Accuracy (%)
88.6
24.4
0.6
Correct (%)
95.9
82.5
44.6
T (using approx. y logn)
15
55
105
Kendall- Tau
0.490
RT (forward) + (reverse) greedy
Accuracy (%)
18.0
Out of Memory
Correct (%)
59.5
T (using approx. log m)) Kendall-Tau
25
0.954
RI (forward) Il (reverse) + greedy
Accuracy (%)
91.1
Out of Memory
Correct (%)
96.4
Table 7: Results of sorting 100 4-digit MNIST images using various denoising schedules with the combination of RS, GPL and beam search consistently applied.
Denoising Schedule
[0, 15]
[0, 8, 9, 15]
[0,7, 8, 8,9, 15]
[0, 8,,78810,1
[0, 8, 10, 15]
Kendall-Tar f
0.000
0.316
0,000
0.000
0.646
Accuracy (%)
0.0
0.0
0.0
0.0
4.5 69.8
Correct (%)
1.0
39.6
1.0
1.0
G.4 LATENT LOSS JIGSAW PUZZLE
In the original setup of the Jigsaw Puzzle experiment using the Gumbel-Sinkhorn network (Mena et al. 2018), the permutations MSE loss and does not use the ground truth permutation label. However, our loss function (12 actually (implicitly) uses the ground truth permutation that maps the shuffled image patches their original order. Therefore, for fair comparison with the Gumbel-Sinkhorn nerwork in the Jigsaw Puzzle experiment, we modify our loss function so that it does not use the ground truth permutation. Recall from Section 3.2 that we defined
latent. That is, the loss function in Gumbel-Sinkhorn is lexel-level
Po(X
Σ D4) GET
(22)
In our original setup, we defined p(X,-1 of) delta distribution 8(X,-1 Qo; Xt), but this
would require that we know the permutation that turns X,-1 to X,, which part the ground truth. So instead, parameterize XX,, of) as Gaussian distribution A \(, I). At the same time, we note that to find the gradient (12). suffices to find the gradient the log of (22). We use the REINFORCE trick (Williams) 1992] to find the gradient of log (X Xx),, which gives
21

---

Published as a conference paper at ICLR 2025
us
Vo log (X
1
E o ofET
Σ p(X, -- (
d,ET
p(X
Σ Xe
TCT
diET
(x)) p( X,, 0f Vo Vogoooloooonnnn log
E Eoo(o((Xi) p(X p(Xe-1XXe
(X
Va log (ot")) (x=")
-m= P (X1-1|X,, (m))
where we have used Monte-Carle estimation in the last step, and o, ..... (N) pa (X2). We further add an entropy regularization term -A- -poloa, [log pe(or (o [x,]] to each of log (X.-)l)))) Using the same REINFORCE and Monte-Carlo trick, we obtain
NV
V (-X-Bpo log po
Alog pe (x. Vo log (o.")lx)).
~
where (1))
of (N) ~ (po(o. Therefore, we have Vo log (X - -A Exo(o..x.) (os
P (X of(m)
N
Alog (o.
Vo Pologpo (x."
(23)
Em (Xe-1lXe,, of
weight
where of .....o, oz (N) ~ Pe(o,)(e)) then substitute in
of =N (xx
for all [N]. Finally, we also subtract the exponential moving average wei ght as control variate for variance reduction, where the exponential moving average is given by ema <- ema_rate ema + (1 ema ema_rate) weight for each gradient descent step.
G.5 TRAINING DETAILS AND ARCHITECTURE HYPERPARAMETERS
Hardware. The Jigsaw Puzzle and Sort 4-Digit MNIST Numbers experiments are trained and evaluated on the NVIDIA A40 GPU. The TSP experiments are trained and evaluated on the NVIDIA A40 and A100 GPU.
Learning Single Permutation List. The SEDD model we use in our experiments has about 25M parameters. use 7 encoder layers, 8 heads, model dimension 512, nerd-rd dimension 2048, and dropout 0.1. We use the uniform transition for the forward process following the original work.
Our model only has about 2M parameters. We use encoder layers, heads, model dimension 128, feed-forward dimension 512, and dropout 0.I. We also a learned embedding for the numbers. For the diffusion process, we use a denoising schedule [0, 8,,8, 10, 15) for n = 100 and [0, 9, 10, 12] for 200.
For = 100, we use a batch size of 512 and 30K training steps on both methods. For n = 200, we batch size of 128 and 30K training steps on both methods. For performance evaluation, randomly sample 2560 sequences for SEDD and 2560 permutations for our method, and we perform their respective decoding processes.
22

---

Published as a conference paper at ICLR 2025
Jigsaw Puzzle. For the Jigsaw Puzzle experiments, use the AdamW optimizer (Loshchilov & Hutter 2019) with weight decay 1e-2, E = 1e-9, and 3 = (0.9, 0.98). We use the Noam learning rate scheduler given in (Vaswani et al. 2023) with 51,600 warmup steps for Noisy MNIST and 46,000 steps for CIFAR-10. We train for 120 epochs with a batch size of 64. When computing the loss (12). use Monte-Carle estimation for the expectation and sample trajectories. For REINFORCE, sampled 10 times for the Monte-Cario estimation in (23), and we used an entropy regularization rate = 0.05 and an ema ema_rate of 0.995. The neural network architecture and related hyperparameters are given in Table B The denoising schedules, with riffle shuffles as the forward process and GPL as the reverse process, are give in Table 9 For beam search, we use beam size of 200 when decoding from GPL, and we use a beam size of 20 when decoding along the diffusion denoising schedule.
Table Jigsaw puzzle neural network architecture and hyperparameters.
Layer
Details
Output channels 32, kernel size 3, padding 1, stride 1
Convolution
Batch Normalization
ReLU
Max-pooling
Pooling 2 Output dimension (din_after_com + 128)/2
Fully-connected
ReLU
Fully-comnected
Output dimension 128
layers, 8 heads, model dimension (dmoseel 128, feed-forward dimension 512, dropout 0.
Transformer encoder
Table Denoising schedules for the Jigsaw Puzzle task, where we use riffle shuffle in the forward process and GPL in the revserse process.
Number of patches per side
Denoising schedule
X 2
0,
3
[0, 0,,35,00
4 X 4
0, 4, 6, 10
5
0,5, 7,
6 x 6
0,6 6, 8, 12
Sort 4-Digit MNIST Numbers. For the task of sorting 4-digit MNIST numbers with n S 100, use the exact training and beam search setup as the Jigsaw Puzzle, except that we do not need to use REINFORCE. The neural network architecture is given in Table TO The denoising schedules, with riflle shuffles as the forward process and GPL as the reverse process, are give in Table II
For 200, we use the cosine decay learning rate schedule with 2350 steps of linear warmup and maximum learning rate 5e-5. The neural network architecture is the same as that of n < 100, with the exception that we use dmodel djeed-forward 768, 12 layers, and 12 heads for the transformer encoder layer. We use the PL distribution for the reverse process. When computing the loss, we randomly sample timestep from the denoising schedule in (q. in Appendix C.2 due to efficiency reasons. All other setups are the same as that of n. < 100.
TSP. For solving the TSP, we perform supervised learning train our SymmetricDifftusers to solve the TSP. Let a² be an optimal permutation, and let X, the list of nodes ordered by a². We note that any cyclic shift of Xo is also optimal. Thus, for simplicity and without loss of generality, always assume o² (1) = 1. the forward process of SymmetricDiftusers, we only shuffle the second to the n node (or component). In the reverse process, mask certain parameters of the reverse distribution so that we will always sample a permutation with (1) =
The architecture details slightly different for TSP-20, where we input both node and edge features into our network. Denote by X, the ordered list of nodes at time t. We obtain Y, E Ruxdu as in Eq. 15 where encodera rg now sinusoidal embedding of the 2D coordinates. Let D, Raxn be the
23

---

Published as a conference paper at ICLR 2025
Table 10: Sort 4-digit MNIST numbers neural network architecture and hyperparameters.
Layer
Details
Output channels 32, kernel size 5, padding 2, stride 1
Convolution
Batch Normalization
ReLU
Max-pooling
Pooling 2
Output channels 64, kernel size 5, padding 2, stride 1
Convolution
Batch Normalization
ReL.U
Max-pooling
Pooling 2 Output dimension (dim aiim_after_con 128)/2
Fully-connected
ReLU
Fully-connected
Output dimension 128
layers, 8 heads, model dimension (dmosel) 128,
Transformer encoder
feed-forward dimension 512, dropout 0.
Table 11: Denoising schedules for the Sort 4-Digit MNIST Numbers task, where riffle shuffle the forward process and GPL. the revserse process.
Sequence Length
Denoising schedule
[0,
5
[0,
[0,
9
5,
[0, 4, 7, 10]
32
0, 5, 12]
[0, 5, 6, 7,
100
0, 8, 10, 15
200
00,9, 10, 12
matrix representing the pairwise distances of points in X,, respecting the order in X,. Let E, E R be the flattened vector of the upper triangular part of Dr. We also apply sinusoidal embedding to E, and add dime__mbb((
it. We call the result F, R[))xdoats
Now, instead of applying the usual transformer encoder with self-attentions,, we alternate between cross-attentions and self-attentions. For cross-attention layers, the previous layer as the query, and we always use = V = Fz. We also apply an attention mask the cross-attention, that each node will only attend edges that it is incident with. For self- attention layers, we always use the node representations from the previous layer as input. We always use an even number of layers, with the first layer being being a self-attention layer structured to produce the required parameters for the reverse distribution as illustrated in Appendix B For hyperparameters, we use 16 alternating layers, attention heads, dmodel = 256, feed-forwart hidden dimension 1024, and dropout rate 0.1.
the node representations from
cots-nattttion layer, and the last layer
For training details on the TSP-20 task. we use the AdamW optimizer (Loshchilov & Hutter 2019) with weight decay le-4, e = 1e-8, and B (0.9, 0.999). We use the cosine annealing learning rate We train for epochs with batch size of 512. When computing the loss 12). use Monte-Carlo estimation for the expectation and sample trajectory. We use a denoising schedule of [0, 4, 5, 7]. with riffle shuffles as the forward process and GPL as the beam search for decoding, and we beam size of 256 both when
scheduler starting from 2e-4 and ending
reverse process. Finally, decoding from GPL and decoding along the denoising schedule.
For the TSP-50 task, tion. We train for 250 epochs with a batch size of 256. We use denoising schedule [o, 5, 6, 7]. We beam size of 768 when decoding from GPI. and decoding along the denoising schedule.
the original architecture stated in Appendix B without the edge informa-
24

---

Published as a conference paper at ICLR 2025
When decoding, we pick the permutation that gives the least TSP tour length in the final step of the beam search. All other setups are identical with that of TSP-20.
G.6 BASELINES IMPLEMENTATION DETAILS
Combel-Sinkhorre Network. We have re-implemented the Gumbel-Sinkhorn Network (Mena at 2018) for application jigsaw puzzles, following the implementations provided in the official repository To ensure a fair comparison, we conducted a thorough grid search of the model's hyper-parameters. The parameters included in our search space as follows,
Table 12: Hyperparameter Search Space for the Gumbel- Network
Hyperparameter
Values
Learning Rate (Ir)
{10-3 10 10-1 10-6) (50) {64, 128}
Batch Size
Hidden Channels
Kernel Size
(3.5)
{0.2, 0.5, 1,2,5] {20)
Number of Sinkhorn Iterations (n sink iter)
Number of Samples
(10)
Diffsort & Error-free Diffsort We have implemented two differentiable sorting networks from the official repository specific to error-free diffsort. For sorting 4-digit MNIST images, error-free diffsort employs Transformerl its backbone, with detailed hyperparameters listed in Table [3 Conversely, Diffsort uses CNN as its backbone, with a learning rate set to 10-3.5, the relevant hyperparameters are outlined in Table 14
For jigsaw puzzle tasks, error-free diffssott continues to utilize a transformer, whereas Diffsort employs CNN. For other configurations, we align the settings with those of tasks having similar sequence lengths in the 4-digit MNIST sorting task. For instance, for 3 x 3 puzzles, we apply the same configuration as used for sorting tasks with sequence length of 9.
Table 13: Hyperparameters for Error-Free Diffsort on Sorting 4-Digit MNIST Numbers
Sequence Length
Steepness
Sorting Network
Loss Weight
Learning Rate
3
10
odd even
1.00
10-1
5
26
odd even
1.00
10
7
31
odd even
1.00
10-4
9
34
odd even
1.00
10-4
15
25
odd even
0.10
10-1
32
124
odd even
0.10
10-1
52
130
bitonic
0.10
10-3.5
100
140
bitonic
0.10
10-3.5
200
200
bitonic
0.10
10-4
TSP. For the baselines for TSP, we first have traditional operations research solvers. Gurobi Gurobi Optimization, LLC, 2023) and Concorde (Applegate al.] (0066) are known exact solvers, while LKH-3 (Helsgaun 2017) is a strong heuristic and 2-Opt (Lin & Kernighan, 1973) is a weak heuristic. For LKH-3, we used 500 trials, and for 2-Opt, we used 5 random initial guesses with seed 42.
For the GCN model (Joshi et al. 2019) we utilized the official repository and adhered closely to its default configuration for the TSP-20 20 dataset. For DIFUSCO (Sun & Yang 2023), we sourced it from
fffreeeenttn ionnntaan iabl e- swap function ions
25

---

Published as a conference paper at ICLR 2025
Table 14: Hyperparameters for Diffsort on Sorting 4-Digit MNIST Numbers
Sequence Length
Steepness
Sorting Network
6
odd even
20
odd even
7
29
odd even
32
odd even
15
25
odd even
bitonic
52
25
bitonic
100
25
bitonic
200
200
bitonic
its official repository and followed the recommended configuration of TSP-50 dataset, with minor adjustment in the batch size. We increased the batch size 512 accelerate the training process. For fair comparison, we also remove the post-processing heuristics in both models during the evaluation.
H LIMITATIONS
Despite the success of this method on various tasks, the model presented in this paper still requires time-space complexity of Onn²) due its reliance on the parametric representation of GPL and the backbone of transformer attention layers. This complexity poses significant challenge in scaling up to applications involving larger symmetric groups Lie groups.
26


Reviews:
[{"review_id": "N7oxQ3jTCS", "reviewer": "Reviewer_H6zv", "summary": "The paper presents SymmetricDiffusers, a novel discrete diffusion model for learning distributions over permutations within symmetric groups. This model addresses the complexity of directly modeling the vast and discrete state space of permutations by decomposing the task into simpler, more manageable transitions.\n\nKey contributions include:\n1) Forward Diffusion Process Using Card Shuffling Methods: Symmetric diffusers introduce noise to permutations using classical shuffling methods (riffle shuffles, random transpositions, and random insertions). These methods facilitate a gradual transformation toward a known noise distribution, simplifying the learning process.\n2) Generalized Plackett-Luce (PL) Distribution for Reverse Diffusion: To return the noisy state to its original distribution, the model leverages a neural network-based generalized PL distribution, enhancing expressiveness and effectively reconstructing complex dependencies within permutations.\n3) Theoretically Grounded Denoising Schedule: An optimized denoising schedule merges reverse steps to boost sampling efficiency and learning performance, reducing computational requirements without sacrificing accuracy.\n\nThe model demonstrates state-of-the-art or comparable results in tasks such as sorting, jigsaw puzzle assembly, and solving traveling salesman problems, validating its effectiveness in permutation-based applications.", "questions": "Comparative Analysis with Other Discrete Diffusion Models:\nQuestion: Could you provide additional comparisons between symmetric diffusers and other discrete diffusion models, such as discrete denoising diffusion probabilistic models (D3PMs), particularly regarding tasks that involve smaller permutation spaces?\nSuggestion: Including benchmarks with other models for more straightforward tasks could help clarify SymmetricDiffusers' unique benefits in terms of accuracy, expressiveness, and computational efficiency. If it outperforms or scales better than current discrete diffusion models, it would provide valuable evidence of its novelty.\n\nScalability for Larger :\nQuestion: How do symmetric diffusers handle scalability as increases, and what strategies do you envision for managing the factorial growth in the state space for large permutations?\nSuggestion: An empirical analysis or theoretical discussion on the scalability limits and potential optimizations (e.g., sparse transition matrices or modular architectures) would provide greater insight into the model’s practicality for large-scale tasks.\n\nDenoising Schedule Efficiency:\nQuestion: How significant is the denoising schedule's impact on computational efficiency and model performance? Have you considered alternative denoising schedules that might further improve efficiency?\nSuggestion: An ablation study on the denoising schedule’s impact on performance and efficiency would be beneficial. Exploring adaptive denoising strategies or alternative merging techniques might also enhance the model’s flexibility and efficiency.\n\nForward and Reverse Diffusion Process Clarifications:\nQuestion: Could you clarify certain technical details of the forward and reverse diffusion processes, especially for readers unfamiliar with symmetric groups? Specifically, could you break down the mathematical formulations for each shuffling method?\nSuggestion: To make these sections more accessible to readers, illustrative examples for the shuffling and reverse diffusion steps or a more detailed explanation of terms such as “stationary distribution” and “mixing time” could be added.\n\nBroader Applicability and Adaptation to Other Discrete Structures:\nQuestion: Can symmetric diffusers be adapted for other structured discrete data beyond permutations, such as specific types of graph structures or other combinatorial tasks?\nSuggestion: A brief discussion of the model’s applicability to other structured discrete data or potential adaptations would help clarify its versatility. For instance, could SymmetricDiffusers’ approach be generalized to data with hierarchical or relational structures?\n\nGeneralized Plackett-Luce Distribution Benefits:\nQuestion: Could you provide more insights or experiments to demonstrate the practical benefits of using the generalized Plackett-Luce (PL) distribution over the standard PL model? In which cases does the generalized PL significantly enhance performance?\nSuggestion: Examples or comparisons that highlight the generalized PL distribution’s expressiveness—especially in complex permutation tasks—would illustrate its value over standard PL models.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: This paper demonstrates notable originality by advancing discrete diffusion models to tackle the problem of learning distributions over permutations in symmetric groups. The task of modeling permutations is inherently challenging due to the factorial growth of the state space, and SymmetricDiffusers introduces an innovative solution by utilizing card shuffling techniques (riffle shuffle, random transpositions, and random insertions) as part of a structured forward diffusion process. Combining classic combinatorial methods with modern neural-based diffusion modeling is an inspired choice well-suited to the discrete and combinatorial nature of permutations. Compared to related works on discrete diffusion—such as Discrete Denoising Diffusion Probabilistic Models (D3PMs), which focus on multinomial categories—SymmetricDiffusers addresses a unique domain with a permutation-focused framework that current D3PMs do not target.\n\nFurthermore, introducing a generalized Plackett-Luce (PL) distribution for the reverse process sets this work apart from other discrete diffusion models, such as score-based continuous-time discrete-state models, which operate on categorical data but lack this flexibility. The generalized PL distribution is well suited for structured dependencies in permutations, enabling greater expressiveness and more accurate learning over complex permutations.\n\nQuality: The paper demonstrates high methodological rigor. Each component of SymmetricDiffusers—the forward process using card shuffling, the generalized PL distribution for reverse diffusion, and the denoising schedule—is technically well developed and grounded in established theories of random walks and Markov chains on finite groups. This adds a layer of mathematical credibility to the proposed model, particularly in the careful treatment of transition probabilities and mixing times associated with the shuffling methods.\n\nAdditionally, the experiments are comprehensive and well-chosen to validate the model across diverse tasks, such as sorting, jigsaw puzzle completion, and traveling salesman problems (TSPs). This diversity of functions effectively showcases the robustness and generalizability of SymmetricDiffusers, achieving state-of-the-art or comparable results in each case. Unlike other discrete diffusion models focused on simpler domains (e.g., categorical image data), SymmetricDiffusers’ application to more complex combinatorial problems uniquely highlights the model’s practical value and robustness.\n\nClarity: The paper is generally well-organized and easy to follow, with a clear introduction to the challenges and requirements of permutation modeling. Complex ideas are supported by notations, figures, and a well-structured flow that gradually builds the reader’s understanding from the background on symmetric groups to the technical construction of SymmetricDiffusers. Visualizing the forward and reverse processes in figure form is particularly helpful, illustrating the model’s approach to diffusion over permutations.\nHowever, compared to some other papers on discrete diffusion (e.g., D3PMs), certain sections—particularly those explaining the forward and reverse processes in detailed mathematical terms—might benefit from additional simplification for readers less familiar with random walks on finite groups. Nevertheless, the overall clarity and structure make the contributions accessible and understandable.\n\nSignificance: SymmetricDiffusers addresses an important area in machine learning by advancing the state of permutation modeling within discrete diffusion frameworks. Distributions over permutations are critical in domains like ranking, combinatorial optimization, and sequence alignment, and current models struggle with the computational complexity posed by large discrete spaces. This work opens up new possibilities for generative modeling within these areas, especially with its robust performance in tasks requiring complex combinatorial reasoning (e.g., TSP).\n\nCompared to recent discrete diffusion works on more general categorical or sequential data, this paper contributes uniquely to discrete generative modeling by directly addressing permutation structures. This is significant because it establishes a pathway for diffusion models to utilize high-complexity tasks beyond traditional applications effectively. By providing an effective means of modeling permutations, SymmetricDiffusers could substantially impact practical applications and inspire future research in discrete-state diffusion for structured data.\n\nThe paper’s strengths lie in its original approach, which combines structured combinatorial methods with diffusion modeling. It also has high-quality methodological rigor, clearly presents complex concepts, and significantly contributes to modeling distributions over large permutation spaces, positioning it as a valuable advancement in discrete diffusion research.", "weaknesses": "Limited Comparative Analysis with Other Discrete Diffusion Models: While SymmetricDiffusers clearly advances permutation modeling, the paper could benefit from a more detailed comparison with other discrete diffusion models, particularly those handling high-dimensional categorical or sequential data, such as Discrete Denoising Diffusion Probabilistic Models (D3PMs). The current related work section mentions other models briefly but does not delve into how SymmetricDiffusers directly compares in terms of handling large discrete spaces. Including more detailed experiments or ablation studies to demonstrate SymmetricDiffusers’ advantages over such models regarding efficiency or performance on simpler permutation tasks could clarify its relative strengths and limitations.\n\nLack of Exploration on Scalability to Larger: The factorial growth of permutations means that scaling to larger sizes can become computationally intensive. The paper currently demonstrates its model on small-to-moderate values (e.g., tasks like 4-digit MNIST sorting). However, it does not provide clear insights into the scalability limits or potential strategies for scaling SymmetricDiffusers to larger values, where the computational load may become a bottleneck. Future work should consider including benchmarks on larger values or discussing ways to optimize the model’s performance, such as using sparse transition matrices or adopting modular architectures.\n\nClarity in the Forward and Reverse Processes: While the paper is generally well-structured, some sections—particularly those detailing the forward and reverse diffusion processes—are highly technical and may be challenging for readers less familiar with symmetric groups and permutation modeling. Additional clarifications or simplified explanations could improve accessibility. Specifically, breaking down the mathematical formulations for each shuffling method and reverse diffusion process with more illustrative examples would make these sections easier to understand. Clearer definitions of key terms, such as “stationary distribution” and “mixing time” in the context of random walks, could also make the content more accessible to a broader audience.\n\nEfficiency of the Denoising Schedule: The paper introduces a theoretically grounded denoising schedule to merge reverse steps and improve efficiency, but it lacks concrete benchmarks or ablation studies to assess its impact. Comparing SymmetricDiffusers with and without the denoising schedule in terms of computational time and performance would provide readers with a clearer understanding of its practical benefits. Additionally, exploring alternative denoising schedules or adaptive strategies that adjust based on task complexity could further optimize the model’s performance.\n\nBroader Applicability and Practical Implications: Although SymmetricDiffusers demonstrates promising results in permutation-specific tasks (sorting, TSP, jigsaw puzzles), the paper could better communicate its broader applicability and potential limitations. For example, could the model be applied to non-permutation-based tasks with discrete structures, such as certain types of graph-based tasks? A brief discussion of the boundaries of SymmetricDiffusers’ applicability and how it might adapt to related yet distinct discrete structures would clarify its versatility and limitations.\n\nSuggestions for Improvement:\n1) Enhance comparative experiments by including Symmetric Diffusers alongside other discrete diffusion models (e.g., D3PMs) on more straightforward permutation tasks for more explicit benchmarks.\n2) Expand on the model’s scalability for larger spaces, possibly with benchmarks on larger permutation spaces or theoretical discussions on extending to high-dimensional settings.\n3) Provide additional clarifications and illustrative examples for complex forward and reverse diffusion processes sections.\nInclude an ablation study on the denoising schedule to quantify its impact on performance and efficiency.\n4) Discuss the model's broader applicability to other discrete structures beyond permutations, providing readers with insight into its potential versatility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4rK748I6hr", "reviewer": "Reviewer_7CmY", "summary": "This paper presents a discrete diffusion model for learning a distribution over permutations ($S_n$). They present several choices of forward noising process, including transpositions, insertions, and the rifle shuffle. For the reverse noising process, they define inverse processes corresponding to each of the three proposed forward processes, as well as a generalization of the Plackett-Luce Distribution that is more expressive than the original (e.g. it can represent a delta distribution). They train their diffusion models via a variational lower bound, estimated via Monte Carlo samples since they cannot obtain an analytic form. They derive a noise schedule, based on merging adjacent diffusion steps for certain inverse distributions, and run experiments comparing different versions of their diffusion model to differentiable sorting methods and test tasks including sorting MNIST digits and solving traveling salesman instances.", "questions": "1. As discussed in the Weaknesses section, what exactly is the motivation of using diffusion model for these problems, if ultimately only a single learned permutation is required per input? This is my primary question. \n2. For the right choice of parameters, can the reverse processes actually represent the exactly correct distributions induced by the corresponding forward diffusion process?\n3. One minor point of confusion was that the abstract claims to learn a distribution over $S_n$, but the concrete objects that are dealt with are ordered sets of objects (stored in an $n$ by $d$ matrix). Would it be accurate to refer to this method as *conditional* diffusion? If not, how could the architectures best be modified to output a distribution over raw permutation matrices? \n4. As noted on line 154, $\\mathcal{S}$ does not change across steps — why enforce this for diffusion models? Does this make something easier? Is it potentially restrictive in terms of what distributions can be represented after a given number of steps?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is the first work to define a diffusion model over $S_n$. Their enumeration, and mathematical treatment, of possible forward and reverse noising processes is thorough, as are the experiments and ablations. The presentation is quite clear in general. The application of a diffusion model to problems that require outputting a permutation in a differentiable matter — even if a distribution over $S_n$ is not strictly required — seems to be creative and novel. The experimental results are generally good, particularly at increased sequence lengths.", "weaknesses": "To me, the main piece missing from this paper is a motivation of why machine learning applications require generative models over permutations in the first place, rather than just outputting a single permutation, which is what the tested datasets seem to satisfy. Is this related to differentiability? The introduction is well-written, but would improve from being more concrete and grounded in the machine learning literature. To be specific, the claim on line 35 that “Therefore, studying probabilistic models over $S_n$ through the lens of modern machine learning is both natural and beneficial” feels a bit unjustified. Differentiable sorting is mentioned in the related works, but a discussion of why “such methods often focus on finding the optimal permutation instead of learning a distribution over the finite symmetric group” and what the tradeoffs are would be helpful. \n\nThe experiments also show stronger performance for longer sequence lengths, but the quadratic scaling with longer sequence lengths remains an open direction of improvement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L36fbgA7b2", "reviewer": "Reviewer_nrB8", "summary": "This paper explores a novel problem: learning probability distributions over finite symmetric groups. As an example, the problem can be thought of as learning to sort, with important differences:\n\n- instead of finding just the most optimal permutation, the focus is on learning a *distribution* over all possible rankings, so that rankings \"closer\" to the optimal permutation are more likely to be generated.\n- the problem formulation is general enough to cover rankings over any finite group of elements and not just a set of $n$ integers as long as one trains the algorithm using appropriate data (for instance, pictures of numbers instead of just numbers). \n\nThe authors propose *SymmetricDiffusers*, a discrete diffusion model (with a transformer-based architecture in this case) trained to recover the target permutation in several steps after the original (target permutation) is converted into noise over a number of steps. This decomposition of the recovery process eases the otherwise very difficult problem of learning to directly come up with the optimal permutation out of $n!$ possible states.\n\nWith this context, the authors' study reveals several insights into potentially best practices with regard to the training process, *e.g.,* choice of the forward shuffling method (riffle shuffle because of its fast mixing time), and when it might be feasible to merge steps during the denoising process. The authors validate their findings via state-of-the-art results on three benchmarks: MNIST sorting, jigsaw puzzles, and the Travelling Salesman Problem.", "questions": "1. Have the authors considered evaluating OOD performance (e.g., feeding colored, or otherwise font-shifted MNIST into a model that was trained on grayscale images) ? Do they anticipate a drop in performance in that setting?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The style of writing is extremely clear and structured. \n2. The paper makes a wealth of interesting contributions: introducing a novel perspective on an important problem that is currently underexplored in the ML community; proposing a method to solve said problem; revealing insights that will help further research; and finally, validating their method on three benchmarks. The authors also share their code. The general nature of the problem means that the findings here can potentially open up a whole new set of possibilities.", "weaknesses": "I don't see obvious weaknesses. But, as the authors also acknowledge (in the conclusion and Appendix G), there is potential for improving scalability (w.r.t. $n$) and possibly extending the method to finite groups beyond $S_n$. \n\n**(Nit: typos)**\n- Abstract (line 009): groups --> group?\n- line 472: performances --> performance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0nNEDkjmSt", "reviewer": "Reviewer_iCzn", "summary": "The paper studies the important question of sampling from the finite symmetric group $S_n$ via diffusion models. This is structurally different from prior works on diffusion models which concentrate on sampling from product spaces (i.e, sampling without replacement). The authors clearly introduce various noising Markov chains for sampling a uniform distribution over $S_n$ and show how to parametrize the forward and reverse processes. This is the main technical contribution of the paper. The training loss is then based the variational lower bound of D3PM (Austin et al). The paper achieves strong performance in solving jigsaw puzzles of MNIST and CIFAR 10, sorting numbers from MNIST and solving TSP problems.", "questions": "Address the points raised in the weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and the Markov chains used to obtain uniformly random permutations are clearly surveyed. The construction of the parameterization of the reverse process and the loss function is nice and the algorithm outperforms existing methods in the empirical tasks.", "weaknesses": "1. Experiments are very small scale, comprising of sorting 4 digit MNIST, solving 20 node TSPs and solving jigsaw puzzles of CIFAR-10 data. \n\n2. There is no substantive theoretical contribution other than introducing the parametrization for the reverse processes. \n\n3. The reverse process for random transposition is not very expressive. Suppose the reverse transposition is $(1,2)$ with probability $0.5$ and $(2,3)$ with probability $0.5$. This simple distribution cannot be expressed using the model. \n\n4. Note that $S_n \\subseteq [n]^n$. If the input data comprises of only permutations, then the network should learn to sample from a distribution whose samples are permutations and the standard framework of diffusion models applies. This simple baseline has not been considered. \n\n[Q related to 4] The authors also mention that representing a transition matrix over $S_n$ requires $n!\\times n!$ sized matrix. However, the authors themselves give a succinct description/ representation of the forward transition matrix in the paper. The authors should elaborate why it is not possible to use this representation algorithmically.\n\n**Minor:**\nIn proposition 1, should it be changed to \"the GPL distribution can represent a delta distributions in the limit\" instead of \"exactly\"?\n\n*additional references:*\n[1] Generating a random permutation with random transpositions by Diaconis and Shahshahani\n[2] Simplified and Generalized Masked Diffusion for Discrete Data by Shi et al.\n[3] Glauber Generative Model: Discrete Diffusion Models via Binary Classification by Varma et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oCl7twbKrw", "reviewer": "Reviewer_hmq7", "summary": "This authors proposes a diffusion framework that learns a distribution over permutations.\nThey propose a set of different reverse and forward diffusion processes to realize this goal.\nThey validate their approach on various different experiments such as TSP, four-digit MNIST, CIFAR and noisy MNIST.", "questions": "1. How do the proposed reverse and forward processes compare to naive discrete denoising diffusion models?\nWhat is the difference between your method and other differentiable sorting baselines in Tab. 1?\n\n2. What is the sequence length $n$ in the four-digit MNIST dataset?\n\n3. Why are experiments limited to TSP-20 instances rather than including larger TSP instances? TSP-20 is considered a very small problem size. This experiment is the most interesting to me and it would be interesting to see how this method performs on larger TSP instances. Can you provide an additional comparison on TSP-100?\n\n4. Can your method be combined with the framework proposed in [1] to solve TSP using diffusion models without requiring data from classical solvers?\n\n\n\n## References\n[1] Sanokowski, Sebastian, Sepp Hochreiter, and Sebastian Lehner. \"A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization.\" Forty-First International Conference on Machine Learning.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents an interesting approach with mostly clear writing\n2. Several innovative reverse and forward processes are proposed", "weaknesses": "1. The claim of state-of-the-art performance on TSP-20 appears too bold given the limited scope\n2. Several experimental details lack sufficient clarity (see Questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)