arXiv:2606.18283cs.LG2026-06

用概率路由替代注意力计算,实现线性内存的长序列建模。

Gaussian Mixture Attention: Linear-Time Sequence Mixing via Probabilistic Latent Routing

论文配图:Gaussian Mixture Attention: Linear-Time Sequence Mixing via Probabilistic Latent Routing
图 1 · 摘自论文原文
  • 通过学习的高斯混合组件进行隐式路由,避免显式计算全连接注意力矩阵。
  • 内存复杂度降为O(NK),在固定K下实现线性扩展,实测长文本任务表现良好。
  • 模型可解释性强,适合需要高效且透明注意力机制的研究场景。

标准点积注意力的密集词元间交互是扩展Transformer至长序列的核心瓶颈。本文提出高斯混合注意力(Gaussian Mixture Attention, GMA),一种基于概率路由的序列混合机制,将显式的成对查询-键比较替换为通过K个学习得到的高斯混合组件进行路由。查询和键被映射到共享潜在路由空间中的后验责任向量;它们的重叠定义了隐式的责任空间相似度,而值则写入并读取一个包含K个槽位的潜在记忆。利用矩阵乘法的结合律,GMA避免生成所需的N×N相似度矩阵,转而使用两个责任矩阵,其主导激活存储规模为O(NK),而非固定K下的O(N²)。我们提出了双向与因果变体,提供了高斯混合组件的端到端可微参数化,并分析了其责任调制梯度结构、非负低秩相似度解释及局部路由稳定性。实验表明,GMA展现出预期的固定K线性内存扩展,在长序列分类任务上优于基线注意力模型;因果GMA在WikiText-103上优于测试的线性/随机特征注意力变体,但当前实现仍落后于优化的因果SDPA和Mamba。对学习到的责任分析显示各组件被广泛使用,且与表面词元类别有中等程度对齐,支持GMA作为概率性、可解释的固定K线性时间注意力替代方案,而非对优化的softmax注意力或状态空间模型的通用替代。

原文摘要 · Abstract (English)

The dense token-to-token interaction pattern of standard dot-product attention remains a central bottleneck in scaling Transformer architectures to long contexts. We introduce \textbf{Gaussian Mixture Attention (GMA)}, a probabilistic attention-style sequence mixer that replaces explicit pairwise query--key comparison with routing through $K$ learned Gaussian mixture components. Queries and keys are mapped to posterior \textit{responsibility} vectors over a shared latent routing space; their overlap defines an implicit responsibility-space affinity, while values are written into and read from a $K$-slot latent memory. By exploiting the associativity of matrix multiplication, GMA avoids materializing the induced $N\times N$ affinity matrix and instead uses two responsibility matrices whose dominant activation storage scales as $\mathcal{O}(NK)$ rather than $\mathcal{O}(N^2)$ for fixed $K$. We formulate bidirectional and causal variants of GMA, provide an end-to-end differentiable parameterization of the Gaussian mixture components, and analyze its responsibility-modulated gradient structure, constrained non-negative low-rank affinity interpretation, and local routing stability. Empirically, GMA exhibits the intended fixed-$K$ linear memory scaling and is competitive with attention-style baselines on long-context classification, while causal GMA improves over tested linear/random-feature attention variants on WikiText-103 but remains behind optimized causal SDPA and Mamba in the current implementation. Analysis of learned responsibilities further shows broad component usage and moderate alignment with surface-form token categories, supporting GMA as a probabilistic, interpretable, fixed-$K$ linear-time attention-style alternative rather than a universal replacement for optimized softmax attention or state-space models.

注意力机制线性复杂度概率路由长序列建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。