跳转至

新架构探索

RetNet 如何通过多尺度保留机制实现训练并行和推理高效？与 Transformer 对比。

RetNet 概述

RetNet（Retentive Network）由微软提出，旨在兼具 Transformer 的训练并行性、RNN 的推理高效性以及长程建模能力。它采用 多尺度保留（Multi-Scale Retention） 机制替代自注意力。

多尺度保留机制 标准注意力为 O(N2)，但在 RetNet 中，保留机制定义为：

image.png

多尺度：使用多个不同的 γ 值（对应不同的保留头），每个头捕获不同时间尺度上的依赖（大 γ 捕获长程，小 γ 关注局部）。各头输出拼接，类似多头注意力。

与 Transformer 对比

  • 训练：RetNet 的并行模式与 Transformer 训练代价相当（都是 O(N²d)），但因无需 softmax，计算略快。

  • 推理：Transformer 需 KV 缓存，内存随长度线性增长；RetNet 递归模式只需固定大小的状态 St，推理内存恒定，延迟极低。在长序列上 RetNet 速度优势显著。

  • 性能：在语言建模上，RetNet 与同等规模 Transformer 竞争甚至更优，尤其在长序列任务上。但 Transformer 在需要精确 token 级检索的任务中仍有优势。

因此，RetNet 是朝着 “训练并行、推理高效” 目标迈出的重要一步，但依然保留了类似注意力的矩阵分解，与线性注意力有相似之处。

RWKV 是如何结合 RNN 和 Transformer 优势的？它的注意力机制是怎样的？

RWKV (Receptance Weighted Key Value) 是一个开源的无注意力架构，旨在像 RNN 一样高效推理，同时像 Transformer 一样可并行训练。

核心思想

RWKV 使用一种时间衰减和 token 交互的混合机制，类似用“指数衰减的记忆”取代注意力。

注意力机制（WKV 算子）

给定输入序列，RWKV 计算四个向量：接收度 R、权重 W、键 K、值 V（均由线性变换得到）。

其核心操作为：

image.png

这里 ww 是学习到的通道级时间衰减因子（正数，使得越远的历史权重越小），uu 是当前 token 自身的偏置（提升当前 token 影响力）。这本质上是带有可学习时间衰减的加权平均，形式上与带有遗忘门的 RNN 相似。

image.png

如何结合 RNN 和 Transformer 优势：

  • 像 RNN：推理时可状态化。将分子分母维护为移动状态，每步更新时只需 O(1) 计算和内存，无 KV 缓存，极适合流式生成和边缘设备。

  • 像 Transformer：训练时，通过使用卷积或并行前缀和技巧（parallel scan），可同时处理整个序列，像 Transformer 一样并行计算 WKV 算子，无需像传统 RNN 一样串行。这使得 RWKV 训练高效。

  • 无限上下文：由于状态固定大小，理论上可处理无限长度序列（实际上受状态容量限制）。

性能：RWKV 在语言建模和常识推理上与同规模 Transformer 竞争，在长序列和低延迟场景中优势明显。但它在需要精确回忆遥远特定信息的任务（如针大海捞针）中可能逊色于 Transformer 的全注意力。

状态空间模型（SSM）的基本原理是什么？S4 模型是如何处理长序列的？

SSM 基本原理

image.png

S4（Structured State Space）模型

S4 通过将 A 矩阵初始化为特殊的 HiPPO 矩阵（可最优地压缩历史信息），并将 SSM 计算转化为卷积形式，实现了长序列的高效训练。具体而言：

image.png

S4 处理长序列的优势：凭借全局卷积形式，训练速度快且内存效率优于注意力。同时，SSM 的状态 hh 固定大小，推理内存恒定。S4 在长序列分类（如 Pathfinder、Long Range Arena）上取得了突破性成绩，证明了 SSM 在极长范围依赖建模上的能力。

Mamba 提出的选择性扫描相比普通 SSM 有什么突破？它解决了什么问题？

Mamba 的突破：选择性扫描（Selective Scan）

image.png

