Hierarchical Shift Mixing -- Beyond Dense Attention in Transformers

📄 arXiv: 2601.22852v1 📥 PDF

作者: Robert Forchheimer

分类: cs.LG

发布日期: 2026-01-30

备注: 11 pages, 10 pdf figures

💡 一句话要点

提出分层移位混合（HSM），在Transformer中实现线性复杂度Token混合，提升效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer 注意力机制 线性复杂度 分层混合 长序列建模

📋 核心要点

  1. Transformer的softmax注意力机制计算复杂度高，限制了其在大规模序列上的应用。
  2. 论文提出分层移位混合（HSM）框架，通过分层方式进行token混合，降低计算复杂度。
  3. 实验表明，HSM变体性能接近softmax注意力，混合架构优于GPT基线，并降低计算成本。

📝 摘要（中文）

随着Transformer架构在大语言模型中的应用，基于softmax的注意力层因其二次时间计算复杂度而面临越来越多的审查。许多尝试用复杂度较低的方法替代它，但通常以性能下降为代价。我们引入了分层移位混合（HSM），这是一个通用的token混合框架，它将成对token交互分布在Transformer层中，而不是在每一层中密集地计算它们。HSM实现了线性时间复杂度，同时保持了对特定混合函数的不可知性。我们表明，即使是简单的HSM变体也能达到接近softmax注意力的性能，并且将HSM与softmax注意力相结合的混合架构可以优于GPT风格的Transformer基线，同时降低训练和推理期间的计算成本。

🔬 方法详解

问题定义：Transformer模型中的自注意力机制具有二次方的时间复杂度，这限制了其在处理长序列时的效率。现有的替代方案通常以牺牲模型性能为代价来降低计算复杂度。因此，如何在保持甚至提升模型性能的同时，降低Transformer的计算复杂度是一个关键问题。

核心思路：论文的核心思路是将token之间的交互分散到不同的Transformer层中进行，而不是像传统自注意力机制那样在每一层都进行密集的两两交互。通过分层移位混合（HSM），每个token只与部分token进行交互，从而降低了计算复杂度。

技术框架：HSM框架包含多个Transformer层，每一层都包含一个混合函数。该混合函数负责将输入token进行混合，生成新的token表示。关键在于，不同层之间的混合方式不同，每一层只关注部分token之间的交互。通过多层的分层混合，最终实现所有token之间的有效交互。整体流程是从输入序列开始，经过多层HSM处理，最终输出序列表示。

关键创新：最重要的创新点在于将token混合操作从单层密集计算转变为多层分层计算。这种分层方式使得计算复杂度从二次方降低到线性级别，同时保持了模型性能。HSM框架对具体的混合函数具有不可知性，可以灵活地选择不同的混合函数。

关键设计：HSM的关键设计包括分层策略和混合函数选择。分层策略决定了每一层token交互的范围和方式。混合函数则负责具体的token混合操作，可以选择简单的加权平均，也可以选择更复杂的神经网络。论文中提到，即使使用简单的混合函数，HSM也能取得不错的效果。此外，论文还探索了HSM与softmax注意力混合使用的架构，以进一步提升性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，简单的HSM变体可以达到接近softmax注意力的性能。更重要的是，将HSM与softmax注意力相结合的混合架构，在降低计算成本的同时，优于GPT风格的Transformer基线。这些结果表明HSM是一种有效的token混合方法，可以在保持甚至提升模型性能的同时，降低计算复杂度。

🎯 应用场景

该研究成果可应用于各种需要处理长序列数据的场景，例如自然语言处理中的长文本建模、语音识别中的长语音建模、以及基因序列分析等。通过降低计算复杂度，HSM使得Transformer模型能够处理更长的序列，从而提升模型性能和应用范围。此外，该方法还可以降低训练和推理成本，使得大规模模型的部署更加可行。

📄 摘要（原文）

Since the introduction of the Transformer architecture for large language models, the softmax-based attention layer has faced increasing scrutinity due to its quadratic-time computational complexity. Attempts have been made to replace it with less complex methods, at the cost of reduced performance in most cases. We introduce Hierarchical Shift Mixing (HSM), a general framework for token mixing that distributes pairwise token interactions across Transformer layers rather than computing them densely within each layer. HSM enables linear-time complexity while remaining agnostic to the specific mixing function. We show that even simple HSM variants achieve performance close to softmax attention, and that hybrid architectures combining HSM with softmax attention can outperform a GPT-style Transformer baseline while reducing computational cost during both training and inference.