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.