---
title: Doubly Stochastic Attention in Transformers
url: https://www.emergentmind.com/topics/doubly-stochastic-attention
type: topic
---

# Doubly Stochastic Attention in Transformers

Searching arXiv for recent papers on doubly stochastic attention and Sinkhorn-based attention.
Querying arXiv: "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 [2110.11773; 2009.14308; 2604.07925; 2306.06119].

## 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\) [2504.16275].

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 \(PXW_V\), or in the equivalent notation,
\[
\mathrm{DoublyStochasticAttn}(Q,K,V)=\mathrm{Sinkhorn}(QK^\top/\sqrt d)\cdot V
\]
[2604.07925; 2507.20453].

The standard Sinkhorn construction starts from a positive matrix such as \(K^{(0)}=\exp(S)\) and alternates row and column normalizations,
\[
N_R(K)_{ij}=\frac{K_{ij}}{\sum_{j'}K_{ij'}},\qquad
N_C(K)_{ij}=\frac{K_{ij}}{\sum_{i'}K_{i'j}},
\]
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 \(\mathcal B_n\) minimizing the KL divergence to the input matrix [2110.11773; 2504.16275].

## 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
\[
\pi_{ij}=\frac{\exp(\mathbf q_i^\top \mathbf k_j/\sqrt D)}{\sum_{j'}\exp(\mathbf q_i^\top \mathbf k_{j'}/\sqrt D)},
\]
each query row sums to one, but nothing constrains column sums, so some keys can satisfy \(\sum_i\pi_{ij}\approx 0\) and are effectively ignored. The doubly-normalized construction in “Attention that does not Explain Away” first column-normalizes \(Z_{ij}=e^{q_i^\top k_j}\), then row-normalizes, and proves that for every key \(j\),
\[
\sum_{i=1}^S \pi_{ij}\ge \frac1S,
\]
so \(\epsilon\)-explaining-away with \(\epsilon\ll 1/S\) cannot occur [2009.14308].

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 [2502.07962].

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 [2507.20453].

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 [2604.07925].

## 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,
\[
\min_{A\in\mathbb R_+^{N\times N}}
\sum_{ij} C_{ij}A_{ij}+\varepsilon\sum_{ij}A_{ij}\log A_{ij}
\quad\text{s.t.}\quad
A\mathbf1_N=\mathbf1_N,\;
A^\top\mathbf1_N=\mathbf1_N,
\]
with alternating updates for scaling vectors \(u\) and \(v\). The computational bottleneck emphasized there is that each Sinkhorn iteration costs \(\mathcal O(N^2)\) and must be executed sequentially, giving total cost \(\mathcal O(SN^2)\) when \(S\) iterations are used [2502.07962].

“Sinkformers” develops a more structural theory. There, Sinkhorn normalization is written as a KL projection onto the set \(U\) of doubly stochastic matrices, and residual self-attention iterations are analyzed in the infinite-depth, mean-field limit. Under the symmetry assumption
\[
W_K^\top W_Q = W_Q^\top W_K = -\,W_V,
\]
the self-attention dynamics induced by the unnormalized kernel \(k^0=\exp(c)\) or by the fully Sinkhorn-normalized kernel \(k^\infty\) 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 \(\partial_t \rho=\Delta \rho\) [2110.11773].

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 \(K\to\infty\) converges exponentially in \(K\) to the unique matrix in \(\mathcal B_n\) minimizing KL divergence to the input matrix, whereas QR and QontOT always produce exact DSMs [2504.16275].

## 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 \(\tfrac1n\mathbf1\mathbf1^\top X\) through the residual
\[
\mathrm{res}(X)=X-\frac1n\mathbf1\mathbf1^\top X.
\]
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 \(L\) under Sinkhorn normalization [2604.07925].

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 \(\|\cdot\|_2\) rather than a mixed \(\ell_{1,\infty}\) norm [2604.07925].

The same study identifies skip connections as crucial. In the path decomposition, inserting a skip at layer \(\ell\) omits the stochastic-matrix factor \(P^\ell\), 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 [2604.07925].

## 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 [2009.14308] | Column-normalize \(Z_{ij}=e^{q_i^\top k_j}\), then row-normalize | Avoids explaining-away; \(\sum_i \pi_{ij}\ge 1/S\) |
| Sinkformer [2110.11773] | Iterative row/column normalization of \(\exp(S)\) | Wasserstein gradient-flow and heat-diffusion theory |
| Reaction-prediction DS attention [2306.06119] | Two doubly stochastic self-attention mappings via Sinkhorn | Satisfies electron-counting and symmetry rules |
| ESPFormer [2502.07962] | Expected Sliced Transport Plans with SoftSort | Fully parallelizable; no iterative Sinkhorn normalization |
| QR-based DSM [2504.16275] | Thin QR followed by \(A=Q\odot Q\) | Exact DSM |
| QontOT / QDSFormer [2504.16275] | Variational quantum circuit \(U(\theta)\) reshaped into \(P(\theta)\) | First parametric DSM layer |
| LOTFormer [2509.23436] | Two entropic OT problems glued through a learnable pivot | Doubly-stochastic, rank at most \(r\), \(O(nr)\) application |

The reaction-prediction formulation is the clearest example of task-specific structure. There, two doubly stochastic attention maps \(W^+\) and \(W^-\) are used so that the predicted electron redistribution \(\Delta \hat E=W^+-W^-\) 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 \(D\)-sum constraint on the aggregate, and exact Sinkhorn enforcement is preferred to a soft penalty for numerical precision [2306.06119].

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 \(r\ll n\), 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 \(r\). 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 [2502.07962; 2509.23436; 2504.16275].

## 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 \(85.8\%\) to \(88.3\%\), and Point Cloud Transformer mean accuracy from \(92.5\%\) to \(92.7\%\); 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 \(\le 10\%\) on a V100 GPU [2110.11773].

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 \(92.0\%\), \(76.4\%\), and \(93.5\%\), compared with the next-best Softmax values \(91.4\%\), \(75.6\%\), and \(92.0\%\); on Imagenette, Doubly Stochastic attention leads in all three corrupted regimes with \(91.0\%\), \(84.3\%\), and \(99.9\%\), preserving almost its full clean-data accuracy under combined corruption [2507.20453].

The rank-preservation experiments in the 2026 analysis quantify the same trend with residual metrics. For the median normalized residual \(\|\mathrm{res}(P_t)\|_2/\|P_t\|_2\) at depth \(t=10\), AG’s News reports \(0.08\) and \(0.30\) for Softmax without and with skips, versus \(0.15\) and \(0.60\) for Sinkhorn; at end of training, the self-attention output residual on AG’s News is \(0.02\) and \(0.20\) for Softmax, versus \(0.05\) and \(0.45\) 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 [2604.07925].

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 [2306.06119]. 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 \(90.0_{\pm 0.15}\%\) with QontOT versus \(89.1_{\pm 0.73}\%\) with Sinkhorn, and on MedMNIST the mean accuracy over 7 datasets is \(74.3\%\) for QontOT versus \(70.8\%\) for Sinkhorn [2504.16275].

Implementation guidance varies by formulation and task. Reported iteration counts range from a single column-then-row normalization in DNAS, to \(L=3\)–\(5\) Sinkhorn steps in Sinkformers, to \(8\)–\(20\) iterations in the robustness study, and \(30\)–\(50\) 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 \(1.5\)–\(2\times\) 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 [2009.14308; 2110.11773; 2507.20453; 2604.07925].

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 [2502.07962; 2504.16275; 2509.23436; 2604.07925].

Source: https://www.emergentmind.com/topics/doubly-stochastic-attention