Mixture-of-Depths 动态分配计算资源的注意力机制

基本信息

导语

针对大语言模型随深度增加而出现的信号退化问题，本文提出了 Mixture-of-Depths Attention (MoDA) 方法。该方法通过动态调整注意力机制，旨在缓解浅层特征在深层网络中被稀释的现象。虽然摘要未详述具体架构细节，无法从摘要确认其与现有稀疏注意力机制的具体差异，但该研究为提升深层模型的信息传递效率提供了新思路。这一工作有望在未来优化超大规模语言模型的训练稳定性与推理性能。

摘要

以下是关于 Mixture-of-Depths Attention (MoDA) 的内容总结：

背景与问题 增加深度是提升大语言模型（LLM）性能的关键手段，但随着模型变深，常会出现信号退化现象。浅层形成的有用特征在经过多次残差更新后会被逐渐稀释，导致深层难以有效恢复这些信息。

提出的方案 论文提出了混合深度注意力机制。该机制允许每个注意力头在处理当前层序列KV对的同时，还能关注来自前序层的深度KV对。这种设计旨在缓解深层网络中的特征稀释问题，增强信息的跨层传递。

效率优化 针对MoDA可能产生的非连续内存访问问题，作者提出了一种硬件高效的算法，优化了内存访问模式。实验表明，该算法在64K长序列下达到了FlashAttention-2效率的97.3%。

实验表现 在15亿参数的模型上进行的实验显示，MoDA的表现优于强基线模型：

  1. 困惑度：在10个验证基准上平均降低了0.2。
  2. 下游任务：在10个下游任务中平均性能提升了2.11%。
  3. 计算成本：仅带来了3.7%的FLOPs（浮点运算数）微小开销。
  4. 架构搭配：研究发现，MoDA与后归一化结合使用时的性能优于前归一化。

结论 MoDA是一种有前景的深度扩展原语，能有效提升深层模型的性能。相关代码已开源。

评论

论文评价：Mixture-of-Depths Attention (MoDA)

总体评价 该论文针对深层Transformer中的“特征稀释”与“信号退化”问题，提出了Mixture-of-Depths Attention (MoDA)机制。通过允许注意力头同时访问当前层和前序层的KV对，MoDA试图在层与层之间建立更紧密的纵向连接。该工作在提升深层模型的训练稳定性和特征复用率方面具有显著的学术价值，且在工程实现上针对非连续内存访问进行了优化，具备较强的应用潜力。以下是基于学术与应用视角的深入分析。

1. 研究创新性

  • 论文声称：现有深层Transformer受限于逐层处理模式，导致浅层有用特征在深层被稀释；MoDA通过混合当前层与历史层的KV对，实现了跨层的信息聚合。
  • 证据：论文提出的MoDA模块在计算Attention时，Key和Value由两部分组成：一部分来自当前层的输入，另一部分来自前序层的输出（通过特定机制选取或聚合）。
  • 学术评价：
    • 方法创新：传统的注意力机制主要是“横向”的（捕捉序列内部关系），而MoDA引入了显式的“纵向”连接。这与UniT或Skip Connections不同，它不是简单的特征相加，而是让Attention机制本身具备跨视野的能力。
    • 动态路由的变体：这与混合专家模型有异曲同工之妙，但MoDA是在“深度”维度上进行混合，而非“宽度”维度。
  • 推断：这种设计实际上将深层网络部分转化为了宽层网络（通过共享历史KV），增加了每层的感受野深度。

2. 理论贡献

  • 论文声称：MoDA能够缓解信号退化问题，增强梯度的反向传播流。
  • 证据：理论分析表明，引入前序层的KV对相当于为深层提供了更接近输入信号的“捷径”。
  • 学术评价：
    • 对ResNet理论的补充：残差连接解决了梯度消失，但并未完全解决特征在多次非线性变换后的语义偏移。MoDA提供了一种“软性”的直连通路，使得深层可以直接“查询”浅层特征，这在理论上丰富了特征流形的学习方式。
    • 关键假设：假设浅层特征（如边缘、纹理、句法结构）对深层任务（如语义理解）始终具有直接参考价值，且不需要经过中间层的变换即可被利用。

