Title: SymmetricDiffusers: Learning Discrete Diffusion on Finite Symmetric Groups
Venue: NeurIPS 2024
Accepted: no
Keywords: Finite Symmetric Groups

Discrete Diffusion

Permutation

Riffle Shuffles

the Plackett-Luce Distribution

Sorting

Jigsaw Puzzle
Abstract: Finite symmetric groups $S_n$ 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 performances on solving tasks including sorting 4-digit MNIST images, jigsaw puzzles, and traveling salesman problems.

1

Abstract

Finite symmetric groups S n 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 performances on solving tasks including sorting 4-digit MNIST images, jigsaw puzzles, and traveling salesman problems.

1 Introduction

As a vital area of abstract algebra, finite groups provide a structured framework for analyzing symmetries and transformations which are fundamental to a wide range of fields, including combinatorics, physics, chemistry, and computer science. One of the most important finite groups is the finite symmetric group S n , defined as the group whose elements are all the bijections (or permutations) from a set of n elements to itself, with the group operation being function composition.

Classic probabilistic models for finite symmetric groups S n , such as the Plackett-Luce (PL) model [35, 27], the Mallows model [28], and card shuffling methods [9], are crucial in analyzing preference data and understanding the convergence of random walks. Therefore, studying probabilistic models over S n through the lens of modern machine learning is both natural and beneficial. This problem is theoretically intriguing as it bridges abstract algebra and machine learning. For instance, Cayley's Theorem, a fundamental result in abstract algebra, states that every group is isomorphic to a subgroup of a symmetric group. This implies that learning a probability distribution over finite symmetric groups could, in principle, yield a distribution over any finite group. Moreover, exploring this problem could lead to the development of advanced models capable of addressing tasks such as permutations in ranking problems, sequence alignment in bioinformatics, and sorting.

However, learning a probability distribution over finite symmetric groups S n poses significant 31 challenges. First, the number of permutations of n objects grows factorially with n , making the 32 inference and learning computationally expensive for large n . Second, the discrete nature of the data 33 brings difficulties in designing expressive parameterizations and impedes the gradient-based learning. 34

In this work, we propose a novel discrete (state space) diffusion model over finite symmetric groups, 35 dubbed as SymmetricDiffusers . It overcomes the above challenges by decomposing the difficult 36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

problem of learning a complicated distribution over S n into a sequence of simpler problems, i.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 as the most effective. We also provide empirical guidelines on choosing the diffusion length based on the mixing time of the riffle shuffle. Furthermore, we examine potential transitions for the reverse diffusion, such as inverse shuffling methods and the PL distribution, and introduce a novel generalized PL distribution. We prove that our generalized PL is more expressive than the PL distribution. Additionally, we propose a 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-art 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 [37, 11, 4, 38]. Techniques from a variety of different fields, including probability, combinatorics, and representation theory, have been used to study random walks on finite groups [38]. In particular, random walks on finite symmetric groups are first studied in the application of card shuffling, with many profound theoretical results of shuffling established. A famous result in the field shows that 7 riffle shuffles are enough to mix up a deck of 52 cards [4], where a riffle shuffle is a mathematically precise model that simulates how people shuffle cards in real life. The idea of shuffling to mix up a 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 [40, 41, 16, 42] are a 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 a continuous-time continuousstate 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 [3], and graphs [45]. Existing discrete diffusion models are applicable for learning distributions of permutations. However, they focused on cases where the state space is small or has a special ( e.g ., decomposable) structure and are unable to deal with intractable-sized state spaces like the symmetric group. In particular, [3] requires an explicit transition matrix, which has size n ! × n ! in the case of finite symmetric groups and has no simple representations or sparsifications.

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 [13, 8, 5], or uses differentiable swap functions [33, 34, 20]. The Gumbel-Sinkhorn method [29] has also been proposed to learn 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 to be less effective as n grows larger due to their high complexities.

3 Learning Diffusion Models on Finite Symmetric Groups

We first introduce some notations. Fix n ∈ N . Let [ n ] denote the set { 1 , 2 , . . . , n } . A permutation σ on [ n ] is a function from [ n ] to [ n ] , and we usually write σ as ( 1 2 · · · n σ (1) σ (2) · · · σ ( n ) ) . The identity permutation , denoted by Id , is the permutation given by Id( i ) = i for all i ∈ [ n ] . Let S n be the set of all permutations (or bijections) from a set of n elements to itself, called the finite symmetric group , whose group operation is the function composition. For a permutation σ ∈ S n , the permutation matrix Q σ ∈ R n × n associated with σ satisfies e ⊤ i Q σ = e ⊤ σ ( i ) for all i ∈ [ n ] . In this paper, we consider a set of n distinctive objects X = { x 1 , . . . , x n } , where the i -th object is represented by a d -dimensional vector x i . Therefore, a ranked list of objects can be represented as a matrix X = [ x 1 , . . . , x n ] ⊤ ∈ R n × d , where the ordering of rows corresponds to the ordering of objects. We can permute X via permutation σ to obtain Q σ X .

<latexi sh

1_b

64="A

T

BO

/

C

vKI

YG

DH

oM

>

Q

c

V

LS

N

FJ3

rPW

zf

U

qy

X

g

2

Z

w

+

9

j

5d

0

np

u

R

E

m

8

k

7

<latexi sh

1_b

64="

qk

TyImgA

rWLH

p

o

OD

Rj

E

>

CJ3

c

V

NS

M

G

/

9

B

5

d

0Yv

u

2

Z

n

+

8

7

Q

w

P

K

X

f

F

<latexi sh

1_b

64="

Dk

TBpX

v

df9N

F

YA2Ro

>

CJ3

c

V

SgM

GEyq r/Wj

m7I

0

H

u

Z

n

+

U

8

K

L

Q

O

z

5

P

w

<latexi sh

1_b

64="

IL

m

7gFOqP+

YSpMU3J

Q

>A

C

c

VB

8

wHE2n0z

/

Xr

y

v

9

uD

K

k

f

T

j

5

o

ZR

W

G

N

d

<latexi sh

1_b

64="oPdGq

Sj5

L

+VWK

7E

>A

CIX

c

BN

0J

F

2

M

r

R

Y

w

/

9y

n

3

T

f

u

g

Z

U

z

m

8

Q

p

H

<latexi sh

1_b

64="9p

0wJ5

ZBf

H

G

RX

PY

>A

CLn

c

VDNS

Uq7

+

r

E

y

m

OK

Qv

g

W

z3

d

M

2

/

F

u

I

o

8

j

k

T

<latexi sh

1_b

64="FImKq8JB

gz

9N0

o+Ck nrP

>A

3

c

V

SwM

E

2

/Wj

R

5

d

X

u

G

p

yY

Z

Tf

H

D

7

U

L

v

O

Q

<latexi sh

1_b

64="RZ9H3Vpqr7QOv5

k

EI

PzF

Yfg

>A

C

n

c

BdSwJ

2

L

M

To

W

U

u

X

0K

/

D

G

mN

+

j

8

y

<latexi sh

1_b

64="Ig

TOu

JS5v nZ

U+q

d

m

M

>A

C

X

c

VDL

F

2

K

R

Y9

0

H

Q

p

7

j

B

/

o

f8

EyP

3

z

N

W

G

k

r

w

<latexi sh

1_b

64="O+mYu2

0Bo

T

grVdz

Zf8

>A

CPn

c

DNS

G

y

/Wj

R

Q

J

9

K

7

p

E3

q

F

