EPS Dashboard

Authors: Dylan Forde

Summary

This paper optimizes a differentiable attention mechanism based on optimal transport (OT)—essentially balancing attention weights using an iterative Sinkhorn solver—for TPU hardware. The authors freeze the first T steps of Sinkhorn iteration and only backpropagate through a short R=2 refinement tail, which keeps memory usage tractable. They prove an exact block-wise gradient schedule that costs O((T+R)LW) with only O(Ld) input storage, and show the method trains successfully on protein sequence data (Pfam), improving reconstruction and loss over a three-hour TPU run.

Main takeaways:

Standard attention can be replaced by balanced optimal transport, which spreads attention mass more evenly using a Sinkhorn solver.
Training through the full Sinkhorn solver is expensive; stopping gradients after T steps and differentiating only a short R=2 tail makes it tractable.
The backward pass decomposes into four explicit matrix factors, enabling a memory-efficient block-wise schedule on TPU hardware.
On protein sequences (Pfam), the method trains end-to-end and improves metrics over baseline, validating the approach.
The "dustbin" mechanism—adding an extra dimension to handle gaps or out-of-distribution tokens—lifts cleanly into the same framework.

Relevance

Not related to my persona or midtraining work—this is a low-level attention architecture optimization for TPUs. Included because it's a novel attention mechanism, but it doesn't touch behavioral installation, steering, or prompt engineering.

Threat model

Potential threat/caveat for clean result "Stretching turn count, completion length, or system-prompt length at train time fails to amplify marker uptake; the longest system prompt instead leaks across bystander personas (LOW confidence)": this item discusses bias.

Abstract

arXiv:2605.08123v1 Announce Type: new Abstract: We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped $T$-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the production $R=2$ case, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the $R=2$ score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost $O((T+R)LW)$, $O(Ld)$ input storage, and $O(L)$ additional HBM usage for fixed head dimension $d$ and band width $W$. We also formalize the current \texttt{dustbin_block} path as the same balanced surrogate on an augmented support, so the schedule lifts to the gap-aware transport path used in our TPU runs. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within $10^{-5}$--$10^{-10}$. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced $R=2$ run sustains roughly $8.5$ examples per second through a three-hour budget, reaching step $1437$. Held-out Pfam test shards improve reconstruction from $3.17$ to $0.99$ and sparse CE from $5.86$ to $5.69$ relative to step $0$. These results support exact fixed-depth backward theory, a theorem-matching gap-aware bridge, and trainability evidence for the production path.