Exploring the Design Space of Permutation Diffusion LLMs

1 Introduction

• Permutation generation orders a fixed candidate set exactly once, and language models can condition that ordering on textual descriptions and task context.Parallel prediction alone may repeat or omit candidate identifiers, so it does not ensure the required global assignment.Source: §1 · p. 1
• Hard assignment predicts candidate scores at each position and uses Hungarian assignment; soft assignment predicts doubly stochastic matrices and uses Sinkhorn normalization before terminal assignment; ordering scores are sorted into a permutation.The three parameterizations impose assignment constraints at different stages of prediction and decoding.Source: §1 · p. 2
• At one function evaluation, hard- and soft-assignment variants achieve reranking quality comparable to constrained masked diffusion, while ordering-score prediction guarantees valid permutations without Hungarian assignment or Sinkhorn normalization.The experiments use recommendation reranking across three Amazon domains with candidate sets of 10 and 20 items.Source: §1 · p. 2
• The study develops family-specific corruption processes, objectives, and samplers, including structured assignment losses and model-generated soft-state training.It compares multiple hard-assignment losses and two adapted sampling procedures for soft assignment.Source: §1 · p. 2
• The paper characterizes tradeoffs in ranking quality, permutation validity, and iterative behavior across the three diffusion-LLM families.The families distinguish state representations and validity mechanisms.Source: §1 · p. 2

Diffusion models for discrete structures

• Discrete denoising diffusion extends generative modeling to categorical states, while masked diffusion uses absorbing mask states and has connections to masked language modeling.DISCO applies discrete-state continuous-time diffusion to graphs and establishes permutation-equivariant and permutation-invariant properties with respect to node ordering.Source: §2 · p. 2

Cold Diffusion and training on model-generated states

• Cold Diffusion’s restoration-and-recorruption perspective motivates samplers adapted here to assignments, where hard corruption rearranges candidates and soft corruption interpolates between doubly stochastic matrices.The paper also measures stability across NFE budgets rather than assuming that extra evaluations improve quality.Source: §2 · pp. 2–3

Generative models of permutations

• Prior permutation models use shuffle-based discrete diffusion, continuous soft ranks, or flows over matrices with unit row and column sums, generally with task-specific architectures.This work instead uses a pretrained language-model backbone and compares per-step assignment projection, soft-state recurrence, and direct ordering-score prediction.Source: §2 · p. 3

Assignment relaxations and neural ordering

• Assignment methods include ListMLE for ordering scores, Sinkhorn normalization for matrix marginals, Gumbel–Sinkhorn relaxations, and Hungarian assignment for discrete optimization.The paper compares how these components impose assignment constraints during language-model sampling.Source: §2 · p. 3

Language models for reranking and recommendation

• Reranking work includes instructional and specialized listwise models, first-token ranking, constrained masked-diffusion reranking, and recommendation models with autoregressive or block-diffusion generation.The paper ranks a fixed candidate set directly, rather than generating candidates or explicit reasoning traces.Source: §2 · p. 3

3.1 Notation and Problem Formulation

• A ranking is a permutation π in the symmetric group, where π(i)=j means candidate j occupies output position i; candidate indices identify presented-list items, not ranks.The context contains user context and candidate descriptions, and the clean target is represented as a candidate-token sequence.Source: §3.1 · p. 4

3.2 Autoregressive Reranking

• The autoregressive reranker predicts candidate-index tokens left to right, requiring K sequential steps; restricting predictions to unused candidate tokens ensures each candidate appears exactly once.Unconstrained greedy decoding may repeat or omit candidates.Source: §3.2 · p. 4

3.3 Masked Diffusion Language-Model Reranking

• The candidate-score matrix has rows for output positions and columns for candidates, and each function evaluation scores all K positions in parallel.At denoising step n, the score matrix is denoted S(n).Source: §3.3 · p. 5

3.3.1 Forward masking and clean-state prediction

• The forward masking process replaces a randomly selected subset of clean ranking tokens with a mask token, ensuring at least one position is masked.The model predicts clean tokens at the masked positions using the unshifted hidden state for each position.Source: §3.3.1 · p. 5

3.3.2 Raw confidence decoding