wM

I

v

5

X

k

U

L

<latexi sh

1_b

64="HM

0/XvgJOz

p

cDI8W

ZP

kQF

>A

C

n

V

NS

GE

U

q

rR

9B

u

Y

K

2

7

y

T

L

fw

5

m

d3

+

o

j

X

⋯

⋯

⋯

Ascending

x

1

x

2

x

3

x

1

x

2

x

3

x

1

x

2

x

3

Probability

Forward Step

Reverse Step

Merged Reverse Step

123

[

x

1

,

x

2

,

x

3

]

/latticetop

{

x

1

,

x

2

,

x

3

}

132

312

321

123

213

231

q

(

X

T

)

132

312

321

123

213

231

q

(

X

1

)

132

312

321

123

213

231

p

data

(

X

0

)

σ

1

σ

t

σ

T

σ

′

T

σ

′

t

σ

′

1

X

0

X

1

X

t

X

T

p

θ

(

σ

′

T

|

X

T

)

p

θ

(

σ

′

t

|

X

t

)

p

θ

(

σ

′

1

|

X

1

)

q

(

X

1

|

X

0

, σ

1

)

q

(

X

t

|

X

t

-

1

, σ

t

)

q

(

X

T

|

X

T

-

1

, σ

T

)

p

(

X

0

|

X

1

, σ

′

1

)

p

(

X

t

-

1

|

X

t

, σ

′

t

)

p

(

X

t

|

X

T

, σ

′

T

)

Our goal is to learn a distribution over S n . We propose learning discrete (state space) diffusion 89 models, which consist of a forward process and a reverse process . In the forward process, starting 90 from the unknown data distribution, we simulate a random walk until it reaches a known stationary 91 'noise' distribution. In the reverse process, starting from the known noise distribution, we simulate 92 another random walk, where the transition probability is computed using a neural network, until it 93 recovers the data distribution. Learning a transition distribution over S n is often more manageable 94 than learning the original distribution because: (1) the support size (the number of states that can be 95 reached in one transition) could be much smaller than n ! , and (2) the distance between the initial and 96 target distributions is smaller. By doing so, we break down the hard problem (learning the original 97 distribution) into a sequence of simpler subproblems (learning the transition distribution). The overall 98 framework is illustrated in Fig. 1. In the following, we will introduce the forward card shuffling 99 process in Section 3.1, the reverse process in Section 3.2, the network architecture and training in 100 Section 3.3, denoising schedule in Section 3.4, and reverse decoding methods in Section 3.5. 101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

3.1 Forward Diffusion Process: Card Shuffling

Suppose we observe a set of objects X and their ranked list X 0 . They are assumed to be generated from an unknown data distribution in an IID manner, i.e ., X 0 , X iid ∼ p data ( X, X ) . One can construct a bijection between a ranked list of n objects and an ordered deck of n 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 ⊆ S n be a set of permutations that are realizable by a given shuffling method in one 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,

where q ( X t | X t -1 ) = ∑ σ t ∈S q ( X t | X t -1 , σ t ) q ( σ t ) and the first equality in Eq. (1) holds since X 0 implies X . In the forward process, although the set X does not change, the rank list of objects X t changes. Here q ( σ t ) has the support S and describes the permutation generated by the underlying shuffling method. Note that common shuffling methods are time-homogeneous Markov chains, i.e ., q ( σ t ) stays the same across time. q ( X t | X t -1 , σ t ) is a delta distribution δ ( X t = Q σ t X t -1 ) since the permuted objects X t are uniquely determined given the permutation σ t and X t -1 . We denote the neighbouring states of X via one-step shuffling as N S ( X ) := { Q σ X | σ ∈ S} . Therefore, we have,

<latexi sh

1_b

64="9XkRTfgnQK0j

CJ

58

E

u

+

>A

S

c

VBN

2

mv oZ

IH

p

r

PF

w

W

d

/D

G

M

z

L

y

U

7

q

3

Y

O

<latexi sh

1_b

64="oXFT89DAY

Rr

C

j

0d wyc

>

W

VBNS

E

3

q/

nv

g

I

Q

GL

z

u5m

+

k

7

P

O

H

M

K

f

p

Z

U

2

J

<latexi sh

1_b

64="cvDyAjHJur

S

9Tm wpQ

BM

>

C

V

L

N

F

Uq/

d

k

WI

K

R

0

7

Y

n

ZG

Ef

3

o

O

z

+

g

P

2

5

X

8

<latexi sh

1_b

64="A8O

JGMrq

LI

vQ0TY2D

3

>

C

n

c

V

S

N

F

o

pf

7dB

g

k

U

W

/

j

Z

K

w

zE

P

5

y

X

9

m

u

H

R

+

<latexi sh

1_b

64="yjk+wQo3g

0

J

D

dM

>A

CHn

c

V

LS

NBE

Rm

8

Xr

2

P

T

Y

z

K

W

vO

f

F

U

9

q7

/

5

Z

p

G

I

u

<latexi sh

1_b

64="od

Jm+

z

U7

5M9KT

3OX

C

Gw

>A

I

c

VDLSg

FE

q2v

p

u

k

W

n

Z

H

Q

R

f8

0

/

B

Y

r

Py

j

N

<latexi sh

1_b

64="/LWOr

J

5N

S

+

8vn

XP

T

>A

C

c

VD

gM

F

Uq7

3

I

E

R

uK9

H

G

p

k

Y

d

Q

B2

z

m

0

Z

o

j

f

y

w

<latexi sh

1_b

64="qB

YvAf

Z

0HT5

w

I

>

CQ

c

VDLS

N

FJ

U

7W+

2

R

dO

ygn

8

k

m

z

M

/

j

rX

P

3

p

o

G

u

9

K

E

<latexi sh

1_b

64="mjBc

ISz+59

fdNTw0vRqn

y

O

>A

CQH

V

8

F

U

/W7

L

g

p

Po

M

32

E

Y

G

X

r

D

J

u

k

K

Z

<latexi sh

1_b

64="

qk

TyImgA

rWLH

p

o

OD

Rj

E

>

CJ3

c

V

NS

M

G

/

9

B

5

d

0Yv

u

2

Z

n

+

8

7

Q

w

P

K

X

f

F

<latexi sh

1_b

64="

Dk

TBpX

v

df9N

F

YA2Ro

>

CJ3

c

V

SgM

GEyq r/Wj

m7I

0

H

u

Z

n

+

U

8

K

L

Q

O

z

5

P

w

<latexi sh

1_b

64="

IL

m

7gFOqP+

YSpMU3J

Q

>A

C

c

VB

8

wHE2n0z

/

Xr

y

v

9

uD

K

k

f

T

j

5

o

ZR

W

G

N

d

<latexi sh

1_b

64="qZuY

o2

E

C

/5

w

Fgj7W

>A

IX

c

J

L

Mvr

Q

Nm

n

0y

R

H

p

9

f

T

k

D

G

U

z

KO

P

3

+

8

<latexi sh

1_b

64="

T58

u

HLDK0+

g

QA

o

>

C

n

c

V

NS

M

G

yq r/Wj

R

B

dk

f

7

p

w2

Y

F

9

PU

z

X

v

Z

I

J3

O

m

E

<latexi sh

1_b

64="XrT0MP

S

N

A

jq

OzvWm

YE

>

CHn

c

VDL

g

B

Jy

R

8

2

o

k

3Q+

K

f

F

U

9

d

7

/

5

Z

p

G

I

w

u

<latexi sh

