Learning Distributions over Permutations and Rankings with Factorized Representations
Authors: Daniel Severo
Brian Karrer
Niklas Nolte
Venue: ICLR 2026 Poster (2026)
Keywords: permutations
rankings
fisher-yates
insertion-vectors
ranking
movielens
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We improve training and inference of probabilistic models over permutations and rankings using factorized representations, significantly outperforming the current SOTA.

Abstract: Learning distributions over permutations is a fundamental problem in machine learning, with applications in ranking, combinatorial optimization, structured prediction, and data association. Existing methods rely on mixtures of parametric families or neural networks with expensive variational inference procedures. In this work, we propose a novel approach that leverages alternative representations for permutations, including Lehmer codes, Fisher-Yates draws, and Insertion-Vectors. These representations form a bijection with the symmetric group, allowing for unconstrained learning using conventional deep learning techniques, and can represent any probability distribution over permutations. Our approach enables a trade-off between expressivity of the model family and computational requirements. In the least expressive and most computationally efficient case our method subsumes previous families of well established probabilistic models over permutations, including Mallow's and the Repeated Insertion Model. Experiments indicate our method significantly outperforms current approaches on the jigsaw puzzle benchmark, a common task for permutation learning. However, we argue this benchmark is limited in its ability to assess learning probability distributions, as the target is a delta distribution (i.e., a single correct solution exists). We therefore propose two additional benchmarks: learning cyclic permutations and re-ranking movies based on user preference. We show that our method learns non-trivial distributions even in the least expressive mode, while traditional models fail to even generate valid permutations in this setting.

