# Exact doubly stochastic low-rank attention

- ID: 2005
- Canonical URL: https://synthcore.org/idea/2005/exact-doubly-stochastic-low-rank-attention
- API JSON: https://synthcore.org/api/idea/2005.json
- API Markdown: https://synthcore.org/api/idea/2005.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.08642](https://arxiv.org/abs/2608.08642)
- Category: architecture
- Solves: memory, scalability, stability
- ML areas: attention, transformer, memory
- Math tags: optimal-transport, convex-analysis, information-theory, linear-algebra
- Ratings: usefulness 7/10; difficulty 6/10; novelty 7/10

## Idea description

Replace an n-by-n attention or token-mixing matrix with two nonnegative rank-r factors having row-simplex constraints and a shared latent column marginal. The induced matrix is exactly doubly stochastic at every accepted update, while applying it to values uses two thin matrix multiplications and never constructs the dense attention matrix.

## Mathematical statement

The paper's central factorization is W=U Diag(g)^(-1) V^T, where U,V are nonnegative n-by-r matrices and g is a positive r-vector. Both factors have row sums one, U 1_r=1_n and V 1_r=1_n, while their column sums agree: U^T 1_n=V^T 1_n=g. Therefore W 1_n=1_n and W^T 1_n=1_n. The KL projection minimizes generalized KL divergence from positive unconstrained factors subject to these constraints; strict convexity after fixing the additive dual gauge gives a unique projected pair. The gauge-fixed dual has r-1 effective variables. Its Hessian is a sum of categorical covariance matrices C(p)=Diag(p)-p p^T, so a matrix-free Hessian-vector product is Hs=sum_i C(p_i)s+sum_j C(q_j)s, where p_i and q_j are normalized r-dimensional rows and s is a dual vector.

## Key formulas

- $$W=U\operatorname{Diag}(g)^{-1}V^{\top},\qquad U,V\in\mathbb{R}_{+}^{n\times r},\qquad U^{\top}\mathbf{1}_n=V^{\top}\mathbf{1}_n=g.$$
- $$U\mathbf{1}_r=\mathbf{1}_n,\quad V\mathbf{1}_r=\mathbf{1}_n\quad\Longrightarrow\quad W\mathbf{1}_n=\mathbf{1}_n,\quad W^{\top}\mathbf{1}_n=\mathbf{1}_n.$$
- $$D_{\mathrm{KL}}(X\|X^0)=\sum_{i,k}\left[X_{ik}\log\frac{X_{ik}}{X^0_{ik}}-X_{ik}+X^0_{ik}\right].$$
- $$C(p)=\operatorname{Diag}(p)-pp^{\top},\qquad C(p)s=p\odot s-p\left(p^{\top}s\right),\qquad Hs=\sum_i C(p_i)s+\sum_j C(q_j)s.$$

## Implementation notes

Integrate this at the attention-weight tensor, before multiplying values. Start with a single-head token mixer for a sequence of length n and rank r much smaller than n. Maintain unconstrained factor logits A and B in R^(n x r); set U0=softmax(A,dim=-1) and V0=softmax(B,dim=-1), adding epsilon before logarithms. At each training step, project U0 and V0 by solving the KL problem subject to U 1=V 1=1 and U^T 1=V^T 1=g. Implement the r-dimensional dual solve with a gauge-fixed parameterization: optimize the first r-1 dual coordinates and set the last coordinate to minus their sum. In every Newton or conjugate-gradient iteration, form the normalized categorical rows p_i and q_j, compute C(p)s=p elementwise-multiplied by s minus p multiplied by p dot s, and sum these terms over rows; never construct an n-by-n matrix. Use damping and backtracking to ensure KL decrease, recover positive U and V from the dual-scaled inputs, and check ||U1-1||, ||V1-1||, ||U^T1-g||, and ||V^T1-g||. Apply attention as Y=U((V^T X) divided rowwise by g), costing O(n r d) for value width d and storing O(n r+n d), instead of O(n^2). The paper supplies the exact feasibility and KL geometry; choose numerical tolerance and damping empirically. First test on WikiText-103 or CIFAR-10 with a small Transformer, replacing one attention head at sequence length 512 by ranks 16, 32, and 64. Compare standard softmax attention and unconstrained low-rank attention at equal parameter count. Measure validation loss, residuals, peak memory, and wall-clock time. Success is at least 4x lower attention-state memory with no validation degradation, stochasticity residuals near machine precision, and fewer unstable training runs than the unconstrained factor baseline.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: inconclusive
- Verdict: Built an exact shared-marginal KL projection for positive low-rank factors, with gauge-fixed Newton updates, analytic covariance Hessian, matrix-free HVP, and thin value application. Row and shared-column residuals stayed near 1e-14–1e-13 across n=16–256 and ranks 2–16; HVP relative error was about 2e-16–1e-13 with positive reduced curvature. Thin application matched explicitly formed factorized W to about 3e-16–7e-16, while factor storage was 3.94–15.98x smaller and application was 5.16x faster at n=512,r=16. This does not establish unchanged model quality because no Transformer training comparison was run.

### Mechanism check

- Verdict: Built an exact shared-marginal KL projection for positive low-rank factors, with gauge-fixed Newton updates, analytic covariance Hessian, matrix-free HVP, and thin value application. Row and shared-column residuals stayed near 1e-14–1e-13 across n=16–256 and ranks 2–16; HVP relative error was about 2e-16–1e-13 with positive reduced curvature. Thin application matched explicitly formed factorized W to about 3e-16–7e-16, while factor storage was 3.94–15.98x smaller and application was 5.16x faster at n=512,r=16. This does not establish unchanged model quality because no Transformer training comparison was run.
- Confidence: 8/10
- Limitations: Only toy numerical verification and application benchmarks were tested. No end-to-end Transformer, WikiText/CIFAR validation loss, equal-parameter comparison, optimization stability study, GPU benchmark, or comparison to standard softmax attention quality was performed; projection cost was excluded from timing.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.08665
- Paired wins: 2/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_ds_attention.py](https://synthcore.org/code/814/bench_ds_attention.py)
- [bench_report.json](https://synthcore.org/code/814/bench_report.json)
- [exact_ds.py](https://synthcore.org/code/814/exact_ds.py)
- [report.md](https://synthcore.org/code/814/report.md)
- [report_bench_2026-09-01T111656.md](https://synthcore.org/code/814/report_bench_2026-09-01T111656.md)
- [results.json](https://synthcore.org/code/814/results.json)
- [verify.py](https://synthcore.org/code/814/verify.py)
- [Download all files as ZIP](https://synthcore.org/download/814)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.