1_b

64="+ET

U

K

>A

3

c

V

LS

F

q

v

p

u

k

w

W

G

o

n

Z

H

mz

Q

JR

f8

9

5

/

X

B

7

Y

r

P

N

<latexi sh

1_b

64="

+KARTmw

Y

N5

3n

o

>

CPH

c

VD

S8M

GE2

0z

/

X

I

Z

9

W

QL

q

pO

f

d

vu

F

U

7

g

y

k

B

<latexi sh

1_b

64="mNy

S

uDc

9

opPO8

>A

C

X

V

g

G

Wq7

+

Rv

UK

ZH

jFYw

r

5

Q

B

2J

0

d

E

I

/

f

<latexi sh

1_b

64="

Dk

TBpX

v

df9N

F

YA2Ro

>

CJ3

c

V

SgM

GEyq r/Wj

m7I

0

H

u

Z

n

+

U

8

K

L

Q

O

z

5

P

w

<latexi sh

1_b

64="

qk

TyImgA

rWLH

p

o

OD

Rj

E

>

CJ3

c

V

NS

M

G

/

9

B

5

d

0Yv

u

2

Z

n

+

8

7

Q

w

P

K

X

f

F

<latexi sh

1_b

64="

IL

m

7gFOqP+

YSpMU3J

Q

>A

C

c

VB

8

wHE2n0z

/

Xr

y

v

9

uD

K

k

f

T

j

5

o

ZR

W

G

N

d

<latexi sh

1_b

64="unwo/

SC

M5QT

YpHfK

0

>A

X

c

VBN

8

F

U

7

+

r

I

k

j

v

D2

d

J

y

G

O

9

P

W

m

E

g

L

3

Z

R

<latexi sh

1_b

64="7DARMEOgBITw5fyG

Uj cJ

Q

>

CHn

V

LS

N

m

8

Xr0

2

o

P

k

Y3

+

z

K

W

v

F

9

d

q

/

Z

p

u

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

Note that X t ∈ N S ( X t -1 ) is equivalent to σ t ∈ S and X t = Q σ t X t -1 .

3.1.1 Card Shuffling Methods

We now consider several popular shuffling methods as the forward transition, i.e ., random transpositions , random insertions , and riffle shuffles . Different shuffling methods provide different design choices of q ( σ t ) , 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, a transposition or a swap is a permutation σ ∈ S n such that there exist i = j ∈ [ n ] with σ ( i ) = j , σ ( j ) = i , and σ ( k ) = k for all k / ∈ { i, j } , in which case we denote σ = ( i j ) . We let S = { ( i j ) : i = j ∈ [ n ] } ∪ { Id } . For any time t , we define q ( σ t ) 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, q ( σ t = ( i j )) = 2 /n 2 when i = j and q ( σ t = Id) = 1 /n .

̸