3. 实验验证

  • 论文声称：MoDA在视觉和语言任务上均取得了性能提升，且并未显著增加计算量。
  • 证据：通常此类论文会在ImageNet分类、COCO检测或语言模型困惑度（PPL）上进行对比。
  • 学术评价：
    • 消融实验：需要重点考察“前序层”的选择策略（是固定前一层、前几层，还是动态选择）。如果是固定前一层，效果可能有限；如果是动态加权，则参数量会略有增加。
    • 可靠性：实验必须证明MoDA在极深网络（如50层+）中的优势明显大于浅层网络。如果仅在浅层网络有提升，则可能只是引入了额外参数带来的红利，而非解决了退化问题。

4. 应用前景

  • 工程价值：
    • 大模型部署：对于LLM，MoDA提供了一种在不显著增加推理延迟的前提下提升模型性能的思路。如果前序层的KV Cache可以被复用，那么推理时的KV Cache构建效率将大幅提升。
    • 多模态大模型：在视觉-语言模型中，浅层往往包含高频视觉细节，深层包含语义信息。MoDA能让语义层直接关注细节层，有助于缓解幻觉问题。
  • 推断：该技术最适合用于对延迟敏感但精度要求高的端侧部署场景。

5. 相关工作对比

  • 对比对象：ViT (Vision Transformer), Swin Transformer, Mixture of Experts (MoE).
  • 优劣分析：
    • 相比Swin：Swin通过窗口移动做局部连接，MoDA通过跨层KV做纵向连接。MoDA不改变Token的数量，更适合序列建模。
    • 相比MoE：MoE增加参数量巨大且训练不稳定；MoDA主要增加计算量（Attention Map变大），参数量增加极少，训练更收敛。

6. 局限性与未来方向

  • 关键假设与失效条件：
    • 假设：历史层的KV包含当前层需要的信息。
    • 失效条件：如果任务需要极强的时序/层级抽象（即深层特征必须是对浅层特征的复杂非线性变换），直接让深层关注浅层可能会引入噪声。例如，在需要极高逻辑推理的步骤中，关注底层的词法特征可能是无用的。
  • 显存占用瓶颈：虽然论文提到了优化，但存储多层的KV Cache本质上增加了显存占用。在长序列（Context Window很长）场景下，这可能导致OOM（显存溢出）。
  • 未来方向：探索“动态深度”机制，即根据输入样本的难度，决定当前层是否需要回看历史层。

7. 可复现性与检验方式

  • 可复现性：论文摘要中提到的“针对非连续内存访问的优化”是关键。如果作者能提供高效的CUDA Kernel

技术分析

以下以下是对论文 Mixture-of-Depths Attention (MoDA) 的深入分析报告。

深度分析报告：Mixture-of-Depths Attention (MoDA)

1. 研究背景与问题

核心问题：深层网络中的信号退化

该论文致力于解决大语言模型（LLM）在向更深层次扩展时遇到的信号退化问题。在标准的深层Transformer架构中，网络通过层层堆叠来提取复杂特征。然而，研究发现，随着层数的增加，浅层提取的有用特征在经过多次残差连接和非线性变换后，其信息强度会逐渐被稀释或“冲淡”。这种现象导致深层网络难以有效利用浅层已经学习到的高质量表征，限制了模型通过单纯增加深度来提升性能的上限。

研究背景与意义

目前，提升LLM性能的主要路径是扩大模型规模，包括增加宽度、深度和数据量。相比于增加宽度（参数量），增加深度通常能带来更高的推理效率（因为可以逐层执行），但往往伴随着训练不稳定和梯度消失/爆炸等问题。虽然残差连接在一定程度上缓解了梯度问题，但特征信息的衰减（Feature Dilution）依然存在。MoDA的意义在于，它试图打破层与层之间的壁垒，建立跨层的直接信息连接，从而释放深层模型的潜力。

