Learning Permutation Distributions via Reflected Diffusion on Ranks

Sizhuang He; Yangtian Zhang; Shiyang Zhang; David van Dijk

cs.LGMar 18, 2026

Learning Permutation Distributions via Reflected Diffusion on Ranks

Authors: Sizhuang He, Yangtian Zhang, Shiyang Zhang, David van Dijk

Organizations: Department of Computer Science, Yale University, New Haven, CT, USA.

Abstract

The finite symmetric group S_n provides a natural domain for permutations, yet learning probability distributions on S_n is challenging due to its factorially growing size and discrete, non-Euclidean structure. Recent permutation diffusion methods define forward noising via shuffle-based random walks (e.g., riffle shuffles) and learn reverse transitions with Plackett-Luce (PL) variants, but the resulting trajectories can be abrupt and increasingly hard to denoise as n grows. We propose Soft-Rank Diffusion, a discrete diffusion framework that replaces shuffle-based corruption with a structured soft-rank forward process: we lift permutations to a continuous latent representation of order by relaxing discrete ranks into soft ranks, yielding smoother and more tractable trajectories. For the reverse process, we introduce contextualized generalized Plackett-Luce (cGPL) denoisers that generalize prior PL-style parameterizations and improve expressivity for sequential decision structures. Experiments on sorting and combinatorial optimization benchmarks show that Soft-Rank Diffusion consistently outperforms prior diffusion baselines, with particularly strong gains in long-sequence and intrinsically sequential settings.

Figures & tables

Figure 1: Soft-Rank Diffusion. We define the forward diffusion by relaxing each item’s rank to a continuous soft-rank variable and evolving these soft ranks in a reflected diffusion process [ 19 ] . At each time $t$ , the soft ranks induce a discrete ordering by simple sorting, thereby yielding a forward process in permutation space. For reverse sampling, we couple a discrete denoiser in permutation space (predicting a clean permutation from $X_{t}$ ) with an auxiliary continuous update: we lift the predicted permutation to a grid-aligned soft-rank vector $\hat{Z}_{0}$ , sample an intermediate latent $Z_{s}$ for $s<t$ from a conditional reverse kernel $p(Z_{s}\mid Z_{t},\hat{Z}_{0},Z_{1})$ , and map back to permutation space by sorting $Z_{s}$ , thus stepping backward in time.

Panel of Figure 2 — Figure 2: Sampling in PL/GPL/cGPL. (a) PL and GPL. In PL, each item is assigned a single scalar score; sampling a permutation amounts to repeatedly sampling from the same score vector without replacement, masking selected items and renormalizing at each step. In GPL, each item is assigned a length- $N$ score vector, yielding position-specific logits: at step $i$ we sample according to the $i$ -th column of the score matrix (after masking previously selected items), proceeding sequentially from $i=1$ to $N$ . (b) cGPL. In cGPL, each item is assigned a position-dependent score vector dynamically . Sampling proceeds autoregressively : the score vector at later positions depends on the outcomes sampled at preceding positions. As in GPL, we apply masking and subsequent renormalization to obtain a valid probability distribution over the remaining items.

Panel of Figure 3 — Figure 3: Exact-match accuracy and relative improvement versus SymmetricDiffusers as a function of sequence length (N). Panel 3(a) shows accuracy on a logit scale. Panel 3(b) reports the accuracy ratio of Soft-Rank Diffusion relative to SymmetricDiffusers on a log scale. Starred points mark sequence lengths where SymmetricDiffusers attains zero accuracy; the corresponding ratios are undefined and are clipped to the top of the plotted range for visualization.

Method	TSP-20		TSP-50
	$L(\text{Tour})$ $\downarrow$	Gap $\downarrow$	$L(\text{Tour})$ $\downarrow$	Gap $\downarrow$
SymmetricDiffusers with GPL [ 38 ]	5.4342	0.4169	12.8525	1.2618
Soft-Rank Diffusion with cGPL Ours	3.8732	0.0079	5.7336	0.0078
Soft-Rank Diffusion with Pointer-cGPL Ours	3.8557	0.0034	5.7157	0.0047

Table 2: TSP performance on TSP-20 and TSP-50. We report tour length

L(\text{Tour})

and the optimality gap (both lower is better).

