Doubly Stochastic Attention in Transformers

• Doubly Stochastic Attention is a normalization approach that enforces nonnegative attention matrices with unit row and column sums, ensuring balanced token participation.
• It leverages iterative methods like Sinkhorn normalization (or alternatives such as QR-based and quantum mappings) to mitigate over-concentration, explaining-away, and rank collapse in deep networks.
• Empirical studies and optimal transport interpretations indicate improved robustness and rank preservation across applications in vision, NLP, and structured prediction tasks.

Searching arXiv for papers on doubly stochastic attention and Sinkhorn-based attention. Querying "doubly stochastic attention Sinkhorn Transformer". Doubly stochastic attention is a family of self-attention normalizations in which the attention matrix is constrained to be nonnegative and to have unit row sums and unit column sums, replacing the standard row-stochastic Softmax used in Transformers. In the most common formulation, the scaled score matrix $QK^\top/\sqrt d$ is exponentiated and alternately row- and column-normalized by Sinkhorn’s algorithm; related formulations include doubly-normalized attention, optimal-transport constructions, QR-based maps, and parametric quantum layers. Across the literature, the principal motivations are to avoid explaining-away, reduce over-concentration and entropy collapse, preserve rank more effectively across depth, and impose conservation laws in structured prediction settings (Sander et al., 2021, Ding et al., 2020, Lapenna et al., 9 Apr 2026, Meng et al., 2023).

1. Formal definition and canonical normalization

A square matrix $A\in\mathbb R^{n\times n}$ is doubly stochastic if each row and each column sums to one and all entries are non-negative. Equivalently, with $\mathbf1_n$ the all-ones vector, doubly stochasticity is the conjunction

$$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$

The set of all such matrices is the Birkhoff polytope $\mathcal B_n$ (Born et al., 22 Apr 2025).

In Transformer notation, a standard score matrix is

$$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$

or equivalently $C=QK^\top/\sqrt d$. Doubly stochastic attention replaces row-wise Softmax by a map $P=\mathrm{Sinkhorn}(S)$ satisfying $P\mathbf1=\mathbf1$ and $\mathbf1^\top P=\mathbf1^\top$, and uses $A\in\mathbb R^{n\times n}$0, or in the equivalent notation,

$A\in\mathbb R^{n\times n}$1

(Lapenna et al., 9 Apr 2026, Tamayo-Rousseau et al., 28 Jul 2025).

The standard Sinkhorn construction starts from a positive matrix such as $A\in\mathbb R^{n\times n}$2 and alternates row and column normalizations,

$A\in\mathbb R^{n\times n}$3

until convergence. In log-domain implementations, dual scaling variables are updated with log-sum-exp operations for numerical stability. As the number of iterations grows, the iterates converge to a doubly stochastic matrix; in the formulation reported for quantum doubly stochastic Transformers, the limit is the unique matrix in $A\in\mathbb R^{n\times n}$4 minimizing the KL divergence to the input matrix (Sander et al., 2021, Born et al., 22 Apr 2025).

2. Motivations: explaining-away, over-concentration, and balanced token participation

The earliest motivation in the cited literature is the “explaining-away” effect of standard upper-normalized attention. With

$A\in\mathbb R^{n\times n}$5

each query row sums to one, but nothing constrains column sums, so some keys can satisfy $A\in\mathbb R^{n\times n}$6 and are effectively ignored. The doubly-normalized construction in “Attention that does not Explain Away” first column-normalizes $A\in\mathbb R^{n\times n}$7, then row-normalizes, and proves that for every key $A\in\mathbb R^{n\times n}$8,

$A\in\mathbb R^{n\times n}$9

so $\mathbf1_n$0-explaining-away with $\mathbf1_n$1 cannot occur (Ding et al., 2020).

A second motivation is control of attention concentration. The ESPFormer formulation states that row-stochastic attention need not be column-stochastic, so some keys may receive disproportionately large or small total attention mass; enforcing both marginals to equal one balances how tokens send and receive mass, alleviates “over-concentration” on a few keys, encourages information mixing across every position, and acts as a permutation-equivariant regularizer in deep stacks of transformer layers (Shahbazi et al., 11 Feb 2025).