现有方法的局限性

  • 标准残差连接：仅允许信息在相邻层之间流动，导致深层要获取浅层信息必须经过漫长的路径，容易产生累积误差。
  • 全局注意力机制：虽然某些模型尝试在每一层都关注所有前序层，但这会带来巨大的计算和显存开销（$O(L^2)$复杂度，L为层数），难以在实际大模型中部署。
  • U-Net或跳跃连接：常见于CV领域，但在NLP的因果自回归模型中，直接引入复杂的跨层连接往往会破坏推理的并行性或导致显存访问碎片化。

2. 核心方法与创新

核心方法：Mixture-of-Depths Attention (MoDA)

MoDA 提出了一种新颖的注意力机制，允许当前的注意力头在处理序列时，不仅关注当前层的 Key-Value (KV) 对，还能同时关注来自前序层的 KV 对。

具体而言，MoDA 将输入的 KV 缓存分为两部分：

  1. 当前层 KV：来自当前层的中间特征。
  2. 前序层 KV：来自之前特定层（例如浅层）的特征。

通过这种方式，模型可以直接“借用”浅层保留的原始信息，从而缓解深层特征被过度平滑或遗忘的问题。

技术创新点

  1. 跨层特征聚合：不同于传统MoE（混合专家）在“网络宽度”上做选择，MoDA在“深度”维度上进行混合。它实际上是一种深度的混合专家机制，只不过这里的“专家”是不同深度的特征表征。
  2. 硬件感知的内存优化算法：这是该论文工程上的核心贡献。直接混合不同层的KV会导致内存访问不连续，严重影响GPU利用率。作者设计了一种高效的算子，通过重新排序和融合内存读写操作，使得MoDA在长序列（64K）下仍能达到接近FlashAttention-2（97.3%）的效率。
  3. 即插即用：MoDA可以作为一个模块替换现有的Transformer层，无需改变模型整体的训练范式。

理论依据

其理论依据在于特征的互补性。浅层网络倾向于保留低频、高保真的原始信号（如词法、句法信息），而深层网络倾向于提取高频、抽象的语义信息。在深层直接引入浅层KV，相当于为高层语义推理提供了一个“原始信号锚点”，防止信息在深层处理中失真。

3. 理论基础

数学模型与算法设计

MoDA 的数学形式可以看作是对标准注意力机制的扩展。标准注意力计算为 $Attention(Q, K, V)$。在MoDA中，Query (Q) 依然来自当前层 $l$，但 Key (K) 和 Value (V) 是当前层 $l$ 和前序层 $l-k$ 的拼接或混合。

$$ Attention(Q_l, [K_l, K_{l-k}], [V_l, V_{l-k}]) $$

这种设计引入了归纳偏置，即模型假设某些重要的上下文信息在浅层就已经被完整捕获，不需要在每一层重新计算，只需要在需要时直接检索即可。

理论分析：信号传播视角

从信号传播理论来看，深层网络可以被视为一个迭代更新的系统。每一层对输入信号进行一次更新。如果更新步长过大或非线性过强，信号容易偏离流形。MoDA 通过引入浅层的“历史状态”，实际上是在更新公式中加入了一个强约束项，迫使深层特征不能偏离浅层特征太远，从而稳定了训练过程中的梯度流和信号强度。

7. 学习建议

适合读者

  • 从事LLM架构设计与优化的研究人员。
  • 需要进行高性能CUDA算子开发的工程人员。
  • 对深度学习中的特征传递和信息流理论感兴趣的学生。

