Does Sinkhorn doubly stochastic attention preserve rank better than softmax attention?

Rank collapse means tokens stop being distinguishable

A self-attention layer mixes every token with every other token. In the standard implementation the attention matrix is row-stochastic: each row sums to one, so it reads as a probability distribution over which tokens a given query attends to. Rank collapse is what happens when repeated application of such layers pushes all token vectors toward the same value, so the sequence carries less and less distinct information. The paper measures this with the residual, the part of the representation left after subtracting the shared mean component across tokens; a smaller residual norm means representations are closer to uniform.

This is a different failure mode from entropy collapse, which the paper describes separately as attention distributions becoming highly concentrated so that each query attends to only a small subset of tokens. Rank collapse concerns the diversity of token representations after mixing; entropy collapse concerns how peaked the attention weights are. The paper's experiments target the first, and its central comparison is between two ways of normalizing the attention matrix rather than between temperature settings.

The second normalization is doubly stochastic: rows and columns both sum to one. The paper obtains it by replacing softmax with the Sinkhorn algorithm, an iterative procedure that alternately rescales rows and columns. Empirically the authors report that attention matrices become effectively doubly stochastic after 50 Sinkhorn iterations, and they use that value in all experiments. The motivation given is that doubly stochastic attention acts like entropy regularization, spreading attention more evenly than a row-only normalization.

Theory and experiments answer different questions

The paper separates two claims that are easy to conflate. The theoretical part analyzes a pure self-attention network, meaning stacked attention layers with feed-forward layers and skip connections disabled. The empirical part measures rank decay inside full trained Transformers, where each block contains a self-attention layer, a feed-forward layer, skip connections and layer normalization. A result about the pure network therefore does not automatically describe a trained Transformer, and the paper keeps the two apart.

Within the experiments the authors also measure two distinct quantities. One is the rank decay of the product of attention matrices along a path through the layers. The other is the rank decay of the self-attention output itself. These are related but not identical objects, and the paper reports them in separate subsections rather than treating one as a substitute for the other.

The empirical setting covers three datasets. Text classification uses AG's News, whose articles fall into four categories, with a 24-layer Transformer encoder, model dimension 128, four attention heads and maximum sequence length 128. Image classification uses MNIST and Cats and Dogs with a Vision Transformer of 24 layers, model dimension 128 and four heads; images are split into 4×4 patches for MNIST and 16×16 patches for Cats and Dogs. In every case softmax attention is compared against Sinkhorn attention under the same architecture.

Sinkhorn attention decays more slowly in trained models

In the trained architectures, the measured residual norm under Sinkhorn normalization stays higher than under softmax as depth increases, which is the paper's headline empirical result: doubly stochastic attention mitigates rank collapse relative to standard row-stochastic attention. The authors describe this effect as most evident in the early layers, before rank collapse becomes severe, and in the presence of skip connections.

The skip-connection condition is not a minor detail. The paper states that skip connections are crucial to mitigating rank collapse, consistent with earlier work on softmax attention, and its experiments vary whether skip connections and feed-forward layers are present. The reported advantage of Sinkhorn normalization is therefore tied to a specific architectural setting rather than to depth alone.

The measurements are reported as means over the batch with shaded regions indicating one standard deviation, and the value at layer zero corresponds to the input embeddings before the first attention layer. That baseline matters for reading the curves: the comparison is about how much of the initial residual survives layer by layer, not about absolute representation quality at any single layer.

Pure self-attention still collapses to rank one

The theoretical result is a bound, not a reprieve. For a single-head single-layer attention with Sinkhorn normalization, the residual norm of the output is bounded by a constant factor times the cube of the input residual norm, where the constant involves the sequence length, the query and key dimension, and spectral norms of the weight matrices. Stacking L such layers multiplies the exponent, giving a bound that scales with the input residual raised to the power 3^L.

That exponent is the key number. Because the exponent grows as a power of depth rather than linearly, the residual converges to zero doubly exponentially with depth, and the output approaches a rank-one matrix. The paper notes that the same doubly exponential behavior was already established for softmax attention, so Sinkhorn normalization changes the constants and the empirical behavior but not this asymptotic outcome in the pure network.

The bound holds in a specific regime: the unnormalized attention logits matrix must have spectral norm much smaller than one, which lets the authors locally approximate the Sinkhorn operator by its linearization around zero. The paper is explicit that this is an approximation, and it lists the tightness of these bounds as a limitation, noting that the estimates could potentially be tightened. The result should be read as a statement about that regime, not as a general theorem for arbitrary attention magnitudes.

Random matrix products show no Sinkhorn advantage

A control experiment sharpens what the empirical advantage does and does not mean. The authors generate random matrices, normalize each product either with softmax or with Sinkhorn, and measure the same residual quantity as a function of path depth, averaging over 100 sampled attention paths per depth. When the matrices in the product are uncorrelated, the two normalizations show essentially the same rank decay.

The same null result appears when a random subset of factors is replaced by the identity to mimic skip connections, with half the matrices replaced in the simulation. So the Sinkhorn advantage observed in trained Transformers is not a generic property of doubly stochastic matrices under multiplication.

The authors offer an explanation and label it as such: attention matrices along a path in a Transformer are correlated, because the attention matrix at one layer takes as input the output produced by the previous layer. They also observe that products of random stochastic matrices decay more slowly than attention matrices extracted from a trained Transformer, which they attribute to this input dependence creating a cascading effect that amplifies decay. Understanding how Sinkhorn and softmax attention matrices interact is left as future work, so the mechanism behind the trained-model advantage is not settled by this paper.

Where the Sinkhorn advantage should change your reading

The practical distinction the paper supports is conditional. If you are comparing normalization schemes inside a full Transformer with skip connections, the evidence points toward Sinkhorn doubly stochastic attention as the slower-decaying option, with the largest separation early in the network. If you are reasoning about a stripped-down stack of attention layers with no skip connections and no feed-forward layers, the theory says both normalizations end at rank one, and the choice does not rescue the representation.

The random-matrix control adds a second condition: the advantage depends on the correlated, input-dependent structure of attention matrices in a trained model, not on doubly stochasticity alone. A claim that Sinkhorn normalization generically preserves rank better than softmax would overreach the evidence, since the paper's own uncorrelated test shows parity.

Two limitations bound how far the empirical results travel. The models were trained as full Transformers and the feed-forward layers were disabled only at inference to isolate pure self-attention behavior, so training a pure self-attention network from scratch remains untested. And the theoretical bounds rest on a first-order approximation of the Sinkhorn operator in a small-logit regime. The paper's own priority for future work is comparing the Lipschitz properties of Sinkhorn and softmax attention, which would speak to robustness and training stability rather than to rank alone.

The judgment worth carrying away is therefore a split, not a verdict: Sinkhorn normalization buys slower rank decay in the early layers of trained Transformers with skip connections, while the doubly exponential collapse to rank one in pure self-attention survives the change of normalization.

Sources used in this answer