A third motivation is robustness under corrupted or spurious inputs. In Vision Transformers, row-normalization alone can allow certain spurious keys to dominate a query’s attention under noise, and column-normalization alone can similarly over-emphasize certain queries. By forcing both rows and columns to sum to one, doubly stochastic attention is reported to discourage collapse onto a small subset of token interactions; the uniform marginal constraints spread attention mass more evenly, mitigate heteroskedastic noise, and stabilize learned attention patterns so that no single corrupted patch can disproportionately skew the representation (Tamayo-Rousseau et al., 28 Jul 2025).

The rank-decay analysis adds a related viewpoint: standard row-stochastic attention has been shown to suffer from rank collapse and entropy collapse, whereas doubly stochastic attention acts as a form of entropy regularization and preserves rank more effectively than standard Softmax row-stochastic attention, even though it does not eliminate collapse in pure self-attention (Lapenna et al., 9 Apr 2026).

3. Optimal-transport and continuum interpretations

Sinkhorn-based doubly stochastic attention has a direct optimal-transport interpretation. In the ESPFormer formulation, projecting a nonnegative matrix onto the set of doubly stochastic matrices is presented as an entropy-regularized optimal transport problem,

$\mathbf1_n$2

with alternating updates for scaling vectors $\mathbf1_n$3 and $\mathbf1_n$4. The computational bottleneck emphasized there is that each Sinkhorn iteration costs $\mathbf1_n$5 and must be executed sequentially, giving total cost $\mathbf1_n$6 when $\mathbf1_n$7 iterations are used (Shahbazi et al., 11 Feb 2025).

“Sinkformers” develops a more structural theory. There, Sinkhorn normalization is written as a KL projection onto the set $\mathbf1_n$8 of doubly stochastic matrices, and residual self-attention iterations are analyzed in the infinite-depth, mean-field limit. Under the symmetry assumption

$\mathbf1_n$9

the self-attention dynamics induced by the unnormalized kernel $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$0 or by the fully Sinkhorn-normalized kernel $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$1 define a Wasserstein gradient flow; by contrast, the SoftMax-normalized kernel does not arise from a Wasserstein gradient. In the joint large-sample, infinite-depth, small-bandwidth limit, the Sinkhorn dynamics reduce to the heat equation $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$2 (Sander et al., 2021).

The quantum doubly stochastic Transformer work reframes the same object geometrically. It emphasizes that doubly stochastic matrices form the Birkhoff polytope and distinguishes between approximate Sinkhorn outputs for small numbers of iterations and exact DSM constructions obtained from QR factorization or from a variational quantum circuit. In that presentation, Sinkhorn with $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$3 converges exponentially in $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$4 to the unique matrix in $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$5 minimizing KL divergence to the input matrix, whereas QR and QontOT always produce exact DSMs (Born et al., 22 Apr 2025).

4. Rank decay, residual dynamics, and the role of skip connections

The most explicit depth-wise analysis of doubly stochastic attention is the 2026 rank-decay study. It considers pure self-attention networks with no feed-forward block and no skip connections, and measures convergence toward the rank-one matrix $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$6 through the residual

$$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$7

In this setting, the paper derives single-head and multi-head bounds showing that the residual, and hence rank, decays to zero doubly exponentially with the network depth $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$8 under Sinkhorn normalization (Lapenna et al., 9 Apr 2026).

This is a critical corrective to a common overstatement. Doubly stochastic normalization preserves rank more effectively than standard Softmax row-stochastic attention, but in pure self-attention it does not prevent collapse altogether. The paper explicitly states that rank decays to one doubly exponentially with depth under Sinkhorn normalization, mirroring the previously established Softmax result of Dong et al. (2021), while obtaining the analogous structure under the spectral norm $$A\,\mathbf1_n=\mathbf1_n,\qquad A^\top\mathbf1_n=\mathbf1_n,\qquad A_{ij}\ge 0\ \forall i,j.$$9 rather than a mixed $\mathcal B_n$0 norm (Lapenna et al., 9 Apr 2026).

The same study identifies skip connections as crucial. In the path decomposition, inserting a skip at layer $\mathcal B_n$1 omits the stochastic-matrix factor $\mathcal B_n$2, shortening the product; since products of stochastic matrices converge to a rank-one projector as the number of factors grows, skipping layers interrupts this convergence and preserves rank. Empirically, both Softmax and Sinkhorn benefit, but the gap between the two schemes becomes larger when skips are present. The paper also notes that feed-forward layers and layer-norm layers do not fully prevent rank decay (Lapenna et al., 9 Apr 2026).