• Raw confidence decoding begins with K masked positions and commits the most confident unresolved positions at each evaluation, but independent candidate choices may repeat candidates.Committed tokens remain visible in later evaluations, and later evaluations may reuse previously committed candidates.Source: §3.3.2 · p. 5

3.4 Assignment Operators

• Hungarian assignment selects a one-to-one candidate-position matching, while Sinkhorn normalization produces an approximately doubly stochastic soft assignment rather than a discrete permutation.Constrained masked decoding assigns distinct unused candidates to selected positions and returns a valid permutation for every supported evaluation budget.Source: §3.4 · p. 6

4 Method

• Training corrupts the target permutation while keeping the instruction prefix fixed, and sampling feeds each predicted clean state into the next model evaluation.Table 1 organizes the families’ state representations, corruption processes, objectives, samplers, and decoders.Source: §4 · p. 6, Table 1 · p. 8

4.1 Hard-Assignment Prediction (HA-dLLM)

• HA-dLLM uses discrete candidate-index sequences as states and predicts a K×K matrix of real-valued candidate scores for clean assignments.The “hard assignment” label refers to its state representation and decoder, not its raw model output.Source: §4.1 · p. 7, §4.1.1 · p. 8

4.1.2 Forward process and training objectives

• HA-dLLM subset-shuffling selects T positions and permutes their candidates, and its four alternative objectives are token CE, assignment CE, Sinkhorn NLL, and structured hinge.Assignment CE adds a column-wise term to token CE; Sinkhorn NLL couples row and column marginals through normalization.Source: §4.1.2 · pp. 8–9

4.1.3 Predict–recorrupt sampling and decoding

• HA-dLLM Predict–recorrupt retains shuffles at scheduled corruption levels, predicts clean assignments with Hungarian decoding, and re-corrupts predictions before later evaluations.The final identity-level shuffle yields a permutation, and the appendix states that this sampler preserves permutation validity.Source: §4.1.3 · pp. 9–10

4.2.1 Input and output representation

• SA-dLLM represents corrupted states as continuous assignments in the Birkhoff polytope and applies Sinkhorn normalization to score-predict a clean soft assignment.The model uses candidate-index embeddings to form ranking-slot embeddings from the corrupted assignment.Source: §4.2.1 · pp. 10–11

4.2.2 Forward process and anchored-rollout training

• SA-dLLM anchored-rollout training makes two evaluations, supervising both a noise-endpoint prediction and a re-corrupted model-generated state with a marginal loss.The marginal surrogate scores ground-truth candidates by position rather than computing the exact joint negative log-likelihood over K! permutations.Source: §4.2.2 · p. 11

4.2.3 Predict–recorrupt and Residual-correction sampling

• SA-dLLM sampling reuses a noise matrix and supports Predict–recorrupt or Residual-correction updates, with a single terminal Hungarian projection and no intermediate discrete projection.The soft-state validity condition holds when the noise matrix and predictions are exactly doubly stochastic; finite Sinkhorn normalization only satisfies it approximately.Source: §4.2.3 · pp. 11–12

4.3.1 Input and output representation

• OS-dLLM assigns each current source slot a scalar ordering score, sorts scores into destination order, and trains against the target source order under subset-shuffling corruption.This avoids enumerating K! permutations.Source: §4.3.1 · p. 13, §4.3.2 · p. 13

4.3.2 Forward process and training loss

• OS-dLLM trains all K positions with the Plackett–Luce negative log-likelihood via ListMLE, after sampling and applying a position shuffle to the clean ranking.The training step scores source slots for the corrupted state and computes the permutation loss over all K positions.Source: §4.3.2 · pp. 13–14

4.3.3 Predict–recorrupt sampling and decoding

• At each inference step, stable descending sorting of source-slot scores produces source order, and Predict–recorrupt updates yield a decoder that uses sorting instead of Hungarian assignment.Source: §4.3.3 · p. 14

5.1.1 Datasets and splits

• The study uses Beauty, Sports and Outdoors, and Toys and Games Amazon-2014 data, with 5-core filtering and leave-one-out splits yielding one query per retained user.Each candidate set includes the held-out positive and K−1 highest-scoring nonpositives; only the positive has an observed relevance label, and negatives are ordered by cross-encoder scores as pseudo-labels.Source: §5.1.1 · p. 15, §5.1.2 · p. 15