Published as a conference paper at ICLR 2026
LEARNING DISTRIBUTIONS OVER PERMUTATIONS AND RANKINGS WITH FACTORIZED REPRESENTATIONS
Daniel Severo
Brian Karrer FAIR at Meta br briankarrerdmeta.coo
Niklas Nolte FAIR at Meta no nolte@mer com
FAIR at Meta
dsevero@meta.co com
ABSTRACT
Learning distributions over permutations is a fundamental problem in machine learning, with applications in ranking, combinatorial optimization, structured pre- diction, and data association. Existing methods rely on mixtures parametric families neural networks with expensive variational inference procedures. In this work, we propose a novel approach that leverages alternative representations for permutations, including Lehmer codes, Fisher- Yates draws, and Insertion- Vectors. These representations form bijection with the symmetric group, allowing for un- constrained learning using conventional deep learning techniques, and can represent any probability distribution over permutations. Our approach enables a trade-off between expressivity of the model family and computational requirements. In the least expressive and most computationally efficient case, our method subsumes pre- vious families well established probabilistic models over permutations, including Mallow's and the Repeated Insertion Model. Experiments indicate our method significantly outperforms current approaches on the jigsaw puzzle benchmark, a common task for permutation learning. However, argue this benchmark is limited in its ability to assess learning probability distributions, as the target is a delta distribution (i.e., single correct solution exists). We therefore propose additional benchmarks: learning cyclic permutations and re-ranking movies based user preference. We show that our method learns non-trivial distributions even in the least expressive mode, while traditional models fail to even generate valid permutations in this setting.
context X,,
permutation repr. FF(X)
M
3,, 1, 4,2,
fortse
M
M
6,,,,,,6,,,4
forest
0
M
M
[3, 1, 2,4,4
fortse
1,,1,5,6,2,
forest
M
1,,,,,,466,
fortes
3,,,5,,,222 6,5,6,2,44
forest
181,5,4,66,2 X
fortes
Figure 1: Overview of our method unscrambling the sequence "osftre" autoregressively using one of the representations we consider in this work: Fisher- Yates draws (Fisher & Yates 1953). We condition on a referenceecontext (green) and the current input (blue) to sample values for the masked tokens (white). The model samples a permutation that unscrambles "forest" on the left, and "fortes" on the right. At any point in generation, the partially-man sequence corresponds to some valid permutation.
INTRODUCTION
Learning in the space of permutations is a fundamental problem with applications ranging from rank- ing for recommendation systems (Feng al. 2021), combinatorial optimization, haaming-to-rante

---

Published as a conference paper at ICLR 2026
[Burges 2010), and data cleaning Kaaaassury et al. 2025). Classical probabilistiic models for permu- tations include the Plackett-Luce (Plackett 1975; Luce et al. 1959) and Mallows (Mallows 1957) distributions, which only represent a limited set of probability distributions
permutations (e.g., Plackett-Luce cannot model a delta distribution). These limitations have been addressed in existing literature by considering mixtures Lu & Boutilier (2014), which require expensive varia- tional inference procedures for learning and inference. More recently, several works have proposed methods for learning arbitrary probability distributions over permutations using neural networks, in the framework of diffusion (Zhang et al. 2024) and convex relaxations (Mena et al. 2018 (see Section 2 for an overview).
In this work, we develop models that can represent any probability distribution over permutations and can be trained with conventional deep learning techniques, including any-order masked language modelling (MLM) (Uria et al., 2016 Larochelle & Murray 2011). and autoregressive next-token- prediction (AR or NTP) (Shannon 1948). We leverage alternative representations for permutations (beyond the usual inline notation) that form a bijection with the symmetric group, allowing for unconstrained learning. The representations we consider stem from well-established algorithms in the permutation literature, such as factorial indexing (Lehmer codes (Lehmer 1960)), generating random permutations (Fisher- Yates draws (Fisher & Yates 1953)). and modelling sub-rankings (Insertion- Vinsertion-Vectrrs Vectors (Doignon al. 2004 Lu & Boutilier] (014)); which all have varying support for their sequence-elemenn that are a function of the position in the sequence (Section 3.1).
To trade off compute and expressivity, MLMs have the capability of sampling multiple permutation elements independently with one forward pass through the neural network. Aforementioned represen- tations always produce valid permutations inference time for any amount of compute spent, even in the fully-factorized case when all tokens are unmasked in a single forward-pass.
Decoding the inline notation of the permutation from the representation is trivial in the case of Lehmer and Fisher- Yater (Kunze et al. (2024a)). In Theorem 4.3 we establish relationship between permutation's 's inverse, and its Lehmer and Insertion- Vector representations, which allows to develop fast decoding algorithm for Insertion- Vectors that can be applied in batch, significantly improving inference time compute.
Our methods establishes new state-of oftte-of-the-art results on the common benchmark of solving jigsaw puzzles (Mena et al. 2018 Zhang et al., 2024), significantly outperforming previous diffusion and convex-relaxation based approaches. However, we also argue this benchmark is inadequate to evaluate learning probability distributions over permutations, as each puzzles contains only one permutation that unscrambles it (i.c., the target distribution is delta function). We therefore propose two new benchmarks, which require learning non-trivial distributions: learning cyclic permutations (Section 5.2) and re-ranking dataset (Section 5.3).
of movies based on observed user preference in the MovieLens
In summary, contributions are four-fold. We:
develop new methods for supervised learning of arbitrary probability distri- butions over permutations that (1) assign zero probability to invalid permutations; (2) can trade-off expressivity for compute at sampling time, without re-training: (3) can learn non- trivial, fully-factorized distributions; (4) is trained with conventional language modelling techniques with a cross-entropy loss: (5) is extremely fast at sampling time;
(Section 4.2)
(Section [5.T) establish state-of-the-arr on the common benchmark of jigsaw puzzles, signifi- cantly outperforming current baselines;
(Sectionn5.2]amm |5. and Section (3.3 define two new benchmarks: learning cyclic permutations and re-ranking based on user preference data, that require learning non-trivial distributions;
(Section
(Theorem 4.3 establish new relationship between insertion-vecton vectors,, inverse permutations, and Lehmer codes that result in an efficient decoding scheme for insertion-vectt vectors..
RELATED WORK
Generative models and objectives. We utilize generative models parametrized by transformers Vaswani et al. 2017). as commonly employed in language modeling. Specifically, we utilize Masked Language Modeling (MLM) and next-token prediction (NTP or AR). The concept of NTP goes back
2

---

Published as a conference paper at ICLR 2026
W
F
FY, 3, swap with further
Vo = 0, insert at
Vi 0, insert 2 at
F
FY, 1, swap with further
V/2 1. insert at
F
w
FY, 0. pass
V/ 3, insert at
F
L
w FY, 0. pass (always)
V4 insert
Figure 2: (Left) Illustration for the Fisher- Yates algorithm for shuffling, defining bijection with permutations. In this example, (X) = [3, 1.0.0) = X [4,3,2 Small numbers in the bottom-right corner of each box represent the draw value required to swap the current element with that position. (Right) Illustration of the generative process defined by the insertion vector, for reference permutation Xet = [1, 2, 3, 3,4,5 5). At each step, the current element of the reference is inserted immediately to the left of V,, and values to the right are shifted right one position to accommodate, Small numbers at the bottom-left corner represent the slot index. In this example, V(X) =
[0, 0, 1, 2] = X [2, 5, 1.
as far as Shannon 1948) and has been applied with great success in language modeling within the last decade, see e.g. Radford et al. (2019); Meta (2024) and many more. Popularized through BERT (Devlin et al., 2019), MLM has been identified as a viable tool for language understanding. More
recently, forms MLM have been derived as a special case discrete diffusion (Austin et al.) 2021
He et al. 2022 Kitouni et al. 2024). where the noise distribution is delta distribution on the masked
state, and have shown promise in generative language modeling (Sahoo et al. 2024 Shi et al. 2025
Nie et al. 2025).
Permutation and Preference Modeling. Notable families of distributions over permutations
include the Plackett-Luce distribution (Plackeu 1975; Luce al. 1959) and the (generalized)
Mallows model (Mallows 1957). both of which have restricted expressivity. Doignon et al. (2004);
Lu & Boutilier (2014) propose the Repeated Insertion Model (RIM) and generalized version (GRIM)
to learn Mallows models and mixtures thereof, which itself also uses the same insertion representation
used this paper. These methods are detailed in Section 3.2
A prominent line of related work approaches permutation learning using differentiable ordering. One common strategy is to relax the discrete problem into continuous space-either by relaxing permutation matrices (Grover et al. [0199 Cuturi et al. 2019) by using differentiable swapping methods (Petersen et al. [0022] Kim et al., 2024). A notable baseline for us is the work Mena et al.. (2018), who utilize the continuous Sinkhorn operator regress to specific permutations, rather than distributions over possible permutations.
Using Lehmer codes for permutation learning has been considered by Diallo et al. (020), but only in the AR context and with a different architecture than considered in this work; as well as Malagón
al. (2025) sample solutions certain optimization problems in their framework (see "4.2 Case 2:
The First-Order Marginal Probabilities Model" in their paper). Recently, Sicently,Zhhang et al. (2024) joined the concepts of discrete space difffusion and differentiable shuffling methods propose an expressive generative method dubbed SymmetricDiffusen SymDiff for short. Inspired from random walks on permutations, they identify the riffle shuffle (Gilbert 1955) their forward process. To model the reverse process, the paper introduces generalized version of the Plackett-Luce distribution. This work serves as our most relevant and strongest baseline.
3 BACKGROUND
A short introduction permutations is given in Section A.I
Notation Sequences of random variables are denoted by capital letters X, L, V, and FY. Subscriis X,, L;, V, and FY, indicate their elements. Contiguous intervals are denoted by [n] = [1,2. and [n) = j0, 1..... n 1]. For some set S with elements $j [n],, let Xs (X,, (X......,, Xx, be the set of elements in X restricted to indices in S. For an ordered collection of sets S,, we denote unions Sez Ufey S Sy. The Lehmer code (Lehmer 1960), Fisher-Yates (Fisher & Yates 1953).
3

---

Published as a conference paper at ICLR 2026
and Insertion vectoor (Doignon et al. 2004 Lu & Boutilier 2014) representations of a permutation X will be denoted by [(X), FY (X), and (X), respectively. We sometimes drop the dependence on X when clear from context when defining distributions over these representations directly, All
logarithms are base
3.1 REPRESENTATIONS OF PERMUTATIONS
Lehmer Codes (Lehmer 1960). A Lehmer code is an alternative representation the inline notation of permutation. The Lehmer code (XX) of a permutation X on [n] is sequence of length that counts the number inversions at each position in the sequence. Inversions can be counted to the left or right, with one of the following 2 definitions,
(1) An example of a right-Lehmee code is given in Figure 3, The right-Lehmer code is commonly used index permutations in the symmetric group, as it is bijective with the factorial number system. The i-th element L(X); of the right-Lehmer has domain [n i + 1), and (i) for the left-Lehmer code. necessary and sufficient condition for a Lehmer code to represent valid permutation for its elements to be within their respective domains. The manhattan distance between Lehmer codes relates to the number of transpositions needed to convert between their respective permutations, establishing metric-space interpretation. This is formalized in Theorem B. As direct consequence, the sum Σ, L(X), equals the number adjacent transpositions required to recover the identity permutation, known as Kendall's 's tau distance (Kendall [938). Code to convert between inline notation and right- or left-Lehmer codes is given in Section D.1
Left: L(X), = Kj <i: X, > X or Right: =(X); X, X.H-
Fisher-Ya Yates Shuffle (Fisher & Yates 1953). The Fisher- Yates Shuffle is algorithm commonly used to generate uniformly distributed permutations. The procedure is illustrated in Figure 2 At each step, the element the current index is swapped with randomly selected element to the right, and after steps is guaranteed to produce uniformly distributed permutation if the initial sequence valid permutation. The index sampled at each step, FY, are referred as the "draws". Each resulting permutation X can be produced with exactly unique sequence of draws (XXX). implying the set possible draw-sequences forms bijections with the symmetric group (Fisher & Yates 1953). During the Fisher- Yates shuffle it possible to sample 0, resulting in no swap (see a "pass" step in Figure for an example). If sampling is restricted such that FY; > O, then the procedure is guaranteed to produce cyclic permutation and is known as Sattolo's 's Algorithm (Sattolo 1986).
Decoding a batch of Fisher- Yates representations can parallelized by applying the Fisher- Yates shuffle batch of identity permutations and forcing the draws to equal elements FY, Encoding requires inverting the Fisher- Yates shuffle by deducing which sequence of draws resulted in the observed permutation. An algorithm to do is provided by Kunze al. (2024b) in Appendix C.I, which can be easily made to work in batch. Code to run Fisher- Yates and Sattolo's algorithm is given in Section D.2
3.2 BOUTILIER 2014)
GENERALIZED REPEATED INSERTION MODEL (DOIGNON ET AL.. 2004 LU &
The repeated insertion model (RIM) DDignon et al. 2004) is a probability distribution over per- mutations that makes use of an alternative representation to inline, called insertion- The insertion- vector ((X) defines a generative process for X, relative some reference permutation X To generate X given X and V(X), we traverse the reference from left to right and insert the i-th element of X at slot V(X), E 1)). See Figure? 2 for an example.
RIM uses conditional distribution that is independent of Ves to define the joint over the insertion- vector, i.e., Pv, Ver.N = Pv, | while the Generalized RIM (GRIM) (Lu & Boutilier) 2014) uses full conditional. GRIM can be used to learn probability distributions over permutations conditioned on an observed uub-permutation. For example, for n 4 and an observed sub-permutation [2, 1, 4), set X,et [2. 1, 4, 3] such that conditional probabilities Pvş Vea. can be learned for all permutations agreeing with the observations, i.e.,
V4= = 0
V4- = 1
V4= = 2
V4 =
[3. 2, 1,4
[2, 3, 1,
[2, 1, 3,
[2, 1, 4, 3).
4

---

Published as a conference paper at ICLR 2026
X1
X2
Xs
===3
3
5
4
=
4,5m3
9,4-2
4,1m0
===1
=4=1
2
3
2
0
=
L
L2
La
L4
Ls
Figure 3: Illustration of the right-Lehmer code for permutation x= 28,,,4,1,7 (Left) Each L(X); L, counts the number of elements to the right of X; that smaller than it. (Right) Lehmer code interpreted as sampling without replacement indices.
Note this is not possible with inline, Lehmer, or the Fisher- Yates representations. The same can be achieved if the initial elements in xer are permuted, as long as the values for Ves are changed accordingly, which highlights an invariance a model over insertion-vectore must learn.
Lu Boutilier (2014) the authors use the insertion- vector representation to model user preference data, where the observed aub-permutation represents partial ranking establishing the preference of some user over a fixed set of items. In Section 5.3 we tackle a similar problem on the MovieLens dataset (Harper & Konstan 2015) where we rank a set of movies according to observed user ratings.
4 LEARNING FACTORIZED DISTRIBUTIONS OVER PERMUTATIONS
This section discusses the main methodological contri- bution of this work. MLMs can trade off compute and expressivity by sampling multiple permutation elements with one network function evaluation (or forward pass). In that case, simultaneously sampled elements are condi- tionally independent, which corresponds to an effective loss in modeling capacity. We begin by showing that per- mutations modeled the inline representation suffer most from the degradation of model capacity the number of function evaluations (NFEs) decreases, and can only model delta functions when restricted single NFE. We propose learning in the 3 alternative representations discussed Scction Lehmer codes, Fisher- Yates draws, and Insertion-vectors; which do not suffer the same degra- dation in capacity Note that while these alternative representations also have constraints for the domain of their elements, these constraints are trivially learned by the neural network it only sees valid permutations during training and can infer the domain by setting the appropriate logits to negative infinity. We show the learned conditional distributions defined by these representations are highly interpretable and subsume well known families such as Mallow's 's model (Mallows) 1957) and RIM (Doignon et al. 22004).
apply
f
4
3
3
M
C dloopgot foe caurted sobeur
IMtaaan nar diree
Figure 4: Training our method with MLL during training with the inline no- tation. For other representations, only the blue tokens change.
4.1 MODELLING CAPACITY OF P FOR THE INLINE REPRESENTATION
The masked models considered in this work are of the form.
p(s)
II
Pxs, 1X8e. = II II Px, [Xsey" JES.
x
where S (S .is......,. ..) forms a partitioning of [n), and the number of neural function evaluations (NFEs) is equal to k. Elements sampled independently if their indices belong to the same set
Sj. when conditioned on previous elements Xs The choice of NFEs restrict p(S) to a different family of models through different choices of partitioning S. For example, when limited to 1 NFE, the model is fully-factorized with S = [n]. AR minimizes at full NFEs (i.e., n = k) with S; (1). while MLM places distribution on the partitionings S resulting in mixture model.
5

---

Published as a conference paper at ICLR 2026
We consider the problem of learning distributions over valid permutations by minimizing the cross- entropy,
min E log p subject to (z) 0 if is not valid permutation,
xkx
where the expectation is taken over the data distribution.
Previous works have considered modelling permutations the inline notation where X; can take on any value in [ra.. To produce only valid permutations, it is necessary and sufficient for the support of Pxy1 Nezs to not overlap with that of another index in S, Sc: = Sgi. We can obtain an upper bound on the entropy of any inline model by considering the case when all indices in S; are deterministic except for some j # j', which is uniformly distributed over the remaining candidate indices. Formally, H (Px,) (522) = 0 and (Px, (5e) log(n |S 1). This implies the following for all j St.
H (
Σ log (n 18 1).
(4)
Equation (4) shows the modelling capacity is severely impacted by the number of NFEs. Most importantly: any inline model respecting the constraint in Equation (3) can only represent
a delta function in the case of NFE (i.e., S = [n]), as (( < 0 implies (() =
(Cover 1999). In practice, this manifests sampling time where the model fails to produce valid
permutations as in Section 5.2 At full NFEs the right-hand side of Equation (4) equals log (n!), and
achievable when p) is uniform distribution.
4.2 FACTORIZED REPRESENTATION FOR PERMUTATIONS
Next, we consider learning distributions over permuta- tions with the factorized representations discussed in Sec- tion These representations have different supports for their sequence-elements and allow values to overlap while still producing valid permutations, implying they don't suffer from the representation capacity issuc discussed in Section 4. At full NFEs, these representations can model arbitrary distributions over permutations, while at a single learn non-trivial distributions such the Mallow's model and RIM; contrast to inline which can only represent delta distribution. For this reason, we refer to them as factorized representations.
M
M
M
M
M
M
pertoee
(adi
boold
NFE they
(itscon testerte
Luphs
lught
light
Light
Ayga)
(ight
+
Lehmer Codes. We consider models p(S) over the (right) Lehmer code as defined in Section 3.1 and illus- trated in Figure 3 Left-to-right unmasking of Lehmer code can be interpreted as the sampling without replace- ment (SWOR) indices of its corresponding permutation, as illustrated in Figures and TO In the AR setting, our model subsumes Mallow's weighted model (Mallows) 1957) over the remaining elements (those that have not yet been sampled).
Hend
2
0
Figure 5: Our method during inference in the inline notation for sequence length L = 6 and NFEs = 3. For other repre- sentations, only the blue tokens change.
Remark 4.1. The weighted Mallow's model with weights Wy and dispersion coefficient ø is recovered when PL₃ (ee, (L, lei) x مi for all j E S.. This follows directly from,
PLs, Lee (es, lea) = II PL. (e, x oees 4544
(5)
JES;
where Ejes, wy Ej is the weighted Kendall's distance (Kendall 1938). In particular, when fulli-factorized, it can recover the weighted Mallow's model over the full permutation.

---

Published as a conference paper at ICLR 2026
Kusher-Yates We define the Fisher-Yates code (X) of some permutation X the sequence of draws of the Fisher- Yates shuffle that produces X starting from the identity permutation. For MLM and AR, unmasking in the Fisher- Yates representation corresponds to applying random transpositions to the inline notation. Similar to Lehmer, this can also be viewed as SWOR, except that the list of remaining elements (faded and bright yellow in Figure is kept contiguous by placing the element at the current pointer (bright yellow in Figure 2 in the gap created from sampling.
Desertion-Vectors We train using the insertion-vector representation to define conditional dis- tributions over ttb--ermutations. Similar how Lehmer recover Mallow's weighted model, conditionals can define a RIM (Doignon et al. 2004) over permutations compatible with the currently observed sub-permutation.
Remark 4.2. RIM is subsumed by model when the insertion probabilities are independent of ordering between currently observed elements, i.e., PVs. Ver.
PVay
Xwr
For Lehmer and Fisher- Yates representations there exist efficient algorithms to convert from (encode) and to (decode) inline, but it is not obvious how do so for insertion- vectors. The following theorem allows for an efficient batched algorithm for encoding and decoding, by leveraging known algorithms for Lehmer codes (see Section [DI]..
Theorem 4.3. Let L(X) be the kth element of the left- code, X- 11 the inverse permutation,
and ((X), ): the kth element of the insertion vector of X. Then,
V(X)a = - -----X-X- 1(X-))i
The proof follows from the repeated insertion procedure sampling, without replacement, the positions in which to insert values in the permutation. A full proof is given in Section B.2 Code to encode
and decode between inline and the insertion-vector representation is given in Section D.3 general theorem was proven in Azpeitia al. (2025)
A more
5 EXPERIMENTS
This section discusses experiments with factorized representations, as well as inline, across different losses. We explore experimental settings. First, common baseline of solving jigsaw puzzles of varying sizes, where the target distribution is delta function on the permutation that solves the puzzle. We then propose new settings with more complex target distributions: learning uniform distributions over cyclic permutations, as well as re-ranking movies based on observed user preference. For MLM low NFEs each set in of size n/ n/NFEs (rounded), with the exception of the last set. Shyper-parameten for all experiments are given in Section F An illustration of training given Figure 4 and inference Figure 5
5.1 SOLVING JIGSAW PUZZLES.
We evaluate models the common benchmark of CIFAR-10 jigsaw puzzles using the exact same setup as in Zhang et al. (024). Experimental details are given in Section E For MLM, we use the same architecture Ditf) as Zhang et al. (2024), with the CNN backbone conditioning on the jigsaw tensor. For AR, we modify the architecture to add an additional step that attends to the input sequence as well as the tensor (see Section D.4). All models have roughly 3 million parameters.
Our method significantly outperforms previous diffusion and convex-relaxate representations and losses. Results are shown in Figure 6 MLM can solve the puzzle with NFE (i.e., 1 forward-pan -pass) as the target distribution is delta on the solution, conditioned on the puzzle.
baselines, with all
5.2 LEARNING A UNIFORM DISTRIBUTION OVER CYCLIC PERMUTATIONS
The jigsaw experiment is limited evaluating the complexity of distributions over permutations, as the target is delta function. In this section we propose a new benchmark where the target distribution is uniform over all (n 1)! cyclic permutations of length n = 10.
All cyclic permutations of length n are generated with Sattolo's 's algorithm (Sattolo 1986). and random set of 20% are taken as the training set, resulting in a train set size of (n - 1)./5. Results are shown in Figure 7 where each point represents 10, 000 samples. All models learn fully generalize
7

---

Published as a conference paper at ICLR 2026
2x2
3x3
4x4
%6 Correct
42,58
34,69 25.31
Fisher-Yates
Inline
Insertion
Lehmer
Figure 6: Percentage of CIFAR-10 jigsaw puzzles (test set) correctly reassembled for varying puzzle size, methods, and permutation representation (higher better). SymDiff (Zhang a (2024) and Cumbel-Sinkhorn (Mena al. 2018) significantly under-perform puzzle size increases, while our methods do not. Numbers over SymDiff and Gumbel-Sinkhorn indicate their values on the y-axis, which fall below the plotted range. MLM outperforms AR by wide margin, even while using only 1 NFE.
% Unique
Unique and Valid
96 Liiiqua, Valid, and Cyclic
100
Fisher-Yazee (MLL)
Infine (MLM)
Percentage
Insertion Vector (MLM)
50
(MLMM)
repr.)
10
2
10
neural function evaluations (NFEs)
Figure 7: Performance on cyclic generating task as a function of NFEs (i.e., forward passes), across different representations and losses (higher is better). Each point contains information regarding 10k samples. (Left) Percentage of unique output sequences, including invalid permutations. All representations achieve 100%. (Middle) Percentage of simultaneously unique and valid permutations. Except for Inline, all representations achieve 100%. (Right) Percentage of unique, valid, and cyclic permutations. See discussion in Section 5.2
in the following sense: out of the 10, 000 samples taken, around 20% are in the training set, while the rest are not. All factorized representations can produce valid permutations, even the number of NFEs decreases, including for the fully-factorized case of NFE. Inline suffers produce valid permutations as discussed in Section 4.1 All methods can fully model the target distribution at full NFEs, including inline representations (right-most plot). Both Lehmer and Insertion- Vector representations still produce some cyclic permutations (above the (n 1)!/n! 0.1 bascline) even at NFE. Fisher- Yates can perfectly model the target distribution for any number of NFEs. This is expected, hinted by Sattolo's algorithm: necessary and sufficient condition to generate cyclic permutations in the Fisher- Yates representation is for FY; > 0, as these represent pass the subset of cyclic permutations. For example,
draw. The model produces uniform distribution Lehmer at 5 NFEs has non-zero mass on only 46.1% of the (n - 1)!/n! cyclic permutations. Within those 46.1 %, the probabilities are uniformly distributed, while the remaining have mass.
5.3 RE-RANKING ON MOVIELENS
Our last experiment is concerned with learning distributions over rankings of size n, conditioned on existing user preference data the MovieLens32M dataset (Harper & Konstan 2015). MovieLens contains 32 million ratings across 87, 585 movies by 200, 948 users
0.5 scale from 0.5 to 5.0.
We first filter to keep only movies rated by at least 1, 000 users, and then randomly sample 1, 000 movies from the remaining. Only users that rated at least movies out of the 000 sampled movies are kept. In the smallest setting (n = 50), the dataset totals roughly 18 million ratings across 174 thousand users. The dataset split users into 80% train and 20% validation.
8

---

Published as a conference paper at ICLR 2026
-#
AR
MLM (1 NFE)
Most Watched
RIM (Uniform)
k=2
k=3
k=6
k=9
k=10
NDCG@K
95
90
85
80
5
1
5
5
1
5
5
k=2
k=6
k=12
k=18
k-20
NDCG@K
90
85
5
10
1
5
10
1
5
10
1
5
10
1
5
10
k=2
k=16
k=32
k=48
k=50
NDCG@K
90
7=50
80
5 10
25
1 5
10
25
1 5
10
25 5 10
25
1
10
25
observed rank size
Figure 8: Results for re-ranking conditioned user ratings in MovieLens (higher is better) for varying rank sizes n. See Section 5.3 for a full discussion of the results.
Note that the only information available to the models in this paper are rankings of previous liked items, with no notion of user, user features, even item features. There no guarantee of single "true ranking" of size when conditioned on sub-ranking of size k < as there will likely be two users that have the same preference on a subset movies, but differ in preference when looking at the full ranking of size (i.e., the target is uniform distribution over these rankings of size ni).
During training, we sample n ratings (each for a different movie) from each user. The (shuffled) sequence of n movie ids make up Xyet. The user ratings then used to compute the true ranking (i.e., labels), with ties broken randomly. The input sequence of size 2n, with the first n corresponding to the movie labels (i.e., X prefix), and the last the true ranking in the insertion representation (i.e., V(X), ) labels). We train with MLM and AR to predict the labels conditioned on the prefix, and the labels generated far (i.e., conventional cross-entropy training, or "teacher- forcing").
To evaluate, we sample ratings for each user in the test set (as done in training) and condition the first few movies Vey to predict the remaining Vz,. Note this is possible without training separate conditional models, because the GRIM representation allows us learn all conditionals of the form Pv, Ver. when training with the AR and MLM objectives.
We compare against two baselines: ranking movies by number of users that watched them, and RIM (Doignon et al. 2004) with uniform insertion probabilities; conditioned on the observed ranking Ver- Results are shown in Figure for the NDCG@k k metric (Järvelin & Kekäläinen 2002). NDCG@ @ measures the agreement to the true user ratings, and has maximum value of 1.0. Note that NDCG@* k similar to cross-entropy when the relevance scores are the normalized log- bog-probbbilitess (which case), which is an appropriate metric for a distribution learning task.
AR (I Pv, |V2,)) and MLM (1 NFE, 11 Pv, Iv..) perform similarly, and outperform both basclines in all settings. Note r = and r = 0 are equivalent, as V(X), = 0 with probability 1. The conditional MI.M model at NFE is different from the unconditional MLL model at NFE

---

Published as a conference paper at ICLR 2026
U, Pvş): which is why performance improves a function of the observed rank size r. In this setting, the AR baseline is very strong baseline, which should have very high performance on this task, given that no semantic content information is available to take advantage of.
DISCUSSION AND FUTURE WORK
We present models capable learning arbitrary probability distributions over permutations via alternative representations: Lehmer codes, Risher-Yate, draws, and insertion vectors. These rep- resentations enable unconstrained learning and ensure that all outputs are valid permutations. We train models using auto-regun retrrssegees and masked language modeling techniques, which allow for a trade-off between computational cost and model expressivity. Our approaches achieve state-of- the-art performance the jigsaw puzzle benchmark. However, we also argue this benchmark is insufficient to test permutation-distribuiie introduce two new benchmarks that require learning non-trivial distributions. Lastly, we establish novel connection between Lehmer codes and insertion vectors to enable parallelized decoding from insertion representations.
modelling as the target is deterministic. Therefore, we
The methods in this work explore learning distributions over permutations, where the set of items to be ranked is already known before-hand. An interesting avenue for future work is to model the of items simultaneously. as is the case in real-world recommender systems. Experiments on MovieLens hint the scaling capabilities of these factorized representations beyond simple toy settings, as the size of learned permutations for non-trivial experiments in previous literature has generally been much smaller than that explored theoretical standpoint there is room for more characterization of the properties of these familiis of distributions the low NFE setting.
our largest MovieLens experiment (n = 50). Finally, from
REFERENCES
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2021. URL fictps://arxiz/
abs/
2107. 03006
Mikel Malagón Azpeitia, Aimar Barrena, Hugo Federico fñigo Arroyo, Josu Ceberio Uribe, Ekhine Irurozki, and José Antonio Lozano Alonso. Una visión unificada de transformaciones biyectivas en la optimización de problemas de permutaciones. In Actas del XVI Congreso Español de Meta- heurísticus, Algoritmos Evolutivos y Bioinspirudox (MAEB 2025) 28-30 de mayo, Donostia/San Sebastión, pp. 169-178. Servicio Editorial= Argitalpen Zerbitzua, 2025.
Christopher JC Burges. From ranknet to lambdarank lambdamart: An overview. Learning, 11 (23-581):81, 2010.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Douglas E Critchlow, Michael A Fligner, and Joseph S Verducci. Probability models on rankings. Journal of mathematical psychology, 35(3):294-318, 1991.
Marco Cuturi, Olivier Teboul, and Jean-Phil Vert. Differentiable ranks and sorting using optimal transport, 2019.
11885
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. 2019, URL abs/1810. 04805
Ming-Wei
Aïssatou Diallo, Markus Zopf, and Johannes Fürnkranz. Permutation learning via lehmer codes, In ECAI 2020, pp. 1002-1102. IOS Press, 2020.
Jean-Paul Doignon, Aleksandar Pekeč, and Michel Regenwetter. The repeated insertion model for rankings: Missing link between two subset choice models. Psychometrika, 69(1)333-54, 2004,
Yufei Feng, Yu Gong, Fei Sun, Junfeng Ge, and Wenwu Ou. Revisit recommender system in the permutation prospective. arXiv preprint arXiv: 102. /2057, 2021.
10

---

Published as a conference paper at ICLR 2026
Ronald Aylmer Fisher and Frank Yates. Statistical tables for biological, agricultural and medical research. Hafner Publishing Company, 1953.
E. N. Gilbert. Theory of shuffling. Technical memorandum, Bell Laboratories, 1955.
Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic optimization of sorting networks via continuous relaxations, 2019,
Maxwell Harper and Joseph Konstan. The movielens datasets: History and context. Acm transactions on interactive intelligent systems (tiis, 5(4):1-19, 2015.
Zhengfu He, Tianxiang Sun, Kuanning Wang, Xuanjing Huang, and Xipeng Qiu. Diffusionbert:
Improving generative masked language models with diffusion models, 2022. URL https
//arxiv.org/aba/2221..
15029
Kalervo Järvelin and Jaana Kekäläinen. Cumulated gain-based evaluation of ir techniques. ACM
Transactions on Information Systems (TOIS), 20(4):422 446, 2002.
Jorge K. S. Kamassury, Henrique Pickler, Filipe R. Cordeiro, and Danilo Silva. Cct: A cyclic co- teaching approach to train deep neural networks with noisy labels. IEEE Access, 13:43843-43860, 2025. doi: 10. 109//CCESSSS..// 10.1109/CCESSS 2025.3555222255
Maurice G Kendall. A new measure of rank correlation. Biometrika, 30(1-2):81-93, 1938.
Jungtaek Kim, Jeongbeen Yoon, and Minsu Cho. Generalized neural sorting networks with error-free differentiable swap functions, 2024. htt.hstps////azii.rrrrr///// 23tp://aaaa/3110111100011 07174
Ouail Kitouni, Niklas Nolte, James Hensman, and Bhaskar Mitra. Disk: A diffusion model for structured knowledge, 2024. htL/httss:// http://axvvvv//v/ org. abs 23122232...0 05253
Julius Kunze, Daniel Severo, Jan Mil-Willem van Meent, and James Townsend. Practical shuffle coding. In The Thihty- Anmual Conference Neural Information Processing Systems, 2024a. URL
net. forum? dddmmmdammmmma
Julius Kunze, Daniel Severo, Giulio Zani, Jan-Willen van de Meent, and James Townsend. En tropy coding of unordered data structures. In The Twelfth International Conference on Learning Representations, 2024b. URL http:::s///pee.....eee...ccc http://wwweeeeeeeeeeerrrrraaaaaa 3Ruh
Hugo Larochelle and lain Murray, The neural autoregressive distribution estimator. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pp. 29 29-37. JMLR Workshop and Conference Proceedings, 2011.
Derrick H Lehmer. Teaching combinatorial tricks to computer. Combinatorial Analysis, pp. 179-193, 1960.
Tyler Lu and Craig Boutilier. Effective sampling and learning for mallows models with pairwise- preference data. J. Mach. Learn. Res., 15(1):378338829, 2014.
R Duncan Luce et al. Individual choice behavior, volume 4. Wiley New York, 1959.
Mikel Malagón, Ekhine Irurozki, and Josu Ceberio. A combinatorial optimization framework for probability- algorithms by means of generative models. ACM Transactions on Evolutionary Learning, 4(3):1-28, 2025.
Colin L. Mallows. Non-null ranking models. i. Biometrika, 44(1/2):114 1957.
John Marden. Analyzing and Modeling Rank Data. 01 2014. ISBN 9780429192494, doi: 10.1201/ b16552.
Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permutations with gumbel-sinkhon networks. arXiv preprint arXiv:1802.08666 2018.
Meta. The llaaa herd of models, 2024. URL htt.:/pa.//aavii.rrr//abass///se

---

Published as a conference paper at ICLR 2026
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin. Ji-Rong Wen, and Chongxuan Li. Large language diffusion models, 2025. URL https://
arxiv.org/abs/ wtt.:/orr/abb//22...000 09992
Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Monotonic differentiable sorting networks, 2022. URL https:/// http::://axxvv.rro/// org/ abs/ http//a...//222.//////// 221///2233.000
Robin L Plackett. The analysis of permutations. Journal of the Royal Statistical Society Series C: Applied Statisties, 24(2):193-200 1975.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and lly Sutskever. Language models are unsupervised multitask learners. 2019.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models, 2024. URL https::/ //tp:vv//a.ivv.a/aaa wwtp/vv..rg/aa.oo abs 2406. 07524
Sandra Sattolo. algorithm to generate random cyclic permutation. Information processing letters, 22(6):31--317, 1986.
Claude E Shannon. A mathematical theory of communication. The Bell system technical journal, 27 (3):379-423, 1948.
Jiaxin Shi, Kehang Han, Zhe Wang, Araud Doucet, and Michalis K. Titsias, Simplified and generalized masked diffusion for discrete data, 2025. URL fttps://arxiv.org//a
2446604322
Benigno Uria, Marc-Alexandre Côté, Karol Gregor, lain Murray, and Hugo Larochelle. Neural autoregressive distribution estimation. Journal of Machine Learning Research, 17(205):1-37, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1706.03766 2017. URL http://arxiv.cr
Yongxing Zhang, Donglin Yang, and Renjie Liao. Symmetricdiffusers:: Learning discrete diffusion on finite symmetric groups. arXiv preprint arXiv:2410.02942, 2024.
12

---

Published as a conference paper at ICLR 2026
A BACKGROUND
A.1 PERMUTATIONS
permutation in this context is a sequence X of elements X; E [n]] such that U, = [n] with X having no repeating elements. Permutations are often expressed in inline notation, such as X = [5, 4, 1, 2, 3]. A permutation can also be seen as bijection X [n] + [n]. where ( X; is the element in the inline notation at position i.
transposition is permutation that swaps exactly 2 elements, such as X = 21,2,4,3.
cycle of permutation is the of values resulting from repeatedly applying the permutation, starting from some value. For the previous example, the cycles are (1 - 5 + 3 + 1) and Figure
2). A cyclic permutation a permutation that has only cycle, example is given in
The inverse X, denoted x-1, the permutation such that (()- (x - (x-))x())
A seb-eermttation of permutation X of length is sequence of m S n elements 2, = X that
agrees with X in the ordering of its elements, i.c., 81 < i2 < im. For example, [5, 1, and [4, 1, are sub-permutations of [5, 4, 1, 3], but [4, 3, 2] is not.
See Marden (2014); Critchlow al. (1991 for more complete introduction permutations and ranking models.
B THEOREMS AND PROOFS
B. NEIGHBORING LEHMER CODES DIFFER BY TRANSPOSITION
The following theorem gives metric-space interpretation for Lehmer codes, and how changes in L(X) affect X.
Theorem B.1. For any two permutations X, x', ( - L(X'Il = then X and X' are equal up transposition.
The proof follows from analyzing the list of remaining elements at each SWOR step, and can be seen from a simple example. Consider following Lehmer codes L, L' differing only at L'g = Lg + 1, their SWOR processes, and their resulting permutations X, X'.
4
X,
Lg
X₃
La
4
Xa
L4
4
X4
La
Xa
Note the following facts:
1. transposing 3 and 1 in the initial permutation (first row) and applying the SWOR process of
L results X";
2. the element chosen at step by L3 is adjacent in the list to the element chosen by L/g, as La Lgl
3. steps before 3 unaffected,
their respective inline elements;
4. steps after 3 unaffected, as long as the sampled index does fall in either of the two blocks corresponding to L3 and La (where a change occurred).
In general, for an increment at position j. the only affected elements are those at Lj and L, + 1, implying X and X' differ exactly by the transposition of these elements.
13

---

Published as a conference paper at ICLR 2026
A more general statement can be given for the case of increments beyond 1. Consider L; L, k. All future steps i with elements L, E [L L, + k] are affected, requiring permutation of size + 1 to recover X.
Theorem B.2. For any two permutations X, such that L(X); = L(X"), all i + then X and X' are equal up to a permutation of [[(); L(X', + elements.
B.2 THEOREM
Restating Theorem Let T(X) be the kih element of the left-Lehmer code, X-1 the inverse permutation, and V (XX) Du the kth element of the insertion vector of X. Then,
v(X) = 1-(((X11)
First, let pa be the position of the value k in X, i.c. Xpp By definition of inversion, ph: X Then, note ((X)+ = | < Px, < KJ.. In words; The insertion vector element V(X)* counts the number of elements to the left of the position of value h in X (i.e. pt) that are smaller than k. This can be seen by the following argument: By definition, an insertion vector element V(X); describes in which index to insert an element with the current value & (or + 1, depending on indexing definitions), see Figure (right). Because all previously inserted values are smaller than k and all values inserted later will be larger, the index at the time of insertion is equal to the count of smaller elements to the left of the final position of value & in X, which is Pk-
Recall the definition of the left Lehmer code: L(X)x = H X > X
Define L'(X): = k - L(X)x and notice that
(' k= = -(X = - 16/ < <|x, > X <
since K6j < k3} and x X fj k
Insert the inverse permutation X-1:
L'(X = H < k|xj'
Xx }= 1 < klp, < p
Next, perform a change of variable on j in V(X)
v(X) = 1i < pa XX < k = Hp: < p&ll < k}
where 1= + +
Comparing,
k: - ='(x--)
Köli < k,p, < <PH 10pp < lA, < = V(X)
C LIMITATIONS
The most important limitation of this work is scalability large permutations. A loose bound can be estimated by realizing that we model the permutations with transformer architectures. Therefore, the memory and compute required to train on tasks that require large permutations are quadratic. In particular, common methods in ranking include score functions, which can act on each item individually to produce a score, rather than needing to condition on all items as we do.
In general, since the search space of permutations grows much quicker with length (n!), the scalability often not dominated by memory requirements search is required, rather by the compute needed for the search.
An inherent limitation of the method is that n forward passes through the network are needed to achieve full expressivity over the space of permutations of length n. This is consequence of MLM and AR training, resulting in token-wiss factorized conditional distributions. This is detailed in Section 4.
14

---

Published as a conference paper at ICLR 2026
D CODE
D.1 LEHMER ENCODE AND DECODE
In practice, our left-Lehmer encoding maps an inline permutation L' from Equation interacts more directly with the insertion vector.
because
def 1ehmer Sechee_eencodeeeeeeeeeeeeeeeeeee (perm: Tensor, left: bool = False) -> Tensor: clone ())
lehmer atrch.
n Schmer.size:/ if left:
for i in reversed (range (1, n)
lehmer :i] - (lehmer [ [i]] <= lehmer :i]) .to tinintt)
else:
for range (1, n): Lehnerl:, i:] (lehmer(:) (;, [i 1]] Lehmer [;, i:]) .to (Int)
10
if Ten shape) -- 1:
lehmer lehmer squeeze () elif len (perm.
13
14
lehmer torch. at least 2d (lehmer)
return lehmer
17
18
19
def lchmer_decode (lehmer: Tensor, left: bool False) -> Tensor:
perm = atleast_tlest_d (Lehmer.cconee (-rm.
22
for range (1, n):
23
if left:
24
perali, :i] += (perml:, [i]] <= permi:, (i]).to(ine)
else:
j
27
[t j :] (pernl:, [i]] <- perm[:, 3 +
{]) to (int))
28
if len (lehmer. shape)
30
perm = perm. squeeze () elif len (lehmer. shape) = 2:
31
32
perm = torch. at atlasst_2dd (porm)
perm
D.2 FISHER- YATES ENCODE AND DECODE
def
sester_yatts_eeeee (perm: terch. Tensor) torch, Tensor: len (perm.shape)
priginal_nuum__
perm terch. ateasttteettt_ (perm)
B, peam.
perm base torch.arange (n) unsqueeze (0) repeat ((2,
lt) (term.deevice)))
fisher_yates = torch. zeros_liiie iik (perm) to (perm. device) batch_idx torch. (B) to peermdevicce.
for i in range (n):
10
j = nonzero torch.nonzero (perm[:, [i]] == perm base, as_tuple-True) [1] fisher fbatch_ixx, i]
13
icx torch. stack (Itorch. full_like (1, 1), 11. dim=1)
14
values pase__aase.geeeeeeeeppeeeeeeeees gather idx)
15
swapped_values = torch. tvalues, [1])
perm_base, scatter_ idx, swapped_va values)
15

---

Published as a conference paper at ICLR 2026
if original num dims
19
fisher_yates sishe_yytees squeeze
20
clif original_num_dinee =- 2:
fisher_yates = t at atcastt_actttaaa hat..tttlos_talflhhheeettttr
fisher_yates
24
25
def fisher_yates_ (fisher_yates:: Tensor) -> Tensor:
26
B, = =isher_yatee shape
perm torch. (n) unsqueeze (0) repeat (B,
1)) (fisher____esss cev
28
batch_idx arrne (B) to satteer__yttssdeiicee
29
for i in range (n)
= fisher__yyye l; il
perm[batch idx, il, perm!:, permi:, il, perm idx, j1 return perm
32
D.3 INSERTION- VECTOR ENCODE AND DECODE
def invert_permpperm. Tensor)
Tensor:
return torch, argsort (perm)
def
corrcoodd__erhhhhhoo (perm: Tensor)
Tensor:
inv_perm invert_perm (perm)
insert_ Lehmer encode torch (inv return insert_
pirypper,,, perm, left=True)
10 def rnoortoonneecorrdddooooooot (insert_v: Tensor) Tensor:
11
inv_perm = 1ehmer_decode_ (incoeeertrtccnr___..... left=True)
12
perm = invernt_eeeeee
return perm
D.4 MODIFIED SYMDIFF-AR
We modify the following function in
https://gtthub..com/Sttt-- himmmettsccllffuuss/////to 6eaf htt.://.........2aaaa2a.22666665
models. Modelss.prPL33577-1333222
The first N elements of embd correspond to the embeddings of the puzzle pieces computed with the
CNN backbone, while the following N are the token embeddings of the input. The attention mask (embd_attn_ne) guarantees all tokens attend the puzzle pieces, but the inputs can be attended to causally pe permattn.maskk is causal, AR case) or fully (MLM).
def apply_layers_seie
self, embd, cime_embd, attn attn_mask-None, perm_att__ask-llone perm_embd-None
smbd.sszee (1)
time embd time emod. unsqueeze (-2) embo - embd + time_embb
embó_ embó_attn_mask None if perm_embd is not None:
10
embd = torch. cat ([embd, perm_ samm_embd], dim=1) embd = self, self.perm caarr__os_enccdeee (embd)
13
if perm_attn_mask
anbd_attn_mask not Nene:
14
15
torch. zeros ( (2 N, 2 wtw.atttm_maadeiiiieaeeee
N) to (bool) to (pern_attn_mask
16

---

Published as a conference paper at ICLR 2026
embd attn mask ;, :N] True
18
emboL_atnn_mask [N:, NJ perm_atth__ass embd_attn_mask "embd_atto_mask
19
20
for layer in self.encoder__ayeee embo layer (embd, ar__mak-errd_attn_msst
23
24
return embdl:, N : N!
E EXPERIMENTS
E. JIGSAW EXPERIMENTS
Each CIFAR-10 10 image partitioned into a jigsaw puzzle in grid-like fashion. The pieces are scrambled by applying permutation sampled uniformly in the symmetric group. This produces tensor of shape [B, N H/N, W/N), where per dimension) and H and W the original image dimensions (i.e. H W = 32 for CIFAR-10). The images are cropped at the edges if H and W are not divisible by N, as in Zhang et al. (024).
the batch dimension, N the puzzle size (specified
Hyperparameters:
1. learning rate = 3 x 10
2. batch size 1024
3. Model configurations
follow
those
in
https://
[a2-3332888778800000000332777622 inscramble-CIFAs
E.2 CYCLIC EXPERIMENTS
1. learning rate 3 10-4
2. batch size 1024
3. DiT model size:
(a) hidden dimension size 128
(b) number transformer heads 8
(c) time embedding dimension
(d) dropout = 0.05
(e) number transformer layers =
E.3 RERANKING MOVIELENS
1. learning rate = x 10-4
2. batch size 1024
3. DIT model size:
(a) hidden dimension size 256
(b) number of transformer heads
(c) time embedding dimension
(d) dropout 0.05
(e) number of transformer layers 10
F COMPUTE
Our experiments were nodes with single NVidia A-100 GPU. Since the models trained of small scale, no experiment took longer than 2 days to converge. In total, an estimated 10000 GPU hours spent for the research for this paper.
17

---

Published as a conference paper at ICLR 2026
G IMPACT STATEMENT
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
H EXTRA FIGURES
3
1
(a) Cyclic: म = 3, (1 + 3 + + 4 1)
(b) Non eyelicelic: π 5, 1, (1 - 2 + 1), (3 + 4)
Figure 9: Illustration of cyclic vs. non-cyelic permutation.
shuffied
Insertiono
Currnn sub ralinaniigg
M
M
M
BUOOO M,M.M.M
O,M.M,M
0
M
M
58,6b, 1,M,M]
M
condissn Moroes and subbaakking
,, 0, 1, 0,M]]
Code:
8,,6b, 0, 0,
Figure 10: (Left) Decoding lehmer code from left to right represents sampling without replacement. Illustrated on Jigsaw puzzles. (Right) Prediction task on the MovieLens dataset. Insertion-vectonn allow us to define conditionals over sub-rankings corresponding to user preference data.
18


Reviews:
[{"review_id": "6i77Oru14o", "reviewer": "Reviewer_RZJL", "summary": "Defining distributions over permutations is a useful machine learning primitive for applications that involve ranking, assigning, etc.\n\nThis paper provides new, interesting ways to represent distributions over permutations by leveraging a variety of classic methods for encoding permutations (e.g., using swaps, insertions, etc). Each of these can be coupled with a language model (masked, autoregressive) for defining a distribution over permutations.\n\nThe representational capacity of various representations is discussed and experiments on both synthetic and real data are used to compare performance.", "questions": "I found the MLM vs. AR exposition confusing. It took me a while to realize that things are sampled independently within a masked region. It may have been helpful to define MLM\n\nI found the setup for experiment 5.2 (cylic permutations) confusing. In what sense are you assessing whether the uniform distribution was learned? Can't these models be used to assess the likelihood of a given permutation? Could you check that the likelihood scores are approximately uniform (suggesting that the distribution is approximately uniform)?\n\nIn sec 5.3 (MovieLens), why were only certain representation approaches used? I would have expected to see a comparison of all the approaches in Figs 5 and 6. Also, MovieLens has been used for benchmarking a wide variety of learning-to-rank models. I don't expect you to necessarily get SOTA performance, but it would have been useful to understand roughly how well your methods compare.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper invokes a number of interesting, classical formulations of permutations that I was not familiar with. I learned a lot from reading the paper. \n\nThe paper is a good example of how to combine classic algorithmic formulations with neural networks to get novel/flexible distributions over structured objects.", "weaknesses": "I found the  (num function evaluations) NFE formalism confusing. This is used to provide a tradeoff between compute cost and expressivity of the resulting distribution. I understand the NFE=1 and NFE=k regimes, but I don't know how to interpret a value in between and I don't understand how the partitioning was chosen. I also don't know why this variable was so central in the experiments. To me, I was most interested in the different representations, not sweeping over values of NFE.\n\nI have some questions about the experiments. See below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3jR1B0ZZGE", "reviewer": "Reviewer_Cnmj", "summary": "This paper addresses distribution learning over permutations by replacing the usual inline permutation representation with three factorized representations — Lehmer codes, Fisher–Yates draws, and Insertion-Vectors. The key idea is that these representations have element-wise conditional independence, so masked language-modeling (MLM) or autoregressive (AR) training can produce valid permutations with much more expressivity at low numbers of forward passes (NFEs) than inline notation.", "questions": "1. **Scaling to larger n.**  While the proposed approach shows strong results on moderate-sized permutations (e.g., n ≤ 50), its scalability to larger sets remains unclear. It would be helpful to include preliminary experiments on larger-scale settings (e.g., n = 100 or 200) to assess both accuracy and computational cost. Prior work such as Zhang et al. (2024) has demonstrated scaling on sorting tasks; including similar experiments would make the empirical evaluation more comprehensive.\n\n2. **Representation selection and adaptability.**  The paper currently relies on manual choice of representation (e.g., Fisher–Yates for cyclic structures). It would strengthen the contribution to discuss or explore automatic strategies for selecting or combining representations—such as heuristic rules, task-dependent priors, or meta-learning approaches. Providing such guidance would make the method more practical and easier to adopt across diverse domains.\n\n3. **Clarification on non-delta distributions (MovieLens benchmark).**  \nThe paper argues that most prior benchmarks (jigsaw puzzles) effectively correspond to learning *delta distributions* over permutations, whereas the proposed MovieLens re-ranking benchmark involves learning a *non-delta* target distribution. However, the paper does not clearly explain why the MovieLens setup yields a non-degenerate distribution. Could the authors elaborate on how the stochasticity or user-item sampling procedure induces a genuine distribution rather than a single ground-truth permutation? More details or visualizations of this distributional diversity would help clarify the distinction.\n\n4. **Conditional independence under Lehmer codes.**  \nThe authors mention that Lehmer code representations exhibit conditional independence across indices, which supports efficient factorization. However, this property seems to hold only when the target permutation distribution is a delta (i.e., deterministic). For non-delta or multimodal permutation distributions, the dependencies across indices should reappear. Could the authors clarify under what assumptions the conditional independence claim holds, and whether it still approximately holds for stochastic targets such as those in the MovieLens setting?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1) **Clear, well-motivated idea.** \nReplacing inline notation with factorized bijections is simple but powerful: it directly addresses invalid outputs in fully-factorized masked models and provides a tradeoff knob between expressivity and compute.  \n2) **Theoretical insight.** \nThe paper formally characterizes why inline representations collapse to degenerate distributions under 1 NFE and proves a useful identity (Theorem 4.3) that enables efficient batched insertion-vector decoding.  \n3) **Remarkable gains on benchmarks.** \nMethod outperforms diffusion and relaxation baselines on jigsaw tasks; The paper also extend two new benchmarks: the cyclic and MovieLens experiments. Both results further prove that the choice of representation can be used to model beyond the delta distribution and to handle conditional generation tasks via  insertion-vector.", "weaknesses": "1. **Scalability to large n.**  The results of training/serving for very large rankings (e.g., thousands of items) remains unknown. Jigsaw and cyclic (n=10) are illustrative but small; MovieLens experiments are more realistic but limited to subrankings (n=50). Behavior at larger scale is untested.\n\n2. **Representation selection requires prior knowledge.**\nChoosing the best factorized representation (e.g., Fisher–Yates for cyclic structure) is sensible but introduces an extra design choice and possible brittleness when task priors are unknown.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yLL2m1WH74", "reviewer": "Reviewer_UPQC", "summary": "The authors propose three other parameterizations of permutations as opposed to the typical inline notation when modeling distributions over permutations, and perform experiments using generative models to confirm the efficacy of these parameterizations.", "questions": "Can the authors offer specific settings where one would like to learn a distribution over permutations in practice, using a MLM or any other generative model?  I feel there could be more motivation for the generative element of the paper.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The core idea seems novel and natural, as there are seemingly many ways to represent permutations with no intuitive reason to prefer the inline notation.  On the synthetic experiments the results are also quite good.", "weaknesses": "For this last task in particular, there are some issues.  The measurement is NDCG@k which as I understand it doesn’t require a distribution?  Even if one explicitly wanted a distribution over rankings for the sake of, say, uncertainty quantification, I also feel a somewhat straightforward baseline is missing from this task.  Namely, one could learn a distribution over each individual ranking of a film conditioned on a user, and then implicitly induce a distribution over rankings by sampling scores for each film and sorting them.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ORphV6hDou", "reviewer": "Reviewer_ZSYA", "summary": "In this paper, the authors approach the problem of learning probability distributions over permutations. Although this task might look trivial, the number of probability models over permutations is limited (Mallows, Plackett-Luce, Bradley-Terry,…), and each of them imposes a certain structure in the learnt distribution (strong unimodality, …). In this sense, given a sample of permutations, learning the probability distribution that best resembles the features of the samples is challenging. In this work, the authors consider the inversion-vector type representation to transform the permutations, and learn the distributions on top of using masked language models.", "questions": "Taking into account the comments above:\n\nQ1. What is the novelty of the work? Are you merely proposing to use masked language models (MLM) with alternative representations?\n\nQ2. Can you describe any properties of the masked models learned, following Critchlow 1991?\n\nQ3. Would it be possible to extend the idea of MLM as adapted from Uria 2016 in this paper to solve optimization problems?\n\nIf the authors consider the suggested comments and make changes to the paper so that it provides a much more comprehensive overview of the literature and its possibilities, I would be willing to reconsider my score.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "S1. Although the problem dates back 50 years, it is still a relevant and unsolved problem today.\n\nS2. I did not read about Fisher-Yates draws previously, and it was definitely relevant to consider as an alternative representation of permutations.\n\nS3. The paper is very well written.", "weaknesses": "W1. Many relevant references are missing. For me, it is surprising not to refer to general literature on this topic as:\n—> Marden 1995, Analyzing and Modeling Rank Data.\n—> Critchlow et al.1991, Probability models on rankings.\n—> Fligner et al.1998, Multistage ranking models.\n\nIn fact, Critchlow et al. (1991) analyze four properties of probability models on rankings (strong unimodality, L-decomposability, …) that could be interesting to consider in this paper. What can the authors say about the properties of the distributions learnt?\n\nW2. At times, the authors present the use of alternative representations as if it were novel, which it is not. In fact, Mallows' Generalized model itself uses V_j to decompose permutations (incidentally, I have not seen any mention of this model; I recommend that the authors review Fligner 1998).\n\nW3. Furthermore, there is a group of authors who have worked on the use of alternative representations to propose algorithms (some based on NN+RL) that learn probability distributions to sample solutions to permutation problems. A recent work that uses feed-forward networks to estimate probabilities on the inversion vectors is:\n—> Malagon et al. 2024, A Combinatorial Optimization Framework for Probability-Based Algorithms by Means of Generative Models. ACM TELO.\nIn fact, in this work, it is shown that one of the strong points is the possibility to infer in parallel a great number of permutations. So, contribution 5 in Section 4.2 is not entirely new.\n\nW4. Regarding Theorem 4.3, I have to say that this result has been published too; actually, even a more general outcome is proposed. See Table 1 in the work below:\n—> Malagon et al. 2025, Una Visión Unificada de Transformaciones Biyectivas en la Optimización de Problemas de Permutaciones, MAEB2025 (English version accepted for ECAI 2025).\nIn that work, the authors consider inversion vectors, Lehmer codes, and the RIM as a part of the same framework and use them with a unified notation. So, the published version is even more general than that in the present work.\n\nW5. Several times, the authors mention the idea of unconstrained learning. I think it would help if the authors first refer to the mutual exclusivity constraint that appears naturally in permutations, and then suggest that the concept disappears in the inversion vectors space. However, in some sense, inversion vectors have also \"constraints\" since the integer numbers that can appear at each position are strict. However, the authors do not mention this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)