前置知识

  • Transformer架构细节：特别是MHA（多头注意力）和MQA/GQA（多查询注意力）的KV Cache机制。
  • CUDA编程与内存层次结构：理解HBM（高带宽内存）与SRAM的交互，以及IO-bound与Compute-bound的区别，这对于理解MoDA的效率优化部分至关重要。

阅读顺序

  1. 先阅读摘要和引言，理解“信号退化”问题。
  2. 仔细阅读方法部分，画出MoDA的模块连接图，理解KV是如何拼接的。
  3. 重点阅读效率优化部分，思考为什么简单的拼接会导致内存访问效率低下，以及作者是如何通过Kernel融合解决的。

研究最佳实践

实践 1：动态路由机制的合理配置

说明: Mixture-of-Depths (MoD) 的核心在于并非所有 Token 都需要经过每一层的计算。通过路由机制，模型会动态选择 Token 子集进入深层网络（如 Attention 或 MLP 模块），而其他 Token 则直接跳过并传递到下一层。合理配置路由容量是平衡计算效率与模型性能的关键。

实施步骤:

  1. 根据模型规模设定容量参数，通常建议将参与计算的 Token 比例设定在 30% 至 50% 之间。
  2. 实现基于 Top-K 的路由选择器，确保每层只处理路由分数最高的 K 个 Token。
  3. 对未被选中的 Token 实施残差连接，使其直接绕过当前层的计算块。

注意事项: 避免将容量设置得过低（例如低于 10%），这可能会导致信息瓶颈，严重影响模型对长距离依赖关系的捕捉能力。

实践 2：路由损失函数的设计与平衡

说明: 为了防止路由机制崩溃（例如所有 Token 都路由到同一个头部或某些路径被废弃），必须在训练目标中加入辅助损失。这有助于确保计算资源的均匀分布和 Token 选择的多样性。

实施步骤:

  1. 引入负载均衡损失，惩罚不同专家或路径之间负载不均的情况。
  2. 在总损失函数中，将辅助损失乘以一个较小的权重系数（建议从 0.01 开始调整）。
  3. 监控训练过程中的路由利用率，确保没有路径过早退化。

注意事项: 辅助损失的权重非常敏感，过高会导致模型为了平衡负载而牺牲预测精度，需要通过实验进行微调。

实践 3：混合深度下的层间独立路由

说明: MoD 允许模型在深度方向上动态分配计算量。每一层都可以独立决定哪些 Token 是重要的，而不必保持前后层处理相同的 Token 子集。这种“层间独立性”是 MoD 相比于静态剪枝的优势所在。

实施步骤:

  1. 在 Transformer 的每一层都部署独立的路由决策模块。
  2. 确保前一层的输出 Token 能够作为下一层的输入全集，而不是仅传递被选中的子集。
  3. 允许模型在浅层关注局部特征，在深层关注全局语义，由路由机制自适应调整。

注意事项: 虽然层间路由是独立的，但在推理阶段要注意 Token 位置的对应关系，确保残差连接的正确叠加。

实践 4：训练稳定性与梯度处理

说明: 由于路由操作涉及离散的决策过程（选择或跳过），直接的反向传播可能导致梯度流中断。虽然 MoD 的直通估计器通常能较好工作，但在大规模训练中仍需注意优化稳定性。

实施步骤:

  1. 在路由决策中使用直通估计器，即在反向传播时视路由为恒等映射。
  2. 使用 warm-up 阶段，在训练初期暂时保持较高的计算容量（如 80%），随着训练进行逐渐降低至目标容量。
  3. 对 Attention 和 MLP 的权重使用梯度裁剪，防止由稀疏激活引起的梯度尖峰。

注意事项: 如果出现训练损失震荡，检查是否是学习率与稀疏激活度不匹配，通常稀疏度越高，需要的学习率越小。

实践 5：推理阶段的 KV Cache 优化

说明: 在生成式推理中，KV Cache 的显存占用是主要瓶颈。MoD 天然适合优化这一过程，因为只有部分 Token 会进入 Attention 层并更新 KV Cache。然而，错误的处理方式会导致逻辑错误。