解决的问题：

  • 内容感知的记忆：模型可以学会根据上下文有选择地遗忘或保留信息，类似于注意力机制中的“关注”，但保留了 SSM 的效率。

  • 克服标准 SSM 在离散信息回忆和复制任务上的短板：选择性扫描显著提升了在需要精确 token 级记忆的任务上的表现，缩小了与 Transformer 的差距。

  • 硬件感知的并行实现：时间变化 SSM 不能直接用卷积并行训练，Mamba 设计了并行扫描（parallel scan）算法结合 kernel fusion，在 GPU 上高效执行选择性递归，实现了训练和推理的高速。

Mamba 的核心贡献在于将线性 SSM 与输入依赖选择结合起来，打造出首个在语言建模上达到甚至超越同规模 Transformer 的线性复杂度模型，引发了替代 Transformer 的热潮。

Mamba 与 Transformer 在语言建模上的性能对比如何？各自的适用场景？

性能对比

  • 语言建模困惑度：在相同参数规模（如 1.3B, 2.8B）和训练数据下，Mamba 的困惑度达到或略优于强基线 Transformer（如 LLaMA 架构的 Transformer++）。尤其是随着序列长度增加，Mamba 的优势更显著，因其 O(N) 复杂度能更好地处理长上下文。

  • 推理吞吐：Mamba 在长序列生成上吞吐远超 Transformer。由于其无 KV 缓存、每步 O(1)，推理速度在序列长度增加时几乎恒定，而 Transformer 呈线性增长。在 2048 以上的长文本生成中，Mamba 可比 Transformer 快数倍。

  • 特定能力：Transformer 凭借精确的全对全交互，在需要检索极远距离的单个 token（大海捞针）或精确复制序列的任务上可能稍优，但 Mamba 的选择性机制在多数常见长程任务上已接近或持平。

适用场景

  • Transformer：仍适合需要极强上下文学习、复杂多步推理、以及利用庞大的预训练生态（如各种工具链、适配器）的场景。目前多模态模型、大规模助手几乎全是 Transformer。

  • Mamba：特别适合长序列处理（DNA、长文档、音频）、低延迟流式应用（语音助手、实时翻译）、资源受限的边缘部署，以及需要高吞吐批量服务的场景。

Mamba-2 在结构化状态空间对偶性上有什么新贡献？与注意力机制的联系是什么？

Mamba-2 的贡献：结构化状态空间对偶性（Structured SSM Duality）

Mamba-2 揭示了线性注意力、状态空间模型和结构化矩阵之间的理论统一。核心新贡献是提出了 结构化掩码注意力（Structured Masked Attention, SMA） 作为 SSM 的对偶形式。

  • SSM 与注意力等价性：Mamba-2 证明，特定参数化的时间变化 SSM（即选择性 SSM）等价于一种带有半可分离（semi-separable）结构掩码的注意力机制：

image.png

  • 其中 L 是一个半可分离矩阵，即其下三角部分可以由低秩分解快速生成（不像全注意力是一个稠密矩阵）。这意味着状态空间模型实际上是在高效地计算一种结构化的注意力。

  • 统一框架：Mamba-2 引入了 SSD（State Space Duality） 的概念，展示了 SSM 递归形式和注意力矩阵形式之间的双元性。这使得 Mamba-2 可以像注意力一样进行块并行计算，进一步提升训练效率。

  • 改进的架构：利用这种对偶性，Mamba-2 设计了新的算法（如块分解）以及更丰富的头结构（类似于多头注意力），在保持理论效率的同时，提高了模型表达能力。

与注意力机制的联系：它填补了线性注意力与 SSM 之间的理论鸿沟。Transformer 的注意力使用全矩阵 L（因果 mask 的变体），而 SSM 使用半可分离 L。Mamba-2 表明我们可以在这两者之间连续插值，设计出更灵活的结构化注意力，可能兼具 Transformer 的强表达力和 SSM 的高效率。

Mamba 能否替代 Transformer 成为下一代基础架构？目前有哪些局限性？

