△ Mechanism confirmed, baseline not beaten 2026

Exact doubly stochastic low-rank attention

Implementation & benchmark of arXiv:2608.08642 — Exact Rank-Space KL Projection for Shared-Marginal Low-Rank Factors: Application to Doubly Stochastic Clustering

Usefulness7/10

Difficulty6/10

Novelty7/10

Architecture Memory Scalability Stability Attention Transformer Memory Optimal transport Convex analysis Information theory Linear algebra

Source paper: Exact Rank-Space KL Projection for Shared-Marginal Low-Rank Factors: Application to Doubly Stochastic Clustering arXiv:2608.08642 ⓘ · analyzed Sep 1, 2026

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

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.

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.$$

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.

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

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Inconclusive

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

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.

Agent confidence: 8/10

Baseline: Explicit factorized W application: 1.282 ms median and 262144 attention-state entries at n=512,r=16; no end-to-end standard softmax training benchmark.
Idea: Thin U((V^T X)/g) application: 0.249 ms median, 5.16x faster; 16400 factor entries, 15.98x lower state storage; output relative error versus explicit factorized W 6.42e-16; doubly-stochastic residual 7.46e-14.

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.

How to run: python3 verify.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.08642) Sep 2, 2026

No significant difference vs. the tuned baseline (+11.2%, p=0.0867, wins 2 / 8 paired seeds).

Benchmark: Sequence modeling
Model: transformer_tiny
Paired seeds: 8
Baseline mean: 0.4481
Idea mean: 0.4981
Effect (Δ): 0.05 (+11.2%; negative = idea better)
Wins: 2 / 8 paired seeds
p-value: 0.0867 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±12.7%
Mechanism: Confirmed ✓
Practical effect: Inconclusive
Baseline tuning: swept over 3 configs

Limitations:

The experiment used the small standard sequence track budget of 400 training samples, 200 test samples, and 8 epochs rather than WikiText-103 or a full-scale Transformer. Only rank 8 was tested, projection overhead and peak GPU memory were not separately timed, and the factor-state ratio was measured at sequence length 32 rather than as an end-to-end peak-memory measurement.

How to run: python3 bench_ds_attention.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Failed ✗

Agent confidence: 9/10

Implemented and ran the exact doubly stochastic low-rank attention variant on the structurally matched registered sequence track. The mechanism signature was confirmed from trained-model behavior, but test MSE was worse than baseline (0.498095 vs 0.448069; paired delta +0.050026; permutation p=0.08665), so the idea did not produce a significant win.

Baseline: Mean test MSE 0.44806899 over 8 seeds; baseline sweep selected lr=0.001 from {0.001, 0.003, 0.006}.
Idea: Rank-8 low-rank doubly stochastic attention mean test MSE 0.49809538 over 8 seeds; idea wins 2/8 paired seeds.

Benchmark result