Random Insertions. Another shuffling method is to insert the last piece to somewhere in the middle. Let insert i denote the permutation that inserts the last piece right before the i th piece, and let S := { insert i : i ∈ [ n ] } . Note that insert n = Id . Specifically, we have q ( σ t = insert i ) = 1 /n when i = n and q ( σ t = 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 [11], and independently by Reeds [37]. 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 k cards in the top pile is ( n k ) / 2 n for 0 ≤ k ≤ 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 S n obtained through one-step riffle shuffle. Let σ ∈ S n . A rising sequence of σ is a subsequence of σ constructed by finding a maximal subset of indices i 1 < i 2 < · · · < i j such that permuted values are contiguously increasing, i.e ., σ ( i 2 ) -σ ( i 1 ) = σ ( i 3 ) -σ ( i 2 ) = · · · = σ ( i j ) -σ ( i j -1 ) = 1 . For example, the permutation ( 1 2 3 4 5 1 4 2 5 3 ) has 2 rising sequences, i.e ., 123 (red) and 45 (blue). Note that a permutation has 1 rising sequence if and only if it is the identity permutation. Denoting by q RS ( σ ) the probability of obtaining σ through one-step riffle shuffle, it is shown in [4] that

where r is the number of rising sequences of σ . The support S is thus the set of all permutations with at most two rising sequences. We let the forward process be q ( σ t ) = q RS ( σ t ) 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 ( i.e ., the time until the Markov chain is close to its stationary distribution measured by some distance), and there exist quantitative results on their mixing times. Let q ∈ { q RT , q RI , q RS } , and for t ∈ N , let q ( t ) be the marginal distribution of the Markov chain after t shuffles. We describe the mixing time in terms of the total variation (TV) distance between two probability distributions, i.e ., D TV ( q ( t ) , u ) , where u is the uniform distribution.

For all three shuffling methods, there exists a cut-off phenomenon , where D TV ( q ( t ) , u ) stays around 1 for initial steps and then abruptly drops to values that are close to 0. The cut-off time is the time when the abrupt change happens. For the formal definition, we refer the readers to Definition 3.3 of

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

[38]. In [38], they also provided the cut-off time for random transposition, random insertion, and riffle shuffle, which are n 2 log n , n log n , and 3 2 log 2 n respectively. Observe that the riffle shuffle reaches the cut-off much faster than the other two methods, which means it has a 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 as 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 N T ( X ) := { Q σ X : σ ∈ T } . The conditional joint distribution is given by where p θ ( X t -1 | X t ) = ∑ σ ′ t ∈T p ( X t -1 | X t , σ ′ t ) p θ ( σ ′ t | X t ) . To sample from p ( X T |X ) , one simply samples a random permutation from the uniform distribution and then shuffle the objects accordingly to obtain X T . p ( X t -1 | X t , σ ′ t ) is again a delta distribution δ ( X t -1 = Q σ ′ t X t ) . We have

where X t -1 ∈ N T ( X t ) is equivalent to σ ′ t ∈ T and X t -1 = Q σ ′ t X t . In the following, we will introduce the specific design choices of the distribution p θ ( σ ′ t | X t ) .

3.2.1 Inverse Card Shuffling

A natural choice is to use 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 := { σ -1 : σ ∈ S} , from which we can parameterize p θ ( σ ′ t | X t ) .

̸

Inverse Transposition. Since the inverse of a transposition is also a transposition, we can let T := S = { ( i j ) : i = j ∈ [ n ] } ∪ { Id } . We define a distribution of inverse transposition (IT) over T using n +1 real-valued parameters s = ( s 1 , . . . , s n ) and τ such that

̸

where ψ ( s , π ) i = exp ( s π ( i ) ) / (∑ n k = i exp ( s π ( k ) )) and ϕ ( · ) is the sigmoid function. π ij is any permutation starting with i and j , i.e ., π ij (1) = i and π ij (2) = j . π ji is any permutation starting with j and i , i.e ., π ji (1) = j and π ji (2) = i .

Inverse Insertion. For the random insertion, the inverse operation is to insert some piece to the end. Let inverse_insert i denote the permutation that moves the i th component to the end, and let T := { inverse_insert i : i ∈ [ n ] } . We define a categorial distribution of inverse insertion (II) over T using parameters s = ( s 1 , . . . , s n ) such that,

Inverse Riffle Shuffle. In the riffle shuffle, the deck of card is first cut into two piles, and the two piles are riffled together. So to undo a riffle shuffle, we need to figure out which pile each card belongs to, i.e ., making a sequence of n binary decisions. We define the Inverse Riffle Shuffle (IRS) distribution using parameters s = ( s 1 , . . . , s n ) as follows. Starting from the last (the n th ) object, each object i has probability ϕ ( s i ) of being put on the top of the left pile. Otherwise, it falls on 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-Luce Distribution and Its Generalization

Other than specific inverse shuffling methods to parameterize the reverse process, we also consider general distributions p θ ( σ ′ t | X t ) whose support are the whole S n , i.e ., T = S n .

The PL Distribution. A popular distribution over S n is the Plackett-Luce (PL) distribution [35, 27], which is constructed from n real-valued scores s = ( s 1 , . . . , s n ) as follows,

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

for all σ ∈ S n . Intuitively, ( s 1 , . . . , s n ) represents the preference given to each index in [ n ] . To sample from PL s , we first sample σ (1) from Cat( n, softmax( s )) . Then we remove σ (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 σ (1) , . . . , σ ( n ) . By [7], the mode of the PL distribution is the permutation that sorts s in descending order.

The Generalized PL (GPL) Distribution. We also propose a generalization of the PL distribution, referred to as Generalized Plackett-Luce (GPL) Distribution . Unlike the PL distribution, which uses a set of n scores, the GPL distribution uses n 2 scores { s 1 , · · · , s n } , where each s i = { s i, 1 , . . . , s i,n } consists of n scores. The GPL distribution is constructed as follows,

Sampling of the GPL distribution begins with sampling σ (1) using n scores s 1 . For 2 ≤ i ≤ n , we remove i -1 scores from s i that correspond to σ (1) , . . . , σ ( i -1) and sample σ ( i ) from a categorical distribution constructed from the remaining n -i +1 scores in s i . It is important to note that the family of PL distributions is a strict subset of the GPL family. Since the GPL distribution has more parameters than the PL distribution, it is expected to be more expressive. In fact, when considering their ability to express the delta distribution, which is the target distribution for many permutation learning problems, we have the following result.

̸

Proposition 1. The PL distribution cannot exactly represent a delta distribution. That is, there does not exist an s such that p PL = δ σ for any σ ∈ S n , where δ σ ( σ ) = 1 and δ σ ( π ) = 0 for all π = σ . But the GPL distribution can represent a delta distribution exactly.

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 t ∈ R n × d , we use a neural network with parameters θ to construct p θ ( σ ′ t | X t ) . In particular, we treat n rows of X t as n tokens and use a Transformer architecture along with the time embedding of t and the positional encoding to predict the previously mentioned scores. For example, for the GPL distribution, to predict n 2 scores, we introduce n dummy tokens that correspond to the n permuted output positions. We then perform a few layers of masked self-attention ( 2 n × 2 n ) to obtain the token embedding Z 1 ∈ R n × d model corresponding to n input tokens and Z 2 ∈ R n × d model corresponding to n dummy tokens. Finally, the GPL score matrix is obtained as S θ = Z 1 Z ⊤ 2 ∈ R n × n . 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:

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 can not obtain q ( X t | X 0 ) analytically, as is done in common diffusion models [16, 3]. Therefore, we have to run the forward process to collect samples. Fortunately, it is efficient as the forward process only involves shuffling integers. We include more training details in Appendix E.

3.4 Denoising Schedule via Merging Reverse Steps

If 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 t of the reverse process, instead of predicting p θ ( X t -1 | X t ) , we can predict p θ ( X t ′ | X t ) for any 0 ≤ t ′ < t . Given a sequence of timesteps 0 = t 0 < · · · < t k = T , we can now model the reverse process as

To align with the literature of diffusion models, we call the list [ t 0 , . . . , t k ] the denoising schedule . After incorporating the denoising schedule in Eq. (10), we obtain the loss function:

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

t

0.0

0.2

0.4

0.6

0.8

1.0

TV Distance to Unifrom, n=100

(a)

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

0.2

0.4

0.6

0.8

1.0

(b)

Note that although we may not have the analytical form of q ( X t i | X t i -1 ) , we can draw samples from it. Merging is feasible if the support of p θ ( X t i -1 | X t i ) is equal or larger than the support of q ( X t i | X t i -1 ) ; 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 p θ ( σ ′ t | X t ) follows the PL or GPL distribution, as they have whole S n 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 t and X t +1 are similar, we 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, we can measure D TV ( q ( t ) , q ( t ′ ) ) for t = t ′ and D TV ( q ( t ) , u ) , where q ( t ) is the distribution at time t in the forward process and u is the uniform distribution. For riffle shuffles, the total variation distance can be computed exactly. Specifically, we first introduce the Eulerian Numbers A n,r [32], i.e ., the number of permutations in S n that have exactly r rising sequences where 1 ≤ r ≤ n . A n,r can be computed using the following recursive formula A n,r = rA n -1 ,r +( n -r +1) A n -1 ,r -1 where A 1 , 1 = 1 . We then have the following result.

Proposition 2. Let t = t ′ be positive integers. Then

and

̸

Note that Eq. (14) was originally given in [19]. We restate it here for completeness. Once the Eulerian numbers are precomputed, the TV distances can be computed in O ( n ) time instead of O ( n !) . Through extensive experiments, we have the following empirical observation. For the principle 1), choosing T so that D TV ( q ( T ) RS , u ) ≈ 0 . 005 yields good results. For the principle 2), a denoising schedule [ t 0 , . . . , t k ] with D TV ( q ( t i ) RS , q ( t i +1 ) RS ) ≈ 0 . 3 for most i works well. We show an example on sorting n = 100 four-digit MNIST images in Fig. 2.

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 or a few states with high probabilities under p θ ( X 0 |X ) . However, since we can only draw samples from p θ ( X 0 |X ) via running the reverse process, exact decoding is intractable. The simplest approximated method is greedy search, i.e ., successively finding the mode or an approximated mode of p θ ( X t i -1 | X t i ) . Another approach is beam search, which

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

↓

maintains a dynamic buffer of k 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 a hierarchical beam search that performs an inner beam search within n 2 scores at each step of the outer beam search. Further details are provided in Appendix C.

4 Experiments

We now demonstrate the general applicability and effectiveness of our model through a variety of experiments, including sorting 4-digit MNIST numbers, solving jigsaw puzzles, and addressing traveling salesman problems. Additional details are provided in the appendix due to space constraints.

4.1 Sorting 4-digit MNIST Images

We first evaluate our SymmetricDiffusers on the four-digit MNIST sorting benchmark, a wellestablished testbed for differentiable sorting [5, 8, 13, 20, 33, 34]. Each four-digit image in this benchmark is obtained by concatenating 4 individual images from MNIST. For evaluation, we employ several metrics to 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 in Table 3, combining PL with either beam search (BS) or greedy search yields good results in terms of Kendall-Tau 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. Finally, combining GPL and BS achieves the best accuracy in correctly sorting the entire sequence of images. Given that accuracy is the most

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

↓

challenging metric to improve, we selecte GPL and BS for all remaining experiments. More ablation study ( e.g ., denoising schedule) is provided in Appendix E.2.

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 performances drop significantly with long sequences ( e.g ., length of 52 or 100). Meanwhile, DiffSort performs the worse due to the error accumulation of its soft differentiable swap function [20, 33]. In contrast, our method is on par with Error-free DiffSort in sorting short sequences and significantly outperforms others on long sequences.

4.2 Jigsaw Puzzle

We then explore image reassembly from segmented "jigsaw" puzzles [29, 31, 39]. We evaluate the performance using the MNIST and the CIFAR10 datasets, which comprises puzzles of up to 6 × 6 and 4 × 4 pieces respectively. We add slight noise to pieces from the MNIST dataset to ensure background pieces are distinctive. To evaluate our models, we use Kendall-Tau coefficient, accuracy, correctness, RMSE (root mean square error of reassembled images), and MAE (mean absolute error) as metrics.

