Learning Permutation Distributions via Reflected Diffusion on Ranks

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.

Many important problems in AI involve arranging things in the right order — for example sorting images, reconstructing shuffled puzzles, recommending items, or planning routes. But teaching AI systems to generate good orderings is difficult because even small changes can completely alter the final arrangement. In this work, we introduce a new way to model orderings using diffusion models. Instead of repeatedly “shuffling” items in abrupt ways, we relax the position of each item from a discrete number to a continuous one and gradually adds noise in this smoother space. The intuition is that small changes in the representation should lead to small changes in the ordering, making the learning process much more stable than directly manipulating discrete permutations. This creates more stable learning and makes the model much better at handling long and complicated sequences. We also design the model so that each decision can adapt based on the choices already made, which is especially useful for step-by-step problems like route planning. Across several tasks — including number sorting, image puzzle reconstruction, and traveling salesperson problems — our method consistently performs better than previous approaches, especially on larger and more difficult examples.