提出线性时间双随机注意力机制,提升长序列建模效率与稳定性。
LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport
- 基于最优传输视角,用可学习枢轴点构造低秩耦合
- 在视觉与文本任务中实现高精度与低计算开销的平衡
- 适合需要长序列处理的高效模型设计场景
Transformer 在多模态任务中表现优异,但标准 softmax 注意力随序列长度呈二次增长,限制了长序列建模。线性注意力通过核特征映射近似注意力,但多数方法仍为行归一化,易集中在少数标记,影响鲁棒性与信息流动。双随机注意力通过平衡行与列的参与度来缓解此问题,但现有方法常带来显著开销。我们提出 LOTFormer,从注意力作为查询与键测度间耦合的最优传输视角出发,通过条件于支持集小的可学习枢轴测度,强制实现低秩传输计划。求解两个熵正则化传输问题:查询到枢轴、枢轴到键,并组合成一个条件耦合,该耦合被证明是双随机的,秩至多为 $r \ll n$,且在 $O(nr)$ 时间内对值进行操作,无需构建完整的 $n \times n$ 矩阵。枢轴位置与质量端到端学习。在视觉与文本基准上,将 LOTFormer 插入 Swin、DeiT 与 BERT 等标准骨干网络,均展现出优越的准确率-效率权衡。
原文摘要 · Abstract (English)
Transformers have proven highly effective across modalities, but standard softmax attention scales quadratically with sequence length, limiting long context modeling. Linear attention mitigates this by approximating attention with kernel feature maps, yet most attention mechanisms remain row normalized and can over concentrate mass on a few tokens, harming robustness and information flow. Doubly stochastic attention counteracts this by balancing token participation across both rows and columns, but existing approaches often add significant overhead. We propose LOTFormer, a linear time doubly stochastic attention mechanism derived from an optimal transport view of attention as a coupling between query and key measures. LOTFormer enforces a low rank transport plan by conditioning on a learnable pivot measure with small support. We solve two entropic transport problems, queries to pivot and pivot to keys, and compose them into a conditional coupling 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$ matrix. The pivot locations and masses are learned end-to-end. Across vision and text benchmarks, LOTFormer delivers strong accuracy efficiency tradeoffs when plugged into standard backbones including Swin, DeiT, and BERT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。