Method	Metric	Grid Size $n\times n$
		$3\times 3$ ( $N\!=\!9$ )	$4\times 4$ ( $N\!=\!16$ )	$8\times 8$ ( $N\!=\!64$ )
SymmetricDiffusers with GPL [ 38 ]	Kendall - Tau $\uparrow$	0.8312	0.6101	0.0250
	Accuracy $\uparrow$	0.7002	0.2172	0.0000
	Correctness $\uparrow$	0.8662	0.6753	0.0501
Soft-Rank Diffusion with cGPL Ours	Kendall - Tau $\uparrow$	0.9494	0.9320	0.2521
	Accuracy $\uparrow$	0.9221	0.8752	0.0924

Table 3: CIFAR-10 jigsaw benchmark. Bold marks the best value per (metric,

n

). All methods use a 7-layer Transformer backbone with the same hidden width.

Figure 4: Per-position accuracy heatmaps on $8\times 8$ CIFAR-10 jigsaw. Each cell is the fraction of test samples where that grid position is correctly assigned. Left: SymmetricDiffusers (mean $5.0\%$ ). Middle: cGPL (mean $28.5\%$ ). Right: Pointer-cGPL (mean $52.3\%$ ). Random baseline $1/64\approx 1.6\%$ .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Sequence Length	Steepness	Sorting Network	Loss Weight	Learning Rate
9	34	odd even	1.00	$10^{-4}$
15	25	odd even	0.10	$10^{-4}$
32	124	odd even	0.10	$10^{-4}$
52	130	bitonic	0.10	$10^{-3.5}$
75	135	bitonic	0.10	$10^{-3.5}$
100	140	bitonic	0.10	$10^{-3.5}$

Appendix

Table 4: Hyperparameters for Error-Free Diffsort on 4-digit MNIST sorting.

Decoder	Init Distribution	Perm Acc (%)	Kendall $\tau$	Prop Correct (%)
cGPL	Uniform	86.67	0.9303	93.93
cGPL	Gaussian	86.96	0.9317	94.08
Pointer-cGPL	Uniform	87.35	0.9365	94.43
Pointer-cGPL	Gaussian	86.94	0.9347	94.38

Appendix

Table 5: Matched train/eval

p_{\text{ref}}

on CIFAR10 jigsaw (

n=4

\sigma=0.5

T=5

). Both Uniform and Gaussian reference distributions yield essentially identical performance, indicating the choice of

p_{\text{ref}}

has negligible impact on training.

Decoder	Train $p_{\text{ref}}$	Eval: Uniform	Eval: Gaussian
cGPL	Uniform	86.67 / 93.93 / $\tau\!=\!0.930$	87.09 / 94.08 / $\tau\!=\!0.931$
cGPL	Gaussian	86.71 / 94.01 / $\tau\!=\!0.932$	86.96 / 94.08 / $\tau\!=\!0.932$
Pointer-cGPL	Uniform	87.35 / 94.43 / $\tau\!=\!0.937$	87.33 / 94.51 / $\tau\!=\!0.937$
Pointer-cGPL	Gaussian	87.17 / 94.31 / $\tau\!=\!0.935$	86.94 / 94.38 / $\tau\!=\!0.935$

Appendix

Table 6: Mismatched train/eval

p_{\text{ref}}

on CIFAR10 jigsaw (

n=4

\sigma=0.5

T=5

). Each cell shows Perm Acc / Prop Correct / Kendall

\tau

. Performance is preserved even when training and evaluation use different reference distributions, confirming the Brownian bridge formulation drives

z

toward

z_{0}

regardless of starting distribution.

	$[0,1]^{N}$ (ours)	Birkhoff polytope $B_{N}$	Orthogonal group $O(N)$
Dimension	$N$	$(N-1)^{2}$	$N(N-1)/2$
Dynamics	Independent reflected Brownian bridges with closed-form transition kernels	Coupled entries; requires per-step Sinkhorn projection	Riemannian geodesics
Decoding	$\mathrm{argsort}$ (Eq. 3 )	Hungarian or Sinkhorn rounding	Nearest permutation matrix

Appendix

Table 7: Comparison of three continuous relaxations of permutations. The soft-rank relaxation

[0,1]^{N}

used by Soft-Rank Diffusion has linear dimension and admits independent closed-form Brownian-bridge dynamics, while Birkhoff- and

O(N)

