**引言:**Transformer 架构自 2017 年被提出以来,基于 Attention 机制的优势迅速崛起,并成为 LLM 的主流架构。然而,计算成本和 KV 缓存问题等局限让业内认为原始的 Transformer 不会是 AGI 的最终版本,因而涌现出大量对 Attention 的改良、近似乃至替换工作,且热度越来越高。
替换 Transformer 不如爆改 Attention?
深度学习领域围绕着如何改进 Transformer 架构的讨论从未停歇。此前,许多研究者试图寻找能够完全替代 Transformer 的新架构,以解决其在计算效率和内存占用方面的局限性。然而,随着研究的深入,继续使用Transformer但改进注意力机制(Attention Mechanism)的做法又开始升温。
-
随着模型规模的扩大和应用场景的复杂化,Transformer 的局限性不断显现,这些缺陷促使研究者们不断探索新的改进方向
① Transformer 最明显局限在于 Attention 机制的二次方的复杂度带来的计算成本,其需要计算每个词与其他所有词的注意力权重,导致时间复杂度为 O(n²),序列长度增加时计算成本急剧上升。
② KV 缓存问题同样明显,为了加速推理,Transformer 需要缓存每个词的 Key 和 Value,这在处理长序列时会占用大量内存,空间复杂度为 O(n)。
③ Transformer 还存在注意力分散的问题,在长序列中,全局注意力机制可能使重要信息被稀释,导致模型难以聚焦关键信息。
④ 由于 Transformer 依赖额外的位置编码来理解词序,但在处理特别长的序列时,位置编码的可靠性会下降。
⑤ 此外,Transformer 逐词生成的机制导致了训练和推理不一致问题。
-
众多局限引发的共识是,原始的 Transformer 不会是 AGI 的最终版本,但这些局限引出了许多改进工作。
-
由于 Attention 机制被广泛认为是 Transformer 成功的关键,因此许多研究都集中在如何改进自注意力机制,以提高其效率和效果。(信源)(PDF下载)
① 第一种是基于传统 Softmax Attention 的各类改进和变种,涵盖 MHA、MQA、GQA 等。
② 第二种是以 RetNet、GLA 和 Lightning Attention 等工作为代表的 Linear Attention 路线。这些方法通过数学变换将计算复杂度从二次方降到线性(O(n)),同时基本保留了原有的模型表现。
③ 在试图降低 Attention 二次方复杂度至线性复杂度的工作中,还衍生了 Mamba 为代表的态空间模型和以 DeltaNet 为代表的在线学习等路线。
④ 还有类似如 MiniMax-01 的工作尝试结合 Softmax Attention 和 Linear Attention 的,用混合架构发挥两种机制的优势。
由 Softmax Attention 衍生的最新变体有什么特点?
在自然语言处理(NLP)领域,Softmax Attention 机制及其变体已成为提升模型性能和效率的关键技术。然而,随着模型规模的扩大和任务复杂度的提升,Softmax Attention 的时间复杂度为 O(n²)的局限性逐渐显现。这促使研究者们提出了多种变体,以优化 Softmax Attention 的性能和效率。
-
传统 Softmax Attention 主要涉及 Self-Attention、MHA、MQA 和 GQA 几种类型。在此之上,2024 年又新出现了 MLA、MFA 和 TPA 等变体,在减少 KV 缓存,强化模型长序列任务表现上取得突破。(知乎1、知乎2)
① Self-Attention 的基本思想是通过计算输入序列中每个元素与其他所有元素的相关性,来确定哪些元素对当前元素的输出最为重要。但由于需要比价每个元素与其他所有元素,导致计算复杂度高,难以适应长序列。
② MHA(Multi-Head Attention)是 Self-Attention 的进阶扩展,通过将输入分割成多个头(heads),每个头学习不同的特征,从而提升模型的表达能力,但 MHA 的时间复杂度仍然是 O(n²),不擅长长序列。
③ 为了减少 KV 缓存的内存占用,MQA(Multi-Query Attention)通过共享 Key 和 Value 投影矩阵,显著减少了 KV 缓存的大小,而 GQA(Group-Query Attention)则进一步通过分组共享机制,在减少内存占用的同时保留了一定的表达能力。
-
MLA(Multi-head Latent Attention)是深度求索团队在 DeepSeek-V2 工作中提出的对 MHA 的改良,可在后者的基础上减少 90%的 KV 缓存,适应更长的上下文。(论文)
① MLA 通过引入低秩分解来优化 KV 缓存的使用,减少内存需求。其将 Key 和 Value 压缩成一个低秩的隐向量,然后通过特定的映射矩阵将这个隐向量映射回 Key 和 Value,尽在低维潜空间中执行 Attention 计算。
② 此外,由于解耦的旋转位置编码(RoPE)与低秩 KV 压缩不兼容,MLA 采用了特殊的解耦策略。它使用额外的多头查询和共享 Key 来携带 RoPE,避免了在推理过程中重新计算所有前缀 token 的 Key,从而提高推理效率。
③ MLA 的局限在于,由于共享了低维空间,其丢失的维度信息会将 MLA 的表达能力局限在其中某个子空间。有分析指出,在 Deepseek-V2 中,可能较大的参数量导致表达能力过剩,没有出现这个问题。(知乎)
-
MFA(Multi-matrix Factorization Attention)则是阶跃星辰与清华等机构提出的注意力机制。MFA 通过低秩矩阵分解扩展注意力头的数量和维度,减少 KV 缓存的使用。
① MFA 的核心思想是通过扩展注意力头的数量和维度,同时利用低秩矩阵分解来减少 KV 缓存的使用量,从而在不显著增加计算开销的情况下提升模型的表达能力。
② 该工作还提出了扩展版本 MFA-KR,通过将 Key 缓存重新用作 Value(重参数化 Value 投影),进一步减少了 50%的 KV 缓存的使用量。
-
TPA 是清华、期智研究院和 UCLA 顾全全团队等提出的工作,其通过张量分解来表示 Q、K 和 V,只保留分解的版本而不存储完整的静态 KV,可让内存占用节省 90%,且不牺牲性能。(论文)
① TPA 在注意力计算过程中,它摒弃了传统的通过单一线性映射生成每个头的 Q、K、V 的方式,而是把它们分解为多个(与上下文相关的)张量积之和,更加适应资源首先环境中的长序列任务。
② TPA 还可以与 RoPE 无缝集成,直接代替现有模型中的 MHA,且无需额外修改。
结合 softmax Attention 和 Linear Attention 的混合架构有多强?
许多试图优化 Attention 机制的工作都在尝试将其二次方复杂度降低至线性复杂度,有次带来了 Linear Attention 的相关工作。较有代表性的 Mamba 架构基于 SSM 的递归公式成功将 Attention 的复杂度降为线性,但由于并行化支持较弱、通用性不足等问题导致未能在 LLM 领域得到广泛应用。在此之上,Mimimax 团队在近期的工作中将 Linear Attention 和 Softmax Attention 混合使用,实现让模型对 400 万 token 的超长上下文的处理能力。
-
Minimax-01 是该团队在 2025 年 1 月中旬新开源的模型系列。该模型采用混合架构,结合了 Lightning Attention、Softmax Attention 和 MoE,实现与 DeepSeek-V3、GPT-4o 等领先模型相竞争的能力。
-
在模型架构方面,MiniMax-Text-01 以 7:1 的比例结合了 Linear (Lightning) Attention 和 Softmax Attention。每 7 个 Lineanr Attention 块后插入一个 Softmax Attention 块,总共 80 层。
① 模型的隐藏(hidden)层大小为 6144,每层包含 32 个专家,每个专家的前馈网络隐藏维度为 9216。
-
MiniMax-Text-01 采用多阶段训练策略,并通过多种优化方法来提高训练和推理效率。
① 训练流程中,该工作首先构建了高质量预训练语料库,并配合数据实验框架评估不同数据配置对性能的影响,从而优化数据混合策略;然后用三阶段训练策略逐步扩展上下文长度,最终达到 100 万 tokens。
② 计算优化方面,研究者针对 MoE 设计了 EP-ETP 重叠策略来减少通信开销;适应了「数据打包」技术减少计算浪费,并设计了 Varlen Ring Attention 和 LASP+ 算法,以支持长序列训练。
③ 推理优化方面,该工作融合了多个内存密集型内核来优化推理速度,实现了超过 75% 的模型浮点运算利用率。
-
MiniMax-01 系列包含基础语言模型 MiniMax-Text-01 和视觉多模态模型 MiniMax-VL-01 两个模型,表现均可与业内领先的同类型模型相竞争,甚至更优。
① MiniMax-Text-01 在 MMLU、MMLU-Pro、C-SimpleQA、IFEval、GPQA、MATH 等基准测试中,与 GPT-4o、DeepSeek-V3 和 Claude-3.5-Sonnet 等顶尖模型相当,对超过 20 万 token 的上下文处理性能显著优于其他模型。
② MiniMax-VL-01 在 MMMU、MMMU-Pro、ChartQA、DocVQA、AI2D、MathVista、OCRBench 等基准测试中,表现同样与 GPT-4o、和 Claude-3.5-Sonnet 等顶尖模型相当,且具备与 MiniMax-Text-01 类似的长上下文处理优势。
融合多种技术的 Titans 有望继任 Transformer 吗?
近期,同样由谷歌团队开发的 Titans 架构通过长期神经记忆模块与 Attention 机制的结合,实现了短期记忆与长期记忆的协同工作。这种设计使得模型能够处理长达 200 万 tokens 的上下文窗口,远超传统 Transformer 的能力范围,被认为是 Transformer 的「继任者」。(机器之心)
-
Titans 的核心创新在于提出了一种长期神经记忆模块(Long-term Memory Module,LMM),作为元上下文模型,在测试时学习如何将数据记忆到其参数中。
-
LMM 的设计灵感来源于人类对意外事件的记忆机制。模型通过计算输入数据的意外分数(surprise metric),即输入相对于模型预测的梯度,来决定哪些信息需要被记忆或遗忘。其输入数据与过去数据的差异越大,梯度越大。
① 除了长期神经记忆模块,Titans 还引入了自适应遗忘机制,允许模型动态管理有限的记忆容量,从而在处理大规模序列时保持高效。
-
Titans 对 Attention 机制的改进主要体现在其与 LLM 的融合方式上。传统的 Transformer 架构依赖于固定大小的上下文窗口,而 Titans 通过将长期记忆模块与 Attention 结合,使得模型能够在处理当前上下文的同时,动态访问存储在长期记忆中的历史信息。
① 短期记忆由注意力机制负责,能够精确建模当前上下文窗口内的依赖关系;而 LMM 则负责存储历史信息,并在需要时提供支持。
-
Titans 架构有三种变体,分别是将记忆作为上下文(Memory as a Context,MAC)、作为门控(Memory as a Gate,MAG)和作为独立层(Memory as a Layer,MAL)。
① 其中,MAC 将长期记忆和持久记忆作为当前输入的上下文一起输入给 attention;MAG 在记忆模块和滑动窗口 attention 两个分支上进行门控融合;MAL 将记忆模块作为独立的一层,压缩历史信息后再输入给 attention。
-
经测试,在语言建模、常识推理、时间序列预测等任务上,Titan 的表现均优于 Transformer 和 Mamba 等其他架构的 SOTA 模型。
① 在语言建模和常识推理任务中,Titans 的三种变体(MAC、MAG 和 MAL)均取得了优于基线模型的性能,尤其是在处理长序列时,其优势更为明显。
② 在大海捞针任务中,Titans 能够有效处理超过 200 万 tokens 的上下文窗口,并实现比 TTT、Mamba2、DeltaNet 等更高的准确性。
③ 微调设置中,,小型微调版本 Titan(MAC)也有效扩展到超过 200 万的上下文窗口,在 BABILong 基准测试中超越了 GPT-4、Llama3+RAG 和 Llama3-70B 等大模型。
④ 此外,Titans 在时序预测和 DNA 建模等任务中也表现优越。