基于多头门控的多频段残差融合技术，用于提升Transformer语言模型性能

《Neurocomputing》：Multi-band residual fusion with per-head gating for enhanced transformer language models

　　黄志高|潘淼|李全发 中国泉州362000，泉州师范学院福建高等教育信息功能材料重点实验室 摘要：当前Transformer注意力机制主要基于逐令牌的注意力输出进行运算，在注意力聚合后可能会丢失有用的局部序列结构。我们提出了带每头门控的多频段残差融合机制（MBF），这是一种具有频率分析解释能力的可学习因果多滤波模块。MBF为每个注意力头输出应用多个因果一维卷积滤波器分支，并通过可学习的每头门控将它们融合在一起，同时保留原有的注意力路径。在GPT-2小型和GPT-2中型模型上对Penn Treebank、WikiText-2和WikiText-103数据集进行的实验显示，困惑度有所降低（在PTB数据集上最高可降低10%）。控制实验表明，可学习的MBF滤波器的效果最为显著，而定量滤波器分析则表明，其支持的机制是基于最近注意力输出的可学习门控因果滤波，而非纯粹的频谱分解。部署诊断进一步指出，当前实现的成本主要来自临时张量内存分配和启动开销，这为采用融合机制及量化感知实现提供了依据。引言：基于Transformer的语言模型彻底改变了自然语言处理领域，在各种任务中取得了前所未有的性能[1]、[2]、[3]。这些模型的成功在很大程度上得益于它们的自注意力机制，该机制使模型能够捕捉序列数据中的长距离依赖关系和复杂模式。然而，标准的注意力模块并未包含显式的注意力后因果滤波器组，因此在注意力聚合后，有用的局部序列结构可能仍会被低估。可学习因果多滤波处理的理念源自信号处理理论，其中滤波器组为分析序列信号提供了有效的工具[4]。在语言模型中，注意力输出可能包含变化缓慢和变化迅速的局部模式，但我们并不认为频谱范围与具体语言类别之间存在一一对应关系。这一理念促使我们为注意力输出添加多个可学习的因果滤波器分支，并将频率分析作为解释工具，而非一种确定的机制。尽管Transformer架构取得了成功，但在优化语言建模的注意力机制方面仍存在若干挑战。首先，标准注意力机制在聚合后会统一处理所有注意力头，没有为每个注意力头提供通过不同因果滤波器混合最近注意力状态的有效方式。其次，注意力计算的复杂度会随序列长度呈二次方增长，这就需要更高效的方法，在尽可能减少计算成本的同时保持或提升性能。第三，现有的注意力机制可能在将长距离注意力与局部注意力后变换结合到单个紧凑模块中时面临困难。鉴于这些局限性，本文提出了带每头门控的多频段残差融合机制，这是一种新的架构，它通过可学习的因果多滤波处理和每头门控来增强标准注意力输出。与那些替代注意力机制的方法（如FNet）或主要追求计算效率的方法（如FlashAttention、Linformer、Performer）不同，我们的方法通过让输出经过多个可学习的因果滤波器分支进行处理，从而提升了注意力的表征能力，使每个注意力头能够在不同的序列模式下实现专业化。我们的主要贡献如下：•新颖的架构设计：我们为注意力头输出引入了多频段滤波模块，该模块通过多个可学习的一维因果卷积滤波器分支来处理每个注意力头的表示结果（默认分支数和核大小为）。•自适应的每头专业化处理：我们开发了一种门控机制，使得每个注意力头能够学习如何不同程度地使用各个因果滤波器分支，从而实现在不同上下文规模和模式粒度下的自动专业化处理。•残差融合策略：我们设计了一种合理的残差连接方案，在添加带门控的因果滤波器变换的同时保留原有的注意力信息，从而确保训练稳定性并避免性能下降。•全面的实证验证：我们在PTB、WikiText-2和WikiText-103数据集上提供了大量的实验证据，包括在GPT-2中型模型上的扩展实验结果、经过多重检验校正后的显著性分析（Holm/FDR）、效应量报告、因果滤波器对照实验，以及利用内存效率高的注意力内核进行的正交性检测。•部署诊断：我们分析了理论成本、实际吞吐量、导致延迟与FLOPs不匹配的软件层面原因、torch.compile的行为表现、SDPA后端兼容性、原生低位宽部署的兼容性，以及FlashAttention - 2/TransformerEngine在包级部署时的适用范围。实验结果表明，与基线Transformer模型相比，我们的方法取得了显著改进，在Penn Treebank数据集上的困惑度降低了10.3%，在WikiText - 2数据集上的困惑度降低了7.5%。通过系统的消融研究，我们确定了推动这些改进的关键因素，并发现可学习的门控因果滤波是效果最显著的机制。本文的其余部分结构如下：第2节回顾了关于Transformer改进以及频谱或结构化序列处理的相关研究。第3节介绍了Transformer架构和注意力机制的必要背景知识。第4节对我们提出的作为注意力后因果滤波模块的MBF进行了理论分析。第5节详细阐述了我们所提出的多频段残差融合方法。第6节描述了实验设置和实现细节。第7节展示了全面的实验结果及其解读。最后，第8节对本文进行总结，并探讨未来的研究方向。片段内容：Transformer架构的改进：自从最初的Transformer架构被提出以来[1]，人们提出了许多改进措施来提升其性能和效率。这些改进大致可分为三个领域：注意力机制优化、架构修改以及训练优化。注意力机制优化的重点是降低计算复杂度并提高注意力质量。Linformer[5]通过某种方式将注意力的计算复杂度从二次方降低到了线性。Transformer架构：Transformer架构[1]由编码器和解码器块组成，每个块都包含多头注意力机制以及带有残差连接和层归一化的前馈子层。对于语言建模任务，我们关注的是仅解码器的架构，该架构能够自回归地生成文本。每个Transformer块可以表示为：其中 表示第 层的隐藏状态，批量大小为，序列长度为，隐藏维度为。多头注意力机制：多头注意力机制用于计算注意力。语言序列的频域分析：自然语言序列具有丰富的时序动态特性，可以通过信号处理理论的视角对其进行分析。虽然语言建模是基于离散的令牌进行的，但注意力输出会形成连续的隐藏状态序列，此时因果滤波可以为这些序列提供有效的局部变换。工作假设（非正式）：我们将频率作为解读注意力输出的一种手段，其中变化平缓的成分对应于变化缓慢的模式，而变化剧烈的成分则对应于变化迅速的模式。概述：我们提出的带每头门控的多频段残差融合架构通过在应用自适应门控和残差融合之前，让注意力头输出经过多个可学习的因果滤波器分支的处理，从而增强了标准的Transformer注意力机制。图1展示了整个架构。可学习的因果多滤波处理：我们方法的核心创新在于让注意力头输出经过多个可学习的因果卷积滤波器处理。这些滤波器可以在频域中进行分析，但是……数据集：我们在三个语言建模数据集上评估了我们的方法：Penn Treebank（PTB）、WikiText - 2和WikiText - 103，这些数据集提供了具有不同词汇量和序列复杂度的互补评估场景。Penn Treebank（PTB）[31]是语言建模研究领域广泛使用的基准数据集。我们采用标准的划分方式，即42,068条训练句子、3,370条验证句子以及3,761条测试句子。该数据集的词汇量为10,000个单词，平均句子长度约为21个单词。主要结果：表4展示了将我们的多频段残差融合方法与基线Transformer在三个数据集（PTB、WikiText - 2、WikiText - 103）以及两种模型规模（GPT - 2小型和GPT - 2中型）下的主要实验结果对比。我们最佳的配置（第6次运行）在所有评估场景下都取得了显著优于基线的成绩。我们的方法在不同的数据集和模型规模上均能带来一致的改进效果：在PTB数据集上的困惑度降低了10.3%（从44.42降至9.85），在WikiText - 2数据集上的困惑度降低了7.5%。结论：本文提出了带每头门控的多频段残差融合机制，这是一种可在注意力聚合后用于增强Transformer语言模型的可学习因果多滤波模块。MBF会让每个注意力头输出经过多个因果卷积滤波器分支的处理，然后通过带有残差连接的每头门控机制对这些分支进行自适应融合。CRediT作者贡献说明：黄志高：写作——审阅与编辑、写作——初稿撰写、可视化、验证、软件、资源、方法论、形式分析。潘淼：写作——审阅与编辑、软件、方法论、资金获取。李全发：写作——审阅与编辑、可视化、软件。资金支持：本研究未获得公共部门、商业机构或非营利组织方面的任何特定资助。利益冲突声明：作者声明他们不存在任何可能影响本文所汇报工作的已知财务利益或个人关系。黄志高在中国燕山大学获得了通信工程硕士学位。他目前是中国泉州师范学院的物理与信息工程系讲师。他的研究兴趣包括人工智能与机器学习、用于计算机视觉和自然语言处理的深度学习、强化学习与智能系统、可解释的人工智能与模型可解释性、社交媒体数据分析以及网络技术。