Gauge-Free Inverse OT Attention
Implementation & benchmark of arXiv:2609.01278 — An exact and fast solution of the inverse Regularized Optimal Transport problem
Source paper: An exact and fast solution of the inverse Regularized Optimal Transport problem arXiv:2609.01278 ⓘ · analyzed Sep 2, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace ordinary row-softmax attention with a doubly stochastic Sinkhorn attention plan W, and periodically recover a gauge-fixed pairwise cost from W using the exact double-centering identity. Use this recovered cost to initialize or regularize a structured attention score, making the attention geometry identifiable despite arbitrary query and key row and column offsets.
Formulas
Mathematical statement
For source and target marginals s and r, entropic optimal transport has coupling W = diag(u) K diag(v), where K_ij = exp(-C_ij divided by epsilon), C is the cost matrix, epsilon is the entropic temperature, and u and v are Sinkhorn scaling vectors. Consequently, log W = log(u) 1-transpose + 1 log(v)-transpose - C divided by epsilon. Let H_m = I_m - 1 1-transpose divided by m and H_n = I_n - 1 1-transpose divided by n be centering matrices. Double-centering removes the row and column gauge terms, yielding H_m log(W) H_n = -H_m C H_n divided by epsilon. The recoverable gauge-fixed cost is therefore C_hat = -epsilon H_m log(W plus delta) H_n, where delta is a positive numerical floor. Costs differing by a 1-transpose plus 1 b-transpose generate the same transport plan, so only the double-centered cost is identifiable.
Implementation notes
Integrate this into a cross-attention block. Convert attention scores S into a positive kernel K with K_ij = exp(-S_ij divided by epsilon), then run 5 to 20 log-domain Sinkhorn iterations to obtain W with prescribed source and target marginals. Keep value aggregation Y = W V, but add an inverse-cost consistency term or use the recovered cost as a structured bias for the next attention block. The pseudocode is: compute stabilized K; initialize positive u and v; repeat u = s divided elementwise by K v and v = r divided elementwise by K-transpose u; form W = diag(u) K diag(v); set L = log(W plus delta); compute C_inv = -epsilon H_m L H_n; add lambda times the squared Frobenius norm of H_m(S - C_inv)H_n to the training loss. The exact inverse identity is supplied by the paper. The engineer must estimate the effect of finite Sinkhorn convergence, floating-point flooring, and model mismatch. First run a synthetic test with random C, prescribed marginals, and generated W before training a two-layer Transformer on a small copy or patch-classification task. Compare softmax attention with Sinkhorn attention with and without the inverse-cost penalty. The prediction is that the relative centered reconstruction error will remain below 5 percent when the smallest W entry exceeds 10^-5 and Sinkhorn marginal error is below 10^-4. Adding any row-plus-column gauge to C must not change C_inv. Reconstruction error should increase sharply when W entries fall below delta.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: No measurable effect
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 a reproducible NumPy MVP for log-domain Sinkhorn attention and gauge-free inverse-cost recovery. On a 9x7 synthetic OT plan, centered reconstruction error was 7.0e-16, gauge perturbation changed the plan by 3.6e-16 relative and the recovered cost by 7.6e-16; flooring increased error from 4.3e-10 at delta=1e-12 to 0.227 at delta=1e-3. The attention toy benchmark did not show a task win: row-softmax retrieval accuracy was 1.000 versus 0.648 for Sinkhorn, and the inverse penalty was evaluated rather than optimized, so the claimed training benefit remains unverified.
- Agent confidence
- 8/10
- Baseline
- Row-softmax retrieval accuracy 1.000; mean row entropy 2.105
- Idea
- Sinkhorn retrieval accuracy 0.6475; mean entropy 0.348; centered inverse recovery 7.02e-16 after 300 iterations; gauge-plan relative change 3.63e-16; gauge-inverse relative change 7.60e-16
Limitations: Only synthetic NumPy experiments were run. No learned Transformer, backpropagated inverse-cost regularizer, FLOP/speed measurement, noisy-plan study, CUDA implementation, or realistic copy/patch-classification task was tested; the penalty branch is currently diagnostic rather than trained.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2609.01278)
No significant difference vs. the tuned baseline (+0.9%, p=0.7557, wins 4 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.9269
- Idea mean
- 0.9357
- Effect (Δ)
- 0.0087 (+0.9%; negative = idea better)
- Wins
- 4 / 8 paired seeds
- p-value
- 0.7557 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±6.3%
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
The inverse-cost consistency penalty and recovered cost were not used as a learned auxiliary loss or next-block bias; this benchmark tested the core Sinkhorn attention replacement only. The local loop was used because the intervention changes the attention forward pass. No FLOP or wall-clock comparison, larger datasets, alternative sequence lengths, or multi-head implementation was tested.
How to run: python3 bench_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Built an end-to-end Gauge-Free Inverse OT Attention transformer by replacing row-softmax attention with differentiable log-domain Sinkhorn attention, while keeping the transformer architecture, Adam optimizer, batch size, epochs, and paired seeds matched. On the corrected benchmark, baseline test MSE was 0.92694 and the idea was 0.93569, giving delta_mean +0.00875 with permutation p=0.75565; therefore there is no significant win. The trained-model mechanism signature also did not confirm the predicted marginal error below 1e-4: observed row error was 0.00786, although column error was 1.13e-8.
- Baseline
- Sequence/transformer_tiny, tuned lr=0.006; full 8-seed test MSE mean 0.9269421473, std 0.0813036234; per-seed [0.9268826246, 0.8607907891, 0.8572450280, 0.8207906485, 1.0688729286, 0.9266430140, 1.0371272564, 0.9171848893]
- Idea
- Sinkhorn attention, best tested eps=0.5 at lr=0.006; full 8-seed test MSE mean 0.9356912449, std 0.0569688326; per-seed [0.9681993127, 0.8495118022, 0.9008096457, 0.9619098306, 0.9716687799, 0.9182450175, 1.0382701159, 0.8769154549]
- Mechanism signature
- Predicted maximum row/column marginal error <1e-4; observed mean row error 0.0078572490, observed mean column error 1.1294106e-8, mean minimum attention entry 1.0011e-14, confirmed=false
- Paired delta
- Mean idea-baseline MSE delta +0.0087490976; idea wins 4/8 pairs; permutation p-value 0.75565
Benchmark result
No significant difference vs. the tuned baseline (+0.9%, p=0.7557, wins 4 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.9269
- Idea mean
- 0.9357
- Effect (Δ)
- 0.0087 (+0.9%; negative = idea better)
- Wins
- 4 / 8 paired seeds
- p-value
- 0.7557 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±6.3%
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
The inverse-cost consistency penalty and recovered cost were not used as a learned auxiliary loss or next-block bias; this benchmark tested the core Sinkhorn attention replacement only. The local loop was used because the intervention changes the attention forward pass. No FLOP or wall-clock comparison, larger datasets, alternative sequence lengths, or multi-head implementation was tested.
How to run: python3 bench_experiment.py
Artifacts
- bench_experiment.py 6.0 KB View Raw
- bench_report.json 2.5 KB View Raw
- experiment.py 3.4 KB View Raw
- report.md 1.4 KB View
- report_bench_2026-09-02T171818.md 4.9 KB View
- results.json 1.9 KB View Raw