Table 1 presents results comparing our method with the Gumbel-Sinkhorn Network[29], Diffsort [34], and Error-free Diffsort [20]. DiffSort and Error-free DiffSort are primarily designed for sorting high-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 Gumbel-Sinkhorn performs better for tasks involving fewer than 4 × 4 pieces. In more challenging scenarios ( e.g ., 5 × 5 and 6 × 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 NP-complete combinatorial optimization problems which are solved using integer programming or heuristic solvers [2, 12]. There exists a vast literature on learning-based models to solve TSPs [22, 23, 18, 17, 6, 24, 10, 36, 21, 43, 30]. They often focus on the Euclidean TSPs, which are formulated as follows. Let V = { v 1 , . . . , v n } be points in R 2 . We need to find some σ ∈ S n such that ∑ n i =1 ∥ v σ ( i ) -v σ ( i +1) ∥ 2 is minimized, where we let σ ( n +1) := σ (1) . Further experimental details are provided in Appendix B.

We compare with operations research (OR) solvers and other learning based approaches on TSP instances with 20 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 is given by ( predicted length -( GT length ) ) / ( GT length ) . As shown in Table 4, SymmetricDiffusers achieves comparable results with both OR solvers and the state-of-the-art learning-based methods.

5 Conclusion

In this paper, we introduce a 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 is provably more expressive than the PL distribution. We further introduce a theoretically grounded "denoising schedule" to improve sampling and learning efficiency. Extensive experiments verify the effectiveness of our proposed model. In the future, we are interested in generalizing our model to general finite groups and exploring diffusion models on Lie groups.

341

References

David Applegate, Robert Bixby, Vasek Chvatal, and William Cook. Concorde tsp solver, 2006. 342

Sanjeev Arora. Polynomial time approximation schemes for euclidean traveling salesman and 343 other geometric problems. J. ACM , 45(5):753-782, Sep 1998. 344

Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. 345 Structured denoising diffusion models in discrete state-spaces, 2023. 346

Dave Bayer and Persi Diaconis. Trailing the dovetail shuffle to its lair. The Annals of Applied 347 Probability , 2(2):294 - 313, 1992. 348

Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable 349 sorting and ranking. In International Conference on Machine Learning , pages 950-959. PMLR, 350 2020. 351

Xavier Bresson and Thomas Laurent. The transformer network for the traveling salesman 352 problem, 2021. 353

Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Hang Li. Learning to rank: from pairwise 354 approach to listwise approach. In Proceedings of the 24th International Conference on Machine 355 Learning , ICML '07, pages 129-136, New York, NY, USA, 2007. Association for Computing 356 Machinery. 357

Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable ranking and sorting using 358 optimal transport. Advances in neural information processing systems , 32, 2019. 359

Persi Diaconis. Group representations in probability and statistics. Lecture notes-monograph 360 series , 11:i-192, 1988. 361

Zhang-Hua Fu, Kai-Bin Qiu, and Hongyuan Zha. Generalize a small pre-trained model to 362 arbitrarily large tsp instances, 2021. 363

E. N. Gilbert. Theory of shuffling. Bell Telephone Laboratories Memorandum , 1955. 364

Teofilo F. Gonzalez. Handbook of Approximation Algorithms and Metaheuristics (Chapman & 365 Hall/Crc Computer & Information Science Series) . Chapman & Hall/CRC, 2007. 366

Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic optimization of sorting 367 networks via continuous relaxations. In International Conference on Learning Representations , 368 2018. 369

Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. 370

Keld Helsgaun. An extension of the lin-kernighan-helsgaun tsp solver for constrained traveling 371 salesman and vehicle routing problems, Dec 2017. 372

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020. 373

Chaitanya K Joshi, Quentin Cappart, Louis-Martin Rousseau, and Thomas Laurent. Learning 374 tsp requires rethinking generalization. In International Conference on Principles and Practice 375 of Constraint Programming , 2021. 376

Chaitanya K. Joshi, Thomas Laurent, and Xavier Bresson. An efficient graph convolutional 377 network technique for the travelling salesman problem, 2019. 378

Shihan Kanungo. Mixing time estimates for the riffle shuffle. Euler Circle , 2020. 379

Jungtaek Kim, Jeongbeen Yoon, and Minsu Cho. Generalized neural sorting networks with error380 free differentiable swap functions. In International Conference on Learning Representations 381 (ICLR) , 2024. 382

Minsu Kim, Junyoung Park, and Jinkyoo Park. Sym-nco: Leveraging symmetricity for neural 383 combinatorial optimization, 2023. 384

Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional 385 networks, 2017. 386

Wouter Kool, Herke van Hoof, and Max Welling. Attention, learn to solve routing problems!, 387 2019. 388

Yeong-Dae Kwon, Jinho Choo, Byoungjip Kim, Iljoo Yoon, Youngjune Gwon, and Seungjai 389 Min. Pomo: Policy optimization with multiple optima for reinforcement learning, 2021. 390

Shen Lin and Brian W Kernighan. An effective heuristic algorithm for the travelling-salesman 391 problem. Operations research , 21(2):498-516, 1973. 392

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. 393

R. D. Luce. Individual Choice Behavior . John Wiley, 1959. 394

Colin L Mallows. Non-null ranking models. i. Biometrika , 44(1/2):114-130, 1957. 395

Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permuta396 tions with gumbel-sinkhorn networks. In International Conference on Learning Representations , 397 2018. 398

399

400

Yimeng Min, Yiwei Bai, and Carla P. Gomes. Unsupervised learning for solving the travelling salesman problem, 2024.

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

Mehdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. In European conference on computer vision , pages 69-84. Springer, 2016.

OEIS Foundation Inc. The eulerian numbers, entry a008292 in the On-Line Encyclopedia of Integer Sequences, 2024. Published electronically at http://oeis.org/A008292 .

Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Differentiable sorting networks for scalable sorting and ranking supervision. In International conference on machine learning , pages 8546-8555. PMLR, 2021.

Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Monotonic differentiable sorting networks. In International Conference on Learning Representations (ICLR) , 2022.

R. L. Plackett. The analysis of permutations. Applied Statistics , 24(2):193 - 202, 1975.

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 , pages 263-346. Springer Berlin Heidelberg, Berlin, Heidelberg, 2004.

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 , pages 3949-3957, 2017.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research , pages 2256-2265, Lille, France, 07-09 Jul 2015. PMLR.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and 426 Ben Poole. Score-based generative modeling through stochastic differential equations, 2021. 427

Zhiqing Sun and Yiming Yang. Difusco: Graph-based diffusion solvers for combinatorial 428 optimization, 2023. 429

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, 430 Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. 431

Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal 432 Frossard. Digress: Discrete denoising diffusion for graph generation, 2023. 433

Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce434 ment learning. Machine Learning , 8(3-4):229-256, 1992. 435

A Additional Details of the GSR Riffle Shuffle Model 436

437

438

439

440

441

There are many equivalent definitions of the GSR riffle shuffle. Here we also introduce the Geometric Description [4], which is easy to implement (and is how we implement riffle shuffles in our experiments). We first sample n points in the unit interval [0 , 1] uniformly and independently, and suppose the points are labeled in order as x 1 < x 2 < · · · < x n . Then, the permutation that sorts the points { 2 x 1 } , . . . , { 2 x n } follows the GSR distribution, where { x } := x -⌊ x ⌋ is the fractional part of x .

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

B Details of Our Network Architecture

We now discuss how to use neural networks to produce the parameters of the distributions discussed in Section 3.2.1 and 3.2.2. Fix time t , and suppose X t = ( x ( t ) 1 , . . . , x ( t ) n ) ⊤ ∈ R n × d . Let encoder θ be an object-specific encoder such that encoder θ ( X t ) ∈ R n × d model . For example, encoder θ can be a CNN if X t is an image. Let

where time_embd is the sinusoidal time embedding. Then, we would like to feed the embeddings into a Transformer encoder [44]. Let transformer_encoder θ be the encoder part of the Transformer architecture. However, each of the distributions we discussed previously has different number of parameters, so we will have to discuss them separately.

Inverse Transposition. For Inverse Transposition, we have n +1 parameters. To obtain n +1 tokens from transformer_encoder θ , we append a dummy token of 0's to Y t . Then we input ( y ( t ) 1 , . . . , y ( t ) n , 0 ) ⊤ into transformer_encoder θ to obtain Z ∈ R ( n +1) × d model . Finally, we apply an MLP to obtain ( s 1 , . . . , s n , k ) ∈ R n +1 .

Inverse Insertion, Inverse Riffle Shuffle, PL Distribution. These three distributions all require exactly n parameters, so we can directly feed Y t into transformer_encoder θ . Let the output of transformer_encoder θ be Z ∈ R n × d model , where we then apply an MLP to obtain the scores s θ ∈ R n .

The GPL Distribution. The GPL distribution requires n 2 parameters. We first append n dummy tokens of 0's to Y t , with the intent that the j th dummy token would learn information about the j th column of the GPL parameter matrix, which represents where the j th component should be placed. We then pass ( y ( t ) 1 , . . . , y ( t ) n , 0 , . . . , 0 ) ⊤ ∈ R 2 n × d model to transformer_encoder θ . When computing attention, we further apply a 2 n × 2 n attention mask

The reason for having B as an upper triangular matrix of -∞ is that information about the j th 464 component should only require information from the previous components. Let 465

where Z 1 , Z 2 ∈ R n × d model . Finally, we obtain the GPL parameter matrix as S θ = Z 1 Z ⊤ 2 ∈ R n × n . For hyperparameters, we refer the readers to Appendix E.4.

C Additional Details of Decoding

Greedy Search. At each timestep t i in the denoising schedule, we can greedily obtain or approximate the mode of p θ ( X t i -1 | X t i ) . We can then use the (approximated) mode X t i -1 for the next timestep p θ ( X t i -2 | X t i -1 ) . Note that the final X 0 obtained using such a greedy heuristic may not necessarily be the mode of p θ ( X 0 |X ) .

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

for all i ∈ [ n ] . In particular, we have

Beam Search. We can use beam search to improve the greedy approach. The basic idea is that, at each timestep t i in the denoising schedule, we compute or approximate the topk -most-probable results from p θ ( X t i -1 | X t i ) . For each of the topk results, we sample topk from p θ ( X t i -2 | X t i -1 ) . Now we have k 2 candidates for X t i -2 , and we only keep the top k of the k 2 candidates.

However, it is not easy to obtain the topk -most-probable results for some of the distributions. Here we provide an algorithm to approximate topk of the PL and the GPL distribution. Since the PL distribution is a strict subset of the GPL distribution, it suffices to only consider the GPL distribution with parameter matrix S . The algorithm for approximating topk of the GPL distribution is another beam search. We first pick the k largest elements from the first row of S . For each of the k largest elements, we pick k largest elements from the second row of S , excluding the corresponding element picked in the first row. We now have k 2 candidates for the first two elements of a permutation, and we only keep the topk -most-probable candidates. We then continue in this manner.

D Proofs

Proposition 1. The PL distribution cannot exactly represent a delta distribution. That is, there does not exist an s such that p PL = δ σ for any σ ∈ S n , where δ σ ( σ ) = 1 and δ σ ( π ) = 0 for all π = σ . But the GPL distribution can represent a delta distribution exactly.

̸

Proof. Assume for a contradiction that there exists some σ ∈ S n and s such that PL s = δ σ . Then we have

Since each of the term in the product is less than or equal to 1, we must have

which happens if and only if s σ ( j ) = -∞ for all j ≥ 2 . But this contradicts (16).

We then show that the GPL distribution can represent a delta distribution exactly. To see this, we fix σ ∈ S n . For all i ∈ [ n ] , we let s i,σ ( i ) = 0 and s i,j = -∞ for all j = σ ( i ) . Then GPL ( s ij ) = δ σ .

Proposition 2. Let t = t ′ be positive integers. Then

and

̸

Proof. Let σ ∈ S n . It was shown in [4] that

where r 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

as claimed. For (14), replace q ( t ) RS ( σ ) with u ( σ ) = 1 n ! in the above derivations. 502

̸

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

E Additional Details on Experiments

E.1 Datasets

Jigsaw Puzzle. We created the Noisy MNIST dataset by adding i.i.d. Gaussian noise with a mean of 0 and a standard deviation of 0.01 to each pixel of the MNIST images. No noise was added to the CIFAR-10 images. The noisy images are then saved as the Noisy MNIST dataset. During training, each image is divided into n × 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 as 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 on 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-digit MNIST numbers on the fly.

TSP. We take the TSP-20 dataset from [17] 1 . The train set consists of 1,512,000 graphs with 20 nodes, 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 [1]. The test set consists of 1,280 graphs with 20 nodes, with ground truth tour generated by the Concorde solver as well.

E.2 Ablation Studies

Choices for Reverse Transition and Decoding Strategies. As demonstrated in Table 5, 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. It is important to note that strategies such as random transposition and random insertion paired with their respective inverse operations, are less suitable for tasks with longer sequences. This limitation is 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 6, 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 0 accuracy justifies our motivation to use diffusion to break down learning into smaller steps. The result we get also matches with our proposed heuristic in Section 3.4.

E.3 Latent Loss in Jigsaw Puzzle

In the original setup of the Jigsaw Puzzle experiment using the Gumbel-Sinkhorn network [29], the permutations are latent. That is, the loss function in Gumbel-Sinkhorn is a pixel-level 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 to their original order. Therefore, for fair comparison with the Gumbel-Sinkhorn network 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

In our original setup, we defined p ( X t -1 | X t , σ ′ t ) as a delta distribution δ ( X t -1 = Q σ ′ t X t ) , but this would require that we know the permutation that turns X t -1 to X t , which is part of the ground truth. So instead, we parameterize p ( X t -1 | X t , σ ′ t ) as a Gaussian distribution N ( X t -1 | Q σ t X t , I ) . At the same time, we note that to find the gradient of (12), it suffices to find the gradient of the log of (17).

We use the REINFORCE trick [46] to find the gradient of log p θ ( X t -1 | X t ) , which gives us 547

where we have used Monte-Carlo estimation in the last step, and σ (1) t , . . . , σ ( N ) t ∼ p θ ( σ t | X t ) . We 548 further add an entropy regularization term -λ · E p θ ( σ t | X t ) [log p θ ( σ t | X t )] to each of log p θ ( X t -1 | X t ) . 549 Using the same REINFORCE and Monte-Carlo trick, we obtain 550

where σ (1) t , . . . , σ ( N ) t ∼ p θ ( σ t | X t ) . Therefore, we have 551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

where σ (1) t , . . . , σ ( N ) t ∼ p θ ( σ t | X t ) . We then substitute in

for all n ∈ [ N ] . Finally, we also subtract the exponential moving average weight as a control variate for variance reduction, where the exponential moving average is given by ema ← ema_rate · ema + (1 -ema_rate ) · weight for each gradient descent step.

E.4 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.

Jigsaw Puzzle. For the Jigsaw Puzzle experiments, we use the AdamW optimizer [26] with weight decay 1e-2, ε = 1e-9, and β = (0 . 9 , 0 . 98) . We use the Noam learning rate scheduler given in [44] 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), we use Monte-Carlo estimation for the expectation and sample 3 trajectories. For REINFORCE, we sampled 10 times for the Monte-Carlo estimation in (18), and we used an entropy regularization rate λ = 0 . 05 and an ema_rate of 0.995. The neural network architecture and related hyperparameters are given in Table 7. The denoising schedules, with riffle shuffles as the forward process and GPL as the reverse process, are give in Table 8. For beam search, we use a beam size of 200 when decoding from GPL, and we use a beam size of 20 when decoding along the diffusion denoising schedule.