5. Constructions and variants

The literature contains several distinct mechanisms for obtaining doubly stochastic attention, ranging from one-pass double normalization to exact non-Sinkhorn constructions and low-rank transport couplings.

Method	Construction	Stated property
DNAS (Ding et al., 2020)	Column-normalize $\mathcal B_n$3, then row-normalize	Avoids explaining-away; $\mathcal B_n$4
Sinkformer (Sander et al., 2021)	Iterative row/column normalization of $\mathcal B_n$5	Wasserstein gradient-flow and heat-diffusion theory
Reaction-prediction DS attention (Meng et al., 2023)	Two doubly stochastic self-attention mappings via Sinkhorn	Satisfies electron-counting and symmetry rules
ESPFormer (Shahbazi et al., 11 Feb 2025)	Expected Sliced Transport Plans with SoftSort	Fully parallelizable; no iterative Sinkhorn normalization
QR-based DSM (Born et al., 22 Apr 2025)	Thin QR followed by $\mathcal B_n$6	Exact DSM
QontOT / QDSFormer (Born et al., 22 Apr 2025)	Variational quantum circuit $\mathcal B_n$7 reshaped into $\mathcal B_n$8	First parametric DSM layer
LOTFormer (Shahbazi et al., 27 Sep 2025)	Two entropic OT problems glued through a learnable pivot	Doubly-stochastic, rank at most $\mathcal B_n$9, $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$0 application

The reaction-prediction formulation is the clearest example of task-specific structure. There, two doubly stochastic attention maps $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$1 and $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$2 are used so that the predicted electron redistribution $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$3 is “doubly conservative”: every row sum is zero and every column sum is zero, and these properties are preserved by final symmetrization. The multi-head extension either enforces per-head double stochasticity or a $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$4-sum constraint on the aggregate, and exact Sinkhorn enforcement is preferred to a soft penalty for numerical precision (Meng et al., 2023).

The non-Sinkhorn extensions pursue either flexibility or efficiency. ESPFormer replaces iterative Sinkhorn normalization by a differentiable Expected Sliced Transport Plan built from soft sorting, retaining approximate double stochasticity while enabling parallel computation. LOTFormer instead introduces a learnable pivot measure of size $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$5, solves two entropic OT problems between queries and pivot and between pivot and keys, and glues the couplings into an attention matrix that is provably doubly stochastic and has rank at most $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$6. QDSFormer goes in a different direction: QR produces exact DSMs deterministically, whereas the QontOT quantum layer is parametric and explores the unistochastic subset of the Birkhoff polytope (Shahbazi et al., 11 Feb 2025, Shahbazi et al., 27 Sep 2025, Born et al., 22 Apr 2025).

6. Empirical behavior, applications, and implementation practice

Empirical results reported across the cited papers are uniformly framed as comparisons against row-stochastic baselines rather than as claims that doubly stochasticity solves all attention pathologies. In “Sinkformers,” gains are reported across 3D shape, vision, and NLP tasks with negligible extra cost: on ModelNet40, Set Transformer mean accuracy improves from $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$7 to $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$8, and Point Cloud Transformer mean accuracy from $$S=\frac{(XW_Q)(XW_K)^\top}{\sqrt{d_{qk}}}\in\mathbb R^{n\times n},$$9 to $C=QK^\top/\sqrt d$0; additional improvements are reported on IMDb sentiment analysis, IWSLT14 De–En translation, Cats vs Dogs, and MNIST patch-size robustness, while per-epoch runtime overhead is reported as $C=QK^\top/\sqrt d$1 on a V100 GPU (Sander et al., 2021).

Under explicit corruption tests, the ViT robustness study compares Softmax, Linear, Sigmoid, Cosine, and Doubly Stochastic attention on CIFAR-10, CIFAR-100, and Imagenette. Its central conclusion is that Doubly Stochastic attention is the most robust. On CIFAR-10, relative accuracies in corrupted settings are $C=QK^\top/\sqrt d$2, $C=QK^\top/\sqrt d$3, and $C=QK^\top/\sqrt d$4, compared with the next-best Softmax values $C=QK^\top/\sqrt d$5, $C=QK^\top/\sqrt d$6, and $C=QK^\top/\sqrt d$7; on Imagenette, Doubly Stochastic attention leads in all three corrupted regimes with $C=QK^\top/\sqrt d$8, $C=QK^\top/\sqrt d$9, and $P=\mathrm{Sinkhorn}(S)$0, preserving almost its full clean-data accuracy under combined corruption (Tamayo-Rousseau et al., 28 Jul 2025).

