Learning Permutation Distributions via Reflected Diffusion on Ranks

TL;DR

This paper introduces Soft-Rank Diffusion, a novel permutation modeling framework that uses soft ranks and generalized Plackett-Luce denoisers to improve learning and denoising of permutation distributions, especially for long sequences.

Contribution

It proposes a structured soft-rank forward process and contextualized PL denoisers, advancing permutation diffusion methods with smoother trajectories and enhanced expressivity.

Findings

Outperforms prior diffusion baselines on sorting benchmarks

Achieves strong results in long-sequence and sequential tasks

Provides a more tractable and smooth permutation modeling approach

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.