×

Sort 4-Digit MNIST Numbers. For the task of sorting 4-digit MNIST numbers, we use the exact 570 training and beam search setup as the Jigsaw Puzzle, except that we do not need to use REINFORCE. 571 The neural network architecture is given in Table 9. The denoising schedules, with riffle shuffles as 572 the forward process and GPL as the reverse process, are give in Table 10. 573

TSP. For solving the TSP, we perform supervised learning to train our SymmetricDiffusers to solve 574 the TSP. Let σ ∗ be an optimal permutation, and let X 0 be the list of nodes ordered by σ ∗ . We note 575 that any cyclic shift of X 0 is also optimal. Thus, for simplicity and without loss of generality, we 576 always assume σ ∗ (1) = 1 . In the forward process of SymmetricDiffusers, we only shuffle the second 577 to the n th node (or component). In the reverse process, we mask certain parameters of the reverse 578 distribution so that we will always sample a permutation with σ t (1) = 1 . 579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

The architecture details are slightly different for TSP, since we need to input both node and edge features into our network. Denote by X t the ordered list of nodes at time t . We obtain Y t ∈ R n × d model as in Eq. (15), where encoder θ is now a sinusoidal embedding of the 2D coordinates. Let D t ∈ R n × n be the matrix representing the pairwise distances of points in X t , respecting the order in X t . Let E t ∈ R ( n 2 ) be the flattened vector of the upper triangular part of D t . We also apply sinusoidal embedding to E t and add time_embd ( t ) to it. We call the result F t ∈ R ( n 2 ) × d model .