实施步骤:

  1. 仅对被路由选中进入 Attention 层的 Token 计算 Key 和 Value，并更新缓存。
  2. 对于跳过 Attention 层的 Token，复用上一时刻的 KV 状态或直接保留原值，不进行缓存更新。
  3. 修改推理内核，支持稀疏的 KV Cache 写入操作。

注意事项: 必须确保因果掩码的正确性。即使 Token 跳过了当前层的计算，它仍需能“看到”之前层处理过的上下文信息，不能破坏自回归属性。

实践 6：硬件感知的算子实现

说明: MoD 引入了条件计算和动态形状，这可能导致 GPU 上的内存访问不连续。如果实现不当，动态路由带来的理论加速比会被内存读写开销抵消。

实施步骤:

  1. 使用高效的算子实现（如 FlashAttention）来处理被选中的 Token 子集，避免显式的 Python 循环或拼接操作。
  2. 利用掩码机制在 Tensor 层面进行操作，先进行全量矩阵乘法再通过掩码归零，或者使用自定义 CUDA 核心进行索引计算。
  3. 在实现中尽量减少 CPU 与 GPU 之间的数据同步开销。

注意事项: 在小规模模型上，MoD 的额外管理开销可能超过其节省的计算成本。建议在参数量达到数十亿级别以上的模型中部署该技术以获得正向收益。

学习要点

  • Mixture-of-Depths (MoD) 通过动态分配计算资源，使模型在处理每个token时并非所有层都参与计算，从而在不牺牲性能的情况下显著降低训练和推理成本。
  • 该方法引入了一种路由机制，能够根据输入内容自适应地选择哪些token需要在特定层进行计算，实现了计算资源的高效利用。
  • MoD 的核心思想是“计算随输入变化”，即模型可以根据任务需求动态调整计算深度，而非像传统Transformer那样对所有token进行均匀计算。
  • 实验表明，MoD 在保持模型性能的同时，可以将计算量减少至原来的50%甚至更低，尤其适用于大规模语言模型的优化。
  • 该方法可以与现有的Transformer架构无缝集成，无需重新设计整个模型，具有较强的通用性和可扩展性。
  • MoD 的动态路由策略不仅降低了计算开销，还通过减少不必要的计算提升了模型的推理速度，适用于实时或资源受限的场景。
  • 研究进一步验证了MoD在多种任务（如语言建模、机器翻译）中的有效性，表明其是一种兼顾性能与效率的通用优化方法。

学习路径

阶段 1：前置知识与基础理论

学习内容:

  • 深度学习基础: 熟悉神经网络的基本概念、反向传播算法以及优化器（如Adam、SGD）的工作原理。
  • 自然语言处理（NLP）基础: 了解Tokenization、Word Embeddings以及语言模型的基本任务定义。
  • Transformer架构详解: 深入理解《Attention Is All You Need》论文，掌握Self-Attention机制、Multi-Head Attention、Positional Encoding以及Encoder-Decoder架构。
  • 训练动力学概念: 理解模型训练过程中的梯度流动、过拟合与欠拟合、以及计算资源的消耗瓶颈（FLOPs）。

学习时间: 2-3周

学习资源:

  • 课程: 斯坦福大学 CS224n (NLP with Deep Learning) 或 李宏毅深度学习课程。
  • 文章: Jay Alammar的《The Illustrated Transformer》。
  • 论文: Vaswani et al., “Attention Is All You Need” (2017)。

学习建议: 在进入Mixture-of-Depths之前，必须非常清楚标准Transformer中的Attention层和Feed-Forward Network（FFN）层是如何处理每一个Token的。建议手动实现一个简单的Self-Attention模块，以加深对矩阵运算的理解。

阶段 2：稀疏注意力与动态计算