-based relaxations require coupled, more expensive constructions.

Panel of Figure 5 — Figure 5: Model architecture and Pointer-cGPL parameterization for permutation generation. We adopt a standard encoder–decoder Transformer backbone [ 34 ] . In the vanilla cGPL parameterization, the decoder states are mapped to logits via a linear output head (yielding a distribution over candidates at each step). In contrast, Panel 5(b) illustrates Pointer-cGPL , where a bi-affine compatibility module scores each encoded item representation against the current decoder state, producing step-wise logits over the input items that can be interpreted as a pointer distribution.

Figure 6: Accuracy ratio (green; Soft-Rank Pointer accuracy / SymmetricDiffusers accuracy) and runtime ratio (red; Soft-Rank Pointer wall-clock time / SymmetricDiffusers wall-clock time) across problem scales. The accuracy advantage grows superlinearly with $N$ , reaching $4.0\times$ at $N\!=\!16$ and becoming unbounded at $N\!=\!64$ (*baseline accuracy is 0%). Runtime overhead remains moderate: with KV cache optimization (solid red), the slowdown is $1.0\times$ – $4.3\times$ , while the shaded region shows the speedup across diffusion steps (up to $1.55\times$ at $N\!=\!64$ ). The accuracy gains far outpace the runtime cost at every scale.

Figure 7: Patch-wise Manhattan displacement distributions on

8\times 8

CIFAR-10 jigsaw. Manhattan displacement is the grid distance between a patch’s predicted and correct positions. Dashed gray = random baseline (mean

5.2

). Left: SymmetricDiffusers (mean

3.8

). Middle: cGPL (mean

2.8

). Right: Pointer-cGPL (mean

1.6

Method	$3\!\times\!3$	$4\!\times\!4$	$8\!\times\!8$
SymmetricDiffusers	28.0s	21.6s	36.3s
cGPL, w/o cache	31.4s	41.5s	216.6s
cGPL, w/ cache	26.3s	31.4s	131.7s
speedup	1.19 $\times$	1.32 $\times$	1.64 $\times$
Pointer-cGPL, w/o cache	31.6s	51.6s	242.3s
Pointer-cGPL, w/ cache	26.7s	36.7s	156.7s

Appendix

Table 9: Wall-clock inference time. Single H100 GPU,

10{,}000

test samples.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Learning Unbiased Permutations via Flow Matching
May 16, 2026Yimeng Min, Carla P. GomesFlow Matching Combinatorial Optimization
Department of Computer Science, Cornell University, Ithaca, NY, USA
Learning Permutation from Structure Without Supervision
May 25, 2026Ran Eisenberg, Ofir LindenbaumUnsupervised Learning Jigsaw Puzzle Solving
Faculty of Engineering, Bar-Ilan University, Ramat Gan, Israel. · Faculty of Engineering, Bar-Ilan
DenoiseRank: Learning to Rank by Diffusion Models
Feb 17, 2026Ying Wang, Preslav Nakov, Shangsong LiangLearning to Rank
Sun Yat-sen University, Guangzhou, China · MBZUAI, Abu Dhabi, UAE

May 16, 2026J/K move · Enter open · S save

Learning Unbiased Permutations via Flow Matching

Yimeng Min, Carla P. Gomes

Department of Computer Science, Cornell University, Ithaca, NY, USA

Flow Matching Combinatorial Optimization Flow-Based Generative Modeling

Open paper AskPDF arXiv

Learning permutations is fundamental to sorting, ranking, and matching, but existing differentiable methods based on entropy-regularized Sinkhorn produce a single softened solution and collapse under ambiguity. We present PermFlow, a conditional flow matching framework that operates directly on the affine subspace of matrices with unit row and column sums. A closed-form tangent-space projector preserves these constraints exactly along every trajectory, by construction rather than through iterative correction, and a nearest-target coupling routes distinct noisy initializations toward distinct valid permutations. The result is a model that captures multimodal permutation distributions rather than collapsing them to a single mode. On a range of problems, PermFlow achieves high accuracy on unambiguous inputs and recovers both valid permutations under ambiguity.

Abstract

Diffusion Models

Explore similar work

Learning Unbiased Permutations via Flow Matching

Learning Permutation from Structure Without Supervision

DenoiseRank: Learning to Rank by Diffusion Models