Now, instead of applying the usual transformer encoder with self-attentions, we alternate between cross-attentions and self-attentions. For cross-attention layers, we use the node representations from the previous layer as the query, and we always use K = V = F t . We also apply an attention mask to the cross-attention, so that each node will only attend to edges that it is incident with. For selfattention 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 a cross-attention layer, and the last layer 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, 8 attention heads, d model = 256 , feed-forward hidden dimension 1024, and dropout rate 0.1.

For training details on the TSP-20 task, we use the AdamW optimizer [26] with weight decay 1e-4, ε = 1e-8, and β = (0 . 9 , 0 . 999) . We use the cosine annealing learning rate scheduler starting from 2e-4 and ending at 0. We train for 50 epochs with a batch size of 512. When computing the loss (12), we use Monte-Carlo estimation for the expectation and sample 1 trajectory. We use a denoising schedule of [0 , 4 , 5 , 7] , with riffle shuffles as the forward process and GPL as the reverse process. Finally, we use beam search for decoding, and we use a beam size of 256 both when decoding from GPL and decoding along the denoising schedule.

E.5 Baselines Implementation Details

Gumbel-Sinkhorn Network. We have re-implemented the Gumbel-Sinkhorn Network [29] for application on jigsaw puzzles, following the implementations provided in the official repository 2 . To ensure a fair comparison, we conducted a thorough grid search of the model's hyper-parameters. The parameters included in our search space are as follows,

Diffsort & Error-free Diffsort We have implemented two differentiable sorting networks from the official repository 3 specific to error-free diffsort. For sorting 4-digit MNIST images, error-free diffsort employs TransformerL as its backbone, with detailed hyperparameters listed in Table 12. Conversely, Diffsort uses a CNN as its backbone, with a learning rate set to 10 -3 . 5 ; the relevant hyperparameters are outlined in Table 13.

For jigsaw puzzle tasks, error-free diffsort continues to utilize a transformer, whereas Diffsort employs a 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 × 3 puzzles, we apply the same configuration as used for sorting tasks with a sequence length of 9.

TSP. For the baselines for TSP, we first have 4 traditional operations research solvers. Gurobi [14] and Concorde [1] are known as exact solvers, while LKH-3 [15] is a strong heuristic and 2-Opt [25]

is a weak heuristic. For LKH-3, we used 500 trials, and for 2-Opt, we used 5 random initial guesses 618 with seed 42. 619

For the GCN model[18], we utilized the official repository 4 and adhered closely to its default 620 configuration for the TSP-20 dataset. For DIFUSCO[43], we sourced it from its official repository 5 621 and followed the recommended configuration of TSP-50 dataset, with a minor adjustment in the batch 622 size. We increased the batch size to 512 to accelerate the training process. For fair comparison, we 623 also remove the post-processing heuristics in both models during the evaluation. 624

625

F Limitations

Despite the success of this method on various tasks, the model presented in this paper still requires a 626 time-space complexity of O ( n 2 ) due to its reliance on the parametric representation of GPL and the 627 backbone of transformer attention layers. This complexity poses a significant challenge in scaling up 628 to applications involving larger symmetric groups or Lie groups. 629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

658

659

660

661

662

663

664

665

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our abstract and Section 1 accurately summarize the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of the work in Appendix F.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide complete proof for Proposition 1 and Proposition 2 in Appendix D.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In Appendix E, we fully disclose all information to reproduce our experimental results, including dataset preparation, training details, and choices of hyperparameters as well as baselines' implementation details.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We've included codes to reproduce the main results in the supplemental material. We also attach a detailed README file that provides sufficient instructions.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All experimental settings and details are specified in Appendix E.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All reported experimental results are averaged over at least three runs with different random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide information on the computation resources used for our experiments in Appendix E.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We preserve anonymity with the NeurIPS Codes of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: There is no societal impact of the work performed.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited the original paper of our reference code and datasets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
I am unable to review this paper as it lies outside my area of expertise.

STRENGTHS:
I am unable to review this paper as it lies outside my area of expertise.

WEAKNESSES:
I am unable to review this paper as it lies outside my area of expertise.

QUESTIONS:
I am unable to review this paper as it lies outside my area of expertise.