当前局限性

  • 生态和工具链：Transformer 拥有极其完善的训练和推理生态（FlashAttention、PagedAttention、TensorRT-LLM、各种并行策略）。Mamba 等新架构的定制内核、分布式训练支持尚不成熟。

  • 扩展律验证不足：Transformer 的 scaling law 已被大量实验确认，Mamba 在数百亿、千亿参数规模下的稳定性和性能天花板仍待充分探索。

  • 上下文学习能力：Transformer 的全注意力提供了极强的上下文学习（in-context learning）能力，这是当前大模型涌现能力的关键。Mamba 的固定状态压缩是否能支撑同样强大的少样本推理尚存疑问。

  • 检索与精确复制：在需要从超长上下文中精确检索单一事实的任务中，Transformer 的全注意力机制具有理论优势。Mamba 选择性记忆可能面临记忆覆盖。

  • 训练效率：虽然推理极快，但选择性扫描的并行实现复杂度高，目前训练速度在中等长度上仍可能慢于高度优化的 Transformer（如使用 FlashAttention-3）。不过随着优化，这个差距在缩小。

前景：Mamba 很可能不会完全取代 Transformer，但会与其混合或作为特定场景的互补。如同 Jamba 等混合架构，未来基础架构可能是 Transformer 层与 SSM 层的交替组合，或根据序列位置动态选择。

线性 Transformer 与状态空间模型之间有什么理论联系？

两者本质上都试图通过低秩或核方法近似注意力，且都可以被归为 线性注意力 框架。理论联系如下：

image.png

  • Mamba-2 的对偶性：明确指出了 SSM 等价于半可分离的结构化注意力，而线性 Transformer 是非结构化低秩注意力。SSM 的 A 矩阵赋予衰减结构，可以看作是在线性注意力上增加了一个归纳偏置，保证长程记忆的稳定。

  • 性能差异根源：线性 Transformer 没有序列长度的先验衰减，容易造成远端信息覆盖或梯度问题；SSM 通过精心设计的 A 矩阵（如 HiPPO）实现了稳定的记忆压缩。Mamba 的选择性相当于让 B, C 输入依赖，即动态调整特征映射 ϕϕ，补足线性 Transformer 的表达力。

Jamba 等混合架构混合了 Mamba 和 Transformer 层，这样设计的好处是什么？

Jamba 由 AI21 Labs 提出，将 Mamba 层与 Transformer 层交织组成单个模型。

好处：

  • 取长补短：Mamba 层高效处理长序列全局依赖，提供线性复杂度；Transformer 层提供强大的内容感知注意力，能捕捉复杂的 token 交互。混合使得模型既保留了 Transformer 的强大推理和上下文学习能力，又兼具 SSM 的长序列效率。

  • 内存与吞吐平衡：Transformer 层需要 KV 缓存，但可以通过 GQA 等压缩；Mamba 层无缓存。混合减少了整体 KV 缓存大小，提升了推理吞吐，同时在全注意力层中关键的密集交互不会损失。

  • 灵活的 MoE 结合：Jamba 在其混合层之上还引入了 MoE，使模型总参数大但激活参数少。混合架构让 MoE 可以在部分层上部署，减少通信压力。

  • 性能经验：实验证明，同等训练预算下，Jamba 的长序列表现优于纯 Transformer 或纯 Mamba。它展示了异构层的组合可以比单一机制更优，这可能是未来大型基础模型的发展方向。

你认为 Transformer 架构是否已经到了瓶颈？未来最可能突破的方向是什么？

瓶颈分析

Transformer 的二次复杂度并未被 FlashAttention 等彻底消除，仅被缓解。随着上下文窗口向百万 token 迈进，计算和内存开销依然巨大。但 Transformer 的表达力、扩展性和生态使其短中期内仍不可替代。因此，瓶颈在于长序列的极限效率和推理成本。