5.1.2 Candidate construction and target relabeling

• Candidate lists are randomly shuffled and targets relabeled by the inverse permutation; prompts provide purchase history and enumerated candidate descriptions and request each index once in relevance order.The presentation shuffle is separate from diffusion corruption.Source: §5.1.2 · p. 15, §5.1.3 · p. 15

5.2 Metrics

• Evaluation averages each metric over test queries and assigns invalid outputs zero for all relevance metrics; position accuracy measures full target-permutation recovery, not merely the positive item’s rank.Reported measures include Recall@1, Recall@5, MRR, position accuracy, and valid rate. Position and pairwise accuracy are computed from raw outputs, without post hoc projection.Source: §5.2 · pp. 15–16

5.3 Baselines

• Baselines include random valid permutations, constrained autoregressive reranking, and masked diffusion; learned models share the Qwen3-0.6B backbone and the proposed families use their specified assignment or sorting decoders.HA-dLLM uses Hungarian decoding at each step, SA-dLLM only at termination, and OS-dLLM sorting.Source: §5.3 · p. 16

5.4 Parallel Reranking with One Function Evaluation

• At NFE=1, no diffusion family dominates all six settings: SA-dLLM leads on Beauty at K=10, while HA-dLLM variants lead on Sports at both sizes and Toys at K=20.SA-dLLM matches constrained mdLLM on Beauty at K=20; AR comparisons are provisional because its values came from earlier full-test controls.Source: §5.4 · p. 16, Table 3 · p. 17

5.5 Constraint Satisfaction Without Assignment Projection

• At NFE=1, OS-dLLM guarantees valid permutations by sorting, while raw mdLLM has near-zero validity; at NFE=K, raw mdLLM approaches unit validity and has higher Recall@1 in all six settings.Sorting guarantees validity at every budget, whereas raw mdLLM relies on iterative commitment to recover valid permutations.Source: §5.5 · p. 17, Table 4 · p. 17

5.6 Effect of Additional Function Evaluations

• On Beauty with K=20, additional evaluations have parameterization-dependent effects: Recall@1 changes little for constrained mdLLM, Sinkhorn-trained HA-dLLM, and SA-dLLM, but token CE and assignment CE degrade.OS-dLLM improves by 0.5 percentage points in this comparison.Source: §5.6 · p. 17, Table 5 · p. 18
• On Beauty with K=20, OS-dLLM gains 0.497 Recall@1 percentage points from NFE=1 to 20 while losing 0.788 position-accuracy and 0.860 pairwise-accuracy points.Sinkhorn-trained HA-dLLM and SA-dLLM remain stable on relevance and full-order metrics; across six settings, Sinkhorn NLL is the most consistently stable HA-dLLM objective, while OS-dLLM improves in five but stays below the strongest HA-dLLM.Source: §5.6 · p. 17, Table 6 · p. 18

• Hard-, soft-, and direct ordering-score prediction impose permutation constraints through per-step assignment, terminal assignment, and sorting, respectively.The model families also require distinct corruption processes, objectives, and samplers.Source: §6 · p. 18
• Useful rankings can be produced at NFE=1, but additional function evaluations do not consistently improve ranking quality.The conclusion identifies tradeoffs among ranking quality, validity, and iterative stability.Source: §6 · p. 18
• Anchored SA-dLLM remains stable across NFE budgets, whereas OS-dLLM guarantees valid permutations without Hungarian assignment but trails the strongest constrained decoders.The findings support treating corruption, parameterization, objective, and sampler as coupled design choices.Source: §6 · p. 18

A Implementation Details

• The implementation section covers input construction, optimization, checkpoint selection, and evaluation settings.Source: A Implementation Details · p. 21

A.1 Input Serialization and Sequence Lengths

• The serialized autoregressive target has K + 2 tokens: an opening delimiter, K candidate indices, and a closing delimiter.Equation (5) includes only the candidate-token component.Source: §A.1 · p. 21
• Maximum sequence lengths are 2,048 for K = 10 and 4,096 for K = 20; truncation left-truncates the prompt prefix while preserving the complete ranking region.Source: §A.1 · p. 21

