Learning Unbiased Permutations via Flow Matching

TL;DR

PermFlow introduces a flow matching framework that directly models multimodal permutation distributions, overcoming limitations of entropy-regularized methods like Sinkhorn.

Contribution

It presents a novel flow-based approach with exact constraint preservation, enabling accurate learning of permutation distributions under ambiguity.

Findings

PermFlow outperforms Sinkhorn-based methods on ambiguous permutation tasks.

It accurately recovers multiple valid permutations in ambiguous scenarios.

Achieves high accuracy on visual sorting and linear assignment benchmarks.

Abstract

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 visual sorting task with blended-digit ambiguity and a symmetric linear assignment problem, PermFlow achieves high accuracy on unambiguous inputs and recovers both valid permutations under ambiguity, where Sinkhorn-based baselines structurally fail.