未来突破方向

  1. 混合架构：Mamba+Attention+MoE 的组合，在全注意力层中应用极稀疏或极低秩注意力，大部分层用高效 SSM 处理。

  2. 动态计算：输入依赖的层选择或提前退出，让模型自适应地分配计算，浅层处理简单 token，深层处理难 token。

  3. 新位置编码与外推方法：更优的位置编码让短窗口训练的模型直接应用于极长序列，无需大量长文本训练。

  4. 存储器与检索增强：将外部存储（向量数据库、神经记忆）紧密耦合到注意力中，使模型不必将所有知识压缩到参数中。

  5. 硬件-算法协同设计：新一代 AI 加速器可能专门针对稀疏、低秩或 SSM 算子进行优化，从而改变架构设计偏好。

  6. 超越序列的建模：将序列视为图、树等更丰富结构，用结构化注意力或神经算法推理，可能突破纯 token 序列的局限性。

如果让你设计一个全新的序列模型架构，用于取代 Transformer，你最可能借鉴哪些思想？

我会设计一个混合、自适应且记忆增强的架构，融合以下思想：

  • 状态空间模型的线性递归（Mamba 的选择性扫描）：作为大部分层的基础计算单元，确保长序列效率。

  • 稀疏全局注意力层（每 N 层插入一层全注意力或路由注意力）：用于处理需要全局交互和上下文学习的复杂模式，实现“慢思考”与“快思考”结合。

  • 可扩展的外部记忆模块：如 Token 级别的记忆库，使用寻址机制（类似于 Neural Turing Machine），支持读写和遗忘，通过记忆 token 注入各层，实现无限上下文。

  • 自适应计算：每个 token 经过轻量级路由器决定当前层应使用局部 SSM、全局注意力还是直接从记忆读取，用动态深度实现计算预算分配。

  • 多分辨率时间尺度：在不同层或头中使用不同的时间衰减因子和窗口大小，显式捕获从字词到篇章的多尺度依赖。

  • 硬件感知的设计：所有算子都支持块操作和融合内核，适配下一代 GPU 的 Tensor Core 和可能的 SSM 专用指令。

  • 训练方式：采用类似 BERT 的理解任务与自回归生成任务联合预训练，使模型既有强编码能力又能高效解码。

这样设计可以在推理效率、长程建模、上下文学习之间取得更优平衡。

分析一下目前各种新架构在工业界落地的情况和阻力。

现状

  • Transformer 仍绝对主导：GPT-4、Claude、Gemini、LLaMA、Qwen 等所有一线大模型均基于 Transformer。部署工具（TensorRT-LLM、vLLM）和硬件优化（FlashAttention）全围绕 Transformer。

  • Mamba 及其混合架构：有较强的学术和开源社区兴趣，但工业级落地仍处早期。Cartesia（Mamba 作者创办）正在推商业化应用，其他公司（如 AI21 Labs 的 Jamba）也展示了潜力。主要用于长文本、DNA 等垂直场景。

  • RWKV：因其开源模型和低延迟特性，在游戏 NPC、实时对话、边缘设备中有小范围采用，也有中文社区推动。但其性能上限和生态与 Transformer 有差距。

  • 线性注意力、RetNet：除微软内部探索外，未见大规模工业应用。理论限制导致其在复杂任务上不如 Transformer。

落地阻力

  1. 路径依赖与沉没成本：各大公司已在 Transformer 生态投入巨资，迁移成本极高，除非新架构表现出显著超越的优势（≥2-3 倍性能或成本降低）。

  2. Scaling Law 不明：Transformer 的 scaling law 指导投入，新架构能否平滑扩展到万亿参数并无保障，企业不愿冒险。

  3. 软件栈和硬件适配：缺乏统一高效的推理引擎，每推一个新架构都需要从头定制内核，难以跨平台部署。

  4. 能力缺失风险：上下文学习、思维链等关键能力可能与 Transformer 的 dense attention 强相关，新架构若在这些高级能力上落后，则无法成为基座。

  5. 人才和知识壁垒：绝大多数 AI 工程师熟悉 Transformer，对 SSM、RWKV 等新范式的调试、优化经验不足，增加项目风险。

展望：新架构很可能通过混合模式逐步渗透：在 Transformer 主干中插入新模块（如 Mamba 层），渐进迭代。待到某混合架构在多个 benchmark 上整体明显优于纯 Transformer，且训练和推理成本大幅降低时，工业界才会加速迁移。