A.2.1 Initialization and shared optimization

• Each diffusion model starts from the autoregressive checkpoint trained for the same dataset and candidate budget, with shared AdamW optimization and validation-only checkpoint selection.Training runs for at most ten epochs with effective batch size 128; diffusion learning rate is 10−4.Source: §A.2.1 · p. 21

A.2.2 Hard-state models

• Hard-state models sample T uniformly from [K], replace identity shuffles with random transpositions, and use family-specific assignment or ListMLE objectives.HA-dLLM supervises all K positions, while OS-dLLM applies ListMLE to the full order.Source: §A.2.2 · p. 21

A.2.3 Soft-state model

• SA-dLLM uses an anchored objective with two supervised backbone evaluations per example and samples Ntr from a candidate-budget-dependent training set.Ntrain is {2, 5, 10} for K = 10 and {2, 5, 10, 20} for K = 20.Source: §A.2.3 · p. 21

• Evaluation uses N ∈ {1, 2, 5, 10} for K = 10 and N ∈ {1, 2, 5, 10, 20} for K = 20.Source: §A.3 · p. 21
• For SA-dLLM, the NFE=1 prediction is shared across samplers, and each example reuses the same deterministic noise endpoint across budgets and samplers.HA-dLLM and OS-dLLM report only Predict–recorrupt; preliminary hard-state Residual-correction runs are excluded.Source: §A.3 · p. 21
• Full-test evaluations use eight parallel workers and save both predicted and target permutations.Per-worker batch size is 32 for hard-state models and SA-dLLM with K = 10, and 16 for SA-dLLM with K = 20.Source: §A.3 · p. 21

B.1 Analysis of Hard-State Sampling

• The analysis covers Predict–recorrupt updates for HA-dLLM and OS-dLLM under the stated hard-state schedule.Source: §B.1 · p. 22

B.1.1 Validity of Hard-State Predict–Recorrupt Sampling

• Composing decoded clean predictions with retained shuffles preserves permutation validity at every denoising step, regardless of prediction accuracy.Source: §B.1.1 · p. 22
• With a fixed clean prediction, post-evaluation states equal that prediction composed with the retained shuffle, but the sampled initial state need not match the corresponding forward-corrupted state.For a uniform retained shuffle, the two initial states have the same marginal distribution, though they need not coincide for the same shuffle.Source: §B.1.1 · p. 22

B.1.2 Why Additive Residual-Correction Does Not Directly Extend to Permutations

• The permutation-space residual update preserves validity but cancels the prediction within each step, even when predictions vary across the trajectory.Source: §B.1.2 · p. 22
• With generation initialization, the permutation-space update returns the original presentation order independently of the model, despite matching initial-state marginals under a uniform shuffle.Source: §B.1.2 · p. 23
• In the soft-corruption case, the residual update retains the prediction through a displacement term that depends on the clean state.Source: §B.1.2 · p. 23
• For the alternative composition order, prediction dependence can remain, but the stated initialization restricts outputs to even permutations for HA-dLLM and OS-dLLM.The limitation applies to these direct analogues, not arbitrary permutation-valued residual samplers.Source: §B.1.2 · p. 23

B.2 Analysis of Soft-State Sampling

• Both soft-state samplers keep every state in the Birkhoff polytope when initialization and predictions lie there and the same noise matrix is reused.Source: §B.2 · pp. 23–24
• The terminal assignment is the final clean prediction for Predict–recorrupt and a weighted assignment for Residual-correction.Source: §B.2 · p. 24
• With finite Sinkhorn iterations, the feasibility statement holds only up to normalization error.Source: §B.2 · p. 24
• When every prediction equals the same assignment, both samplers reproduce its affine corruption path; in the ideal permutation-valued case, they recover the clean permutation before terminal decoding.Source: §B.2 · p. 24
• These guarantees establish feasible trajectories and consistency with clean-state predictions, not equality with the data distribution or exact reverse-conditional sampling.Reusing the same noise matrix couples each re-corrupted state to the preceding trajectory rather than making it an independent forward-corruption sample.Source: §B.2 · p. 24