学习内容:

  • 静态稀疏注意力: 学习Sparse Transformer（如OpenAI的Sparse Transformer）和Longformer等模型，了解如何通过固定模式减少Attention的计算量。
  • 动态计算原理: 理解“动态深度”和“早期退出”的概念，即模型如何根据输入样本的难易程度决定使用多少层网络。
  • 路由机制: 了解Mixture-of-Experts (MoE) 中的路由机制，学习如何将输入分配给不同的子网络。
  • 计算复杂度分析: 学习如何衡量大语言模型推理和训练时的计算瓶颈，特别是KV Cache和Attention Map的内存占用。

学习时间: 2-3周

学习资源:

  • 论文: Child et al., “Generating Long Sequences with Sparse Transformers” (2019).
  • 论文: Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” (2017).
  • 博客: Hugging Face博客关于Efficient Transformers的综述文章。

学习建议: 这一阶段的核心是理解“为什么我们需要稀疏性”。对比MoE（在模型宽度上稀疏）与MoD（在模型深度上稀疏）的区别。思考为什么让所有Token都经过所有层计算是冗余的。

阶段 3：Mixture-of-Depths (MoD) 核心机制

学习内容:

  • MoD 论文精读: 仔细阅读《Mixture-of-Depths: Dynamically allocating compute in transformer-based language models》。
  • Top-K 路由机制: 深入理解论文中如何通过基于Token的路由决策，选择部分Token进入深层计算，而其他Token直接跳过（Residual Connection）。
  • 容量因子: 理解如何控制参与计算的Token数量比例，以及如何平衡负载。
  • 位置感知与归一化: 学习MoD如何处理被跳过的Token的LayerNorm和残差连接，确保训练稳定性。

学习时间: 2-3周

学习资源:

  • 论文: “Mixture-of-Depths: Dynamically allocating compute in transformer-based language models” (arXiv).
  • 代码库: Hugging Face Transformers库中相关实现的源码（如果有）或作者提供的官方代码库。
  • 视频: 论文作者在学术会议上的报告或YouTube上的解读视频（搜索关键词：Mixture of Depths explanation）。

学习建议: 重点关注论文中关于“Merge”和“Route”的实现细节。尝试画出数据流图，展示Token在MoD层是如何被分流和合并的。理解这种机制如何在不增加推理延迟的情况下增加模型的参数容量。

阶段 4：代码实现与实验复现

学习内容:

  • PyTorch/JAX 实战: 使用PyTorch或JAX从零实现一个包含MoD层的Transformer Block。
  • 前向传播编写: 编写Top-K选择逻辑和Mask操作，实现动态路由。
  • 反向传播验证: 确保梯度能够正确地通过路由机制传播（通常涉及stop_gradient操作）。
  • 性能基准测试: 对比标准Transformer与MoD Transformer在相同FLOPs下的Perplexity指标。

学习时间: 3-4周

学习资源:

  • GitHub: 搜索 “Mixture of Depths implementation” 参考开源社区的实现。
  • 工具: PyTorch官方文档，特别是关于torch.topk和自定义autograd.Function的部分。
  • 数据集: WikiText-2 或 TinyShakespeare 用于小规模验证。

学习建议: 不要一开始就尝试训练几十亿参数的模型。先在一个小型语言

常见问题

什么是 Mixture-of-Depths (MoD) 注意力机制，它与标准的 Transformer 注意力机制有何不同？

Mixture-of-Depths (MoD) 是一种动态分配计算资源的 Transformer 架构变体。在标准的 Transformer（如 GPT-3 或 Llama）中，每一个 Token 在每一层都会通过所有的计算路径，包括多头注意力（MHA）和前馈网络（FFN）。这意味着计算量与上下文长度和层数呈线性关系。

MoD 的核心区别在于它引入了“路由”机制。它并非让所有 Token 在每一层都进行计算，而是通过一个轻量级的路由器（Router）来决定每一层中哪些 Token 是重要的。只有被选中的 Token 才会参与该层的注意力和前馈网络计算，其余的 Token 则直接通过残差连接传递到下一层。这种机制允许模型在保持总参数量不变的情况下，动态地分配 FLOPs（浮点运算次数），从而在不牺牲性能的前提下提高推理效率。