LIMITATIONS:
I am unable to review this paper as it lies outside my area of expertise.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors aim to create a discrete diffusion model that generates permutations. This model can then be used to solve combinatorial problems including jigsaws and travelling salesman problems. To formulate their model they cover a range of forward shuffling strategies and discuss how to parametrize the reverse transition. During sampling, they also use beam search to find high probability samples. They find their method performs competitively on computational experiments.\n\nSTRENGTHS:\nThe idea of using riffle shuffles to create a corruption process over permutations and then parametrizing the time reversal of this process is novel and interesting. I enjoyed reading the paper. I believe further work will build on this as there are many instances in machine learning where needing to learn permutations crops up.\n\nThe paper is quite well written and easy to understand. It is not overloaded with mathematical equations and intuition is given for some concepts.\n\nThe experimental results seem promising as it performs on par or better (especially in high dimensions) than previous methods for learning permutations. I appreciate the ablation studies into the greedy search vs beam search and types of shuffle (in the appendix).\n\nWEAKNESSES:\nI think some further clarification is required for the weaknesses of the random transposition and random insertion style of card shuffling. You later say that you can merge steps of the forward process if each individual step does not induce enough mixing and so the stated weakness that these styles of shuffle have slow mixing seems moot.\n\nYou mention that you do not have access to $q(X_t | X_0)$, and I think it should also be discussed that $q(X_{t-1} | X_t, X_0)$ is also unavailable since this distribution is used in standard diffusion models to re-write the variational bound in a lower variance form, see Appendix A in https://arxiv.org/pdf/2006.11239 .\n\nYou dedicate a lot of space to discussing the various forward noising processes with different shuffling methods, which is quite interesting. However, the ablations with these different styles of shuffle are in the appendix and I think it should be in the main since they have been given such prominence earlier on in the discussion of the method, it is strange they are not included in the main experiments.\n\nI find it difficult to follow the description of the inverse transposition parametrization, there is no intuition given for the functions $\\phi$ and $\\psi$ nor the functional form of $p_{IT}(\\sigma)$. Perhaps this is due to space limitations but since inverse transposition is not in the main experiments (see above point), I think you should either relegate a lot of this to the appendix if you only use the riffle shuffle in practice, or try and shift the wording to properly explain these types of forward and inverse process and have experiments for them in the main.\n\nQUESTIONS:\nWhen you compute the loss function in equation (10), do you do T calls to the neural network because you calculate the full variational bound for every iteration? This would indeed be very memory intensive, isn't there an alternative where although you need to run the full forward process since $q(X_t | X_0)$ is intractable, you could still only compute the loss on some subset of the $p_\\theta(X_{t-1} | X_t)$ reverse transitions.\n\nLIMITATIONS:\nThe authors do a good job of discussing the limitations of various parametrizations of their method.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a discrete diffusion model to learn distribution over the finite symmetric group $S_n$. The forward process is built off of random walks on finite groups (in this case, card shuffles), and the paper learns to reverse this diffusion process with standard discrete diffusion arguments.\n\nSTRENGTHS:\nOverall, I really like the paper's contributions and presentation.\n\n* The idea is a neat application of the discrete diffusion ideas to an important area. In particular, the structure of $S_n$ is sufficiently different from standard image/text datasets as to necessitate this paper.\n\n* The presentation is very good and the contributions are numerous.\n\n* For the experiments listed, the method seems to provide a very strong improvement over baseline methods. In particular, these other methods are based on fundamentally different technology, so this highlights that discrete diffusion can become a very promising direction here.\n\nWEAKNESSES:\nThere are three primary weaknesses. These should all be addressable to some degree, and I'll take any response into consideration when recalibrating my final score.\n\n1. The model proposes to directly learn the reverse transition densities $p_\\theta(X_{t - 1}, X_t)$. The issue with doing this for standard diffusion models is that this seems to hurt model training since it increases the variance of training (as, in particular, one must sample the two $X_{t - 1}, X_t$ for training instead of just one $X_t$). As such, most works use the (ultimately equivalent) mean/score-parameterizations [1, 2]. I would want to hear a bit more about if this would be applicable in the $S_n$ case (and training with this parameterization might improve the model) or if this is not possible.\n\n2. (Related to the above). Since most modern discrete diffusion methods are formulated in continuous time, I think the paper would benefit greatly with a discussion about potentially extending the current methods to this realm. In particular, works like [3, 4, 5] have established a working theory for discrete diffusion in continuous time, so it would be beneficial to discuss how the proposed framework might fit into the established theory.\n\n3. The experiments, while showing good results, do not show that the method is particularly scalable, which seems to be a fundamental problem in prior work that was explicitly mentioned in this paper. In particular, it seems that the maximum value of $n$ in $S_n$ is 100. While some discussion is made here that talks about transformer layers, I think large values of $n$ aren't that big of an issue in transformers due to systems like Flash Attention. So, it should be made more clear if this is a fundamental problem with the existing method, or a larger scale example (even toy) should be presented.\n\n[1] https://arxiv.org/abs/2006.11239\n\n[2] https://arxiv.org/abs/2011.13456\n\n[3] https://arxiv.org/abs/2205.14987\n\n[4] https://arxiv.org/abs/2211.16750\n\n[5] https://arxiv.org/abs/2310.16834\n\nQUESTIONS:\nN/A beyond addressing the weaknesses above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces SymmetricDiffusers, a new approach to learning complex distributions. It works by breaking down the problem into simpler steps: learning how to reverse a transformation using deep neural networks. The authors identify a particularly effective method for this reversal step (the riffle shuffle) and provide guidance on choosing the right length for the process based on mathematical properties. Additionally, they propose a more powerful alternative to a common distribution (the generalized Plackett-Luce distribution) and a theoretically sound strategy for improving efficiency (the denoising schedule). Experiments show that SymmetricDiffusers performs extremely well on various tasks, including sorting images, solving puzzles, and optimizing routes.\n\nSTRENGTHS:\nThe proposed method in general is interesting and seems effective in the examples studied in this paper.\n\nAlso, the problem studied is interesting.\n\nWEAKNESSES:\nMy main concern is the proposed method is not the state-of-the-art. A clear approach is to learn invariant features and equivariant group actions using some existing methods like that proposed in Robin Winter, et al. Unsupervised Learning of Group Invariant and\nEquivariant Representations, NeurIPS 2022. Such an approach. \n\nIn Robin Winter et al.'s paper, the authors have studied the symmetric group and my understanding is disentangling invariant features and equivariant groups enables the design of flexible diffusion models since you only need to perform diffusion modeling in the invariant latent space. \n\nI want to see the comparison of the approach proposed in Robin Winter et al's paper. My understanding is that Robin Winter's approach enables building state-of-the-art diffusion models for molecular generation.\n\nQUESTIONS:\nMy main concern is the proposed method is not the state-of-the-art. A clear approach is to learn invariant features and equivariant group actions using some existing methods like that proposed in Robin Winter, et al. Unsupervised Learning of Group Invariant and\nEquivariant Representations, NeurIPS 2022. Such an approach. \n\nIn Robin Winter et al.'s paper, the authors have studied the symmetric group and my understanding is disentangling invariant features and equivariant groups enables the design of flexible diffusion models since you only need to perform diffusion modeling in the invariant latent space. \n\nI want to see the comparison of the approach proposed in Robin Winter et al's paper. My understanding is that Robin Winter's approach enables building state-of-the-art diffusion models for molecular generation.\n\nLIMITATIONS:\nMy main concern is the proposed method is not the state-of-the-art. A clear approach is to learn invariant features and equivariant group actions using some existing methods like that proposed in Robin Winter, et al. Unsupervised Learning of Group Invariant and\nEquivariant Representations, NeurIPS 2022. Such an approach. \n\nIn Robin Winter et al.'s paper, the authors have studied the symmetric group and my understanding is disentangling invariant features and equivariant groups enables the design of flexible diffusion models since you only need to perform diffusion modeling in the invariant latent space. \n\nI want to see the comparison of the approach proposed in Robin Winter et al's paper. My understanding is that Robin Winter's approach enables building state-of-the-art diffusion models for molecular generation.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-13"}