The rank-preservation experiments in the 2026 analysis quantify the same trend with residual metrics. For the median normalized residual $P=\mathrm{Sinkhorn}(S)$1 at depth $P=\mathrm{Sinkhorn}(S)$2, AG’s News reports $P=\mathrm{Sinkhorn}(S)$3 and $P=\mathrm{Sinkhorn}(S)$4 for Softmax without and with skips, versus $P=\mathrm{Sinkhorn}(S)$5 and $P=\mathrm{Sinkhorn}(S)$6 for Sinkhorn; at end of training, the self-attention output residual on AG’s News is $P=\mathrm{Sinkhorn}(S)$7 and $P=\mathrm{Sinkhorn}(S)$8 for Softmax, versus $P=\mathrm{Sinkhorn}(S)$9 and $P\mathbf1=\mathbf1$0 for Sinkhorn. MNIST and Cats vs Dogs show the same direction of effect. The paper summarizes these numbers as illustrative of the consistent trend that Sinkhorn slows rank collapse more than Softmax, especially when skip connections are active (Lapenna et al., 9 Apr 2026).

Other domains show analogous benefits under different inductive biases. The reaction-prediction framework reports that its doubly stochastic graph-based non-autoregressive decoder consistently improves predictive performance while satisfying the electron-counting rule and the symmetry rule, and does not bring an unbearable additional computational cost (Meng et al., 2023). QDSFormer reports that the hybrid classical-quantum doubly stochastic Transformer consistently surpasses both a standard Vision Transformer and other doubly stochastic Transformers on multiple small-scale object recognition tasks, with improved training stability and lower performance variation; on FashionMNIST, a 2-layer ViT reaches $P\mathbf1=\mathbf1$1 with QontOT versus $P\mathbf1=\mathbf1$2 with Sinkhorn, and on MedMNIST the mean accuracy over 7 datasets is $P\mathbf1=\mathbf1$3 for QontOT versus $P\mathbf1=\mathbf1$4 for Sinkhorn (Born et al., 22 Apr 2025).

Implementation guidance varies by formulation and task. Reported iteration counts range from a single column-then-row normalization in DNAS, to $P\mathbf1=\mathbf1$5–$P\mathbf1=\mathbf1$6 Sinkhorn steps in Sinkformers, to $P\mathbf1=\mathbf1$7–$P\mathbf1=\mathbf1$8 iterations in the robustness study, and $P\mathbf1=\mathbf1$9–$\mathbf1^\top P=\mathbf1^\top$0 iterations in the practical recommendations of the rank-decay study. Log-domain updates are repeatedly recommended for numerical stability. The robustness study states that no explicit additional regularizer was needed and that all other parts of the multi-head block remain unchanged; it also reports a $\mathbf1^\top P=\mathbf1^\top$1–$\mathbf1^\top P=\mathbf1^\top$2 increase in attention layer compute time. The rank-decay study recommends residual connections, notes that warm-starting the scaling vectors can reduce compute, and points to sliced transport, learned scaling, and quantum variational circuits as alternatives when iterative cost is prohibitive (Ding et al., 2020, Sander et al., 2021, Tamayo-Rousseau et al., 28 Jul 2025, Lapenna et al., 9 Apr 2026).

Taken together, these works define doubly stochastic attention less as a single algorithm than as a normalization principle: attention should conserve mass along both query and key dimensions. Within that principle, the literature separates three issues. The first is structural balance, captured by the row-and-column constraints themselves. The second is the mechanism used to impose those constraints, ranging from Sinkhorn scaling to soft sorting, QR factorization, variational quantum circuits, and low-rank transport gluing. The third is the depth-wise behavior of the resulting network: doubly stochasticity regularizes entropy, improves robustness, and preserves rank more effectively than row-stochastic Softmax, but deep pure self-attention still requires skip connections to avoid doubly exponential collapse toward rank one (Shahbazi et al., 11 Feb 2025, Born et al., 22 Apr 2025, Shahbazi et al., 27 Sep 2025, Lapenna et al., 9 Apr 2026).