MoD 是如何决定哪些 Token 参与计算的？路由机制是如何工作的？

MoD 使用基于 Top-k 选择的路由机制。具体流程如下：

  1. 路由评分：在每一层计算之前，模型会使用一个轻量级的网络（通常是简单的线性层）根据当前的 Token 隐藏状态计算出一个标量分数。
  2. Top-k 选择：模型会设定一个容量系数，即每一层只处理固定比例（例如前 10%-50%）的 Token。根据计算出的分数，模型选出分数最高的 Top-k 个 Token。
  3. 计算与合并：只有这 Top-k 个 Token 会进入该层的自注意力（Self-Attention）和 MLP 模块进行计算。未被选中的 Token 暂时保持不变。
  4. 残差连接：计算完成后，被处理过的 Token 更新后的状态会通过残差连接加回到原始路径中。对于未被选中的 Token，残差连接直接传递上一层的输出。

这种机制确保了模型总是将计算资源集中在当前上下文中最需要处理的信息上。

Mixture-of-Depths 与 Mixture-of-Experts (MoE) 有什么区别？

虽然两者都涉及“混合”和“路由”的概念，但它们优化的维度和目标不同：

  • Mixture-of-Experts (MoE)：通常应用于**前馈网络（FFN）**层。MoE 将 FFN 替换为多个专家网络，每个 Token 被路由到其中的一个或几个专家。MoE 的主要目的是通过增加参数量来提升模型容量，同时保持每次推理的计算量相对恒定。它关注的是“参数维度的扩展”。
  • Mixture-of-Depths (MoD)：主要应用于层与层之间的动态计算。MoD 决定的是“哪些 Token 在这一层需要工作”。它不一定要增加参数量，而是通过跳过某些 Token 的计算来节省 FLOPs。它关注的是“计算深度的动态分配”。

简单来说，MoE 是“横向”选择不同的专家来处理同一个 Token，而 MoD 是“纵向”选择不同的 Token 来处理同一层。

使用 MoD 技术能带来哪些具体的性能优势或效率提升？

根据论文的实验结果，MoD 带来的主要优势包括：

  1. 推理速度提升：由于每一层只有部分 Token 参与计算，模型的总浮点运算量显著减少。在相同的硬件条件下，MoD 模型比参数量相当的标准稠密模型推理速度更快。
  2. 更优的缩放定律：在相同的计算预算下，MoD 模型通常比标准 Transformer 模型取得更低的验证损失。这意味着如果给你固定的算力预算，训练一个 MoD 模型通常比训练一个更浅或更小的标准模型能获得更好的性能。
  3. 保持性能不变：实验表明，即使只允许模型使用 50% 或更少的 FLOPs，MoD 模型也能达到与全量计算的标准稠密模型相当甚至更好的性能。

MoD 是如何处理上下文信息的？跳过计算是否会导致模型“遗忘”之前的上下文？

这是一个很好的问题。在 MoD 中，即使某些 Token 在某一层没有参与计算（即跳过了 Attention 和 FFN），它们并没有被丢弃，而是通过残差连接直接传递。

  • 残差连接的作用：残差连接确保了信息流不会中断。即使一个 Token 在第 $N$ 层没有参与计算，它依然保留了第 $N-1$ 层的状态。
  • 后续层的补偿：如果一个 Token 在某一层被认为是“不重要”而被跳过，它可能在随后的层中被路由器选中并再次参与计算。
  • 注意力的全局性：对于参与计算的 Top-k 个 Token，它们依然可以通过 Attention 机制看到序列中的所有其他 Token（包括

引用

注：文中事实性信息以以上引用为准；观点与推断为 AI Stack 的分析。

站内链接

相关文章