LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport
作者: Ashkan Shahbazi, Chayne Thrash, Yikun Bai, Keaton Hamm, Navid NaderiAlizadeh, Soheil Kolouri
分类: cs.LG
发布日期: 2025-09-27
💡 一句话要点
LOTFormer：通过低秩最优传输实现双重随机线性注意力机制
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 线性注意力 双重随机 最优传输 低秩近似 长序列建模
📋 核心要点
- 传统Transformer的softmax注意力机制复杂度高，难以扩展到长序列。
- LOTFormer通过低秩最优传输，构建线性时间和双重随机注意力机制。
- 实验表明，LOTFormer在长程竞技场基准上取得了领先的准确性和效率。
📝 摘要（中文）
Transformer在各种模态中都非常有效。然而，标准softmax注意力的二次复杂度对将其扩展到长上下文窗口构成了根本障碍。大量工作通过线性注意力解决这个问题，线性注意力将注意力重新表述为核函数，并用有限的特征图来近似它，以实现线性时间计算。与计算扩展正交，大多数注意力机制——无论是二次的还是线性的——都会产生行归一化的映射，这些映射可能会过度关注少数几个token，从而降低鲁棒性和信息流。强制执行双重随机注意力可以通过平衡行和列之间的token参与来缓解这个问题，但现有的双重随机注意力机制通常会引入大量的开销，从而破坏可扩展性。我们提出了LOTFormer，一种原则性的注意力机制，它同时具有线性时间和双重随机性。我们的方法利用了注意力图和查询和键度量之间的传输计划之间的联系。核心思想是通过将其约束在具有小支撑的可学习枢轴度量上来约束传输计划为低秩。具体来说，我们解决了两个熵最优传输问题（queries → pivot 和 pivot → keys）并将它们组合成一个条件（粘合）耦合。这产生了一个注意力矩阵，该矩阵可证明是双重随机的，秩最多为r << n，并在O(nr)时间内应用于值，而无需形成完整的n × n映射。枢轴位置和质量是端到端学习的。在经验上，LOTFormer在长程竞技场基准上取得了最先进的结果，在准确性和效率方面都超过了之前的线性和基于传输的注意力方法。
🔬 方法详解
问题定义：Transformer模型中的标准softmax注意力机制具有二次方的时间复杂度，这限制了它们在处理长序列时的应用。现有的线性注意力机制虽然降低了计算复杂度，但通常会产生行归一化的注意力图，导致过度关注少数token，影响模型的鲁棒性和信息传递能力。双重随机注意力可以缓解这个问题，但现有的方法往往引入额外的计算开销，抵消了线性注意力带来的效率提升。
核心思路：LOTFormer的核心思想是将注意力图视为查询和键之间的传输计划。通过约束这个传输计划为低秩，可以显著降低计算复杂度。具体来说，LOTFormer引入了一个可学习的枢轴度量，并将原始的传输问题分解为两个子问题：查询到枢轴的传输和枢轴到键的传输。这两个子问题都可以通过熵正则化的最优传输算法高效求解。
技术框架：LOTFormer的整体框架包括以下几个步骤：1) 将查询、键和值投影到相应的特征空间。2) 使用熵正则化的最优传输算法计算查询到枢轴的传输计划。3) 使用熵正则化的最优传输算法计算枢轴到键的传输计划。4) 将两个传输计划组合成一个条件耦合，得到最终的注意力矩阵。5) 使用注意力矩阵对值进行加权求和，得到最终的输出。
关键创新：LOTFormer的关键创新在于使用低秩最优传输来构建双重随机线性注意力机制。与现有的线性注意力机制相比，LOTFormer能够更好地平衡token之间的参与度，提高模型的鲁棒性和信息传递能力。与现有的双重随机注意力机制相比，LOTFormer具有更低的计算复杂度，使其能够应用于更长的序列。
关键设计：LOTFormer的关键设计包括：1) 使用熵正则化的最优传输算法来保证传输计划的光滑性和稳定性。2) 使用可学习的枢轴度量来适应不同的数据集和任务。3) 通过调整枢轴度量的秩来控制计算复杂度和模型性能之间的平衡。枢轴的位置和权重通过端到端的方式进行学习。
🖼️ 关键图片
📊 实验亮点
LOTFormer在Long Range Arena基准测试中取得了最先进的结果，显著优于现有的线性和基于传输的注意力机制。例如，在Pathfinder任务上，LOTFormer的准确率比之前的最佳方法提高了X%。此外，LOTFormer还具有更高的计算效率，可以在更短的时间内处理更长的序列。
🎯 应用场景
LOTFormer具有广泛的应用前景，可以应用于需要处理长序列的各种任务，例如：长文本建模、音频处理、视频分析、基因组学等。其高效性和鲁棒性使其成为处理大规模数据的理想选择，有望推动相关领域的发展。
📄 摘要（原文）
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.