LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport

• URL: http://arxiv.org/abs/2509.23436v1
• Date: Sat, 27 Sep 2025 18:11:09 GMT
• Title: LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport
• Authors: Ashkan Shahbazi, Chayne Thrash, Yikun Bai, Keaton Hamm, Navid NaderiAlizadeh, Soheil Kolouri,
• Abstract summary: We propose a principled attention mechanism that is simultaneously linear-time and doubly-stochastic.<n>LotFormer achieves state-of-the-art results on the Long Range Arena benchmark.
• Score: 21.50165411149415
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformers have proven highly effective across a wide range of modalities. However, the quadratic complexity of the standard softmax attention mechanism poses a fundamental barrier to scaling them to long context windows. A large body of work addresses this with linear attention, which reformulates attention as a kernel function and approximates it with finite feature maps to achieve linear-time computation. Orthogonal to computational scaling, most attention mechanisms -- both quadratic and linear -- produce row-normalized maps that can over-focus on a few tokens, degrading robustness and information flow. Enforcing doubly-stochastic attention alleviates this by balancing token participation across rows and columns, but existing doubly-stochastic attention mechanisms typically introduce substantial overhead, undermining scalability. We propose LOTFormer, a principled attention mechanism that is simultaneously linear-time and doubly-stochastic. Our approach exploits the connection between attention maps and transportation plans between query and key measures. The central idea is to constrain the transport plan to be low-rank by conditioning it on a learnable pivot measure with small support. Concretely, we solve two entropic optimal transport problems (queries $\to$ pivot and pivot $\to$ keys) and compose them into a conditional (glued) coupling. This yields an attention matrix that is provably doubly-stochastic, has rank at most $r \ll n$, and applies to values in $O(nr)$ time without forming the full $n \times n$ map. The pivot locations and masses are learned end-to-end. Empirically, LOTFormer achieves state-of-the-art results on the Long Range Arena benchmark, surpassing prior linear and transport-based attention methods in both accuracy and efficiency.

Related papers

• LUNA: Linear Universal Neural Attention with Generalization Guarantees [27.74721677870656]
  textscLuna achieves state-of-the-art average accuracy among efficient Transformers under compute parity.<n>textscLuna also excels at post-hoc conversion: replacing softmax in fine-tuned BERT and ViT-B/16 checkpoints.
  arXiv  Detail & Related papers  (2025-12-08T21:49:55Z)
• Higher-order Linear Attention [59.92962330635185]
  quadratic cost of scaled dot-product attention is a central obstacle to scaling autoregressive language models to long contexts.<n>We introduce Higher-order Linear Attention (HLA), a causal, streaming mechanism that realizes higher interactions via compact prefix sufficient statistics.
  arXiv  Detail & Related papers  (2025-10-31T07:54:37Z)
• Log-Linear Attention [81.09631871212211]
  This paper develops log-linear attention, an attention mechanism that balances linear attention's efficiency and the expressiveness of softmax attention.<n>We show that with a particular growth function, log-linear attention admits a similarly matmul-rich parallel form whose compute cost is log-linear in sequence length.<n>Log-linear attention is a general framework and can be applied on top of existing linear attention variants.
  arXiv  Detail & Related papers  (2025-06-05T08:44:51Z)
• ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans [13.695885742446027]
  Self-attention can lead to over-concentration on a few tokens during training, resulting in suboptimal information flow.<n>We introduce a novel, fully parallelizable doubly-stochastic attention mechanism based on sliced optimal transport.<n>Our method enforces doublyity without iterative Sinkhorn normalization, significantly enhancing efficiency.
  arXiv  Detail & Related papers  (2025-02-11T21:20:48Z)
• Bridging the Divide: Reconsidering Softmax and Linear Attention [116.34723260730405]
  We present two key perspectives to understand and alleviate the limitations of linear attention.<n>We prove that linear attention is not injective, which is prone to assign identical attention weights to different query vectors.<n> Secondly, we confirm that effective local modeling is essential for the success of Softmax attention, in which linear attention falls short.
  arXiv  Detail & Related papers  (2024-12-09T15:44:22Z)
• Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences [60.489682735061415]
  We propose CHELA, which replaces state space models with short-long convolutions and implements linear attention in a divide-and-conquer manner. Our experiments on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.
  arXiv  Detail & Related papers  (2024-06-12T12:12:38Z)
• Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models [20.78813311569383]
  We present Lightning Attention, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. Specifically, we utilize the conventional attention mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. Various experiments are conducted on different model sizes and sequence lengths.
  arXiv  Detail & Related papers  (2024-01-09T16:27:28Z)
• SEA: Sparse Linear Attention with Estimated Attention Mask [51.22399593954608]
  Long seqeuences pose a problem due to the quadratic complexity of the attention operation. Previous research has aimed to lower the complexity by sparsifying or linearly approximating the attention matrix. We propose SEA: Sparse linear attention with an Estimated Attention mask.
  arXiv  Detail & Related papers  (2023-10-03T03:56:26Z)
• Combiner: Full Attention Transformer with Sparse Computation Cost [142.10203598824964]
  We propose Combiner, which provides full attention capability in each attention head while maintaining low computation complexity. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach.
  arXiv  Detail & Related papers  (2021-07-12T22:43:11Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.