想象一下，你正在尝试教一个机器人写故事。长期以来，主要有两种方法可以实现这一目标，但两者都存在一个重大缺陷。
两种旧方法
- “单人表演”（仅解码器）： 这就像一个学生读完一本书后立即尝试写出下一句话。他们写作速度非常快，但在写作时必须在大脑中记住迄今为止读过的所有内容。如果故事很长，他们的大脑（内存）就会过载，导致速度变慢。
- “双人团队”（编码器 - 解码器）： 这就像有一个“读者”和一个“写手”。读者先读完整本书，做笔记，然后将笔记交给写手。这节省了写手的大脑算力，但读者非常挑剔。他们只记录大约 15% 的单词（即“被破坏”的部分），而忽略其余部分。这意味着团队学习缓慢，因为他们忽略了故事的大部分内容。
新方案：基于块的双解码器
本文的作者提出了一种新的团队结构，试图兼得两者的优点。他们称之为基于块的双解码器。
以下是其工作原理，使用一个简单的类比：
想象你在阅读一部长篇小说，但不是逐字阅读，而是将其拆分为块（如章节或场景）。
- 上下文解码器（读者）： 这部分读取直到某个点之前的整个故事。它就像一个快速阅读者，扫描前几章以理解背景和人物。由于它只读取“过去”，因此无需一次性将整个书的内容保存在内存中。它会生成一个摘要。
- 生成解码器（写手）： 这部分接收来自读者的摘要以及当前的文本块。它负责撰写故事的下一部分。
魔法技巧：“双重因果”块
秘诀在于他们如何分割故事。他们将文本切成块（chunks）。
- 在单个块内部，写手可以查看该块中的所有单词（就像小组讨论）。
- 但在查看之前的块时，写手只能查看读者提供的摘要，而不能查看原始单词。
为什么这很重要？
- 不再浪费学习机会： 在旧的“双人团队”中，读者忽略了 85% 的单词。在这个新系统中，每一个单词都有机会被学习。模型会对每个 token 进行“评分”，使其学习得更快、更聪明。
- 超高效内存： 当机器人实际撰写故事（推理）时，它不需要记住整本书。它只需要记住来自读者的摘要和当前块。这将所需的内存减少了约三分之二。这就像从背包里背着一座图书馆，变成了只带一张索引卡。
- 速度： 由于“读者”部分仅在开始时运行一次，然后处于空闲状态，因此“写手”部分要轻量得多，速度也更快。这就像在赛跑开始时使用重型引擎，但在冲刺阶段使用轻量级、空气动力学的车身。
他们发现了什么？
研究人员将这种新架构与旧架构进行了测试。
- 训练： 新模型的学习效果几乎与“单人表演”（速度的黄金标准）一样好，并且远优于旧的“双人团队”。
- 写作（推理）： 当需要实际生成文本时，新模型表现出色。它使用的计算机内存显著少于旧的“双人团队”，速度也更快，同时依然非常智能。
结论
该论文声称，通过将工作拆分为两个解码器并将文本拆分为块，他们创建了一个能够从每一个单词中学习（不同于旧的高效模型）的模型，但在需要写作时不会因内存问题而受阻（不同于旧的快速模型）。这是一种获得高性能机器人的方法，而无需依赖庞大且昂贵的计算机来运行。
技术摘要：基于块的双解码器
问题陈述
本文探讨了 Transformer 架构设计中训练效率与推理效率之间的根本权衡。仅解码器模型（Decoder-only models）提供了卓越的训练效率和可扩展性，但推理成本高昂，尤其是在 KV 缓存内存和每 token 计算方面。相反，编码器 - 解码器模型（Encoder-decoder models）在推理时能显著节省资源（减少 KV 缓存），但由于预训练目标（如 span corruption，其中仅约 15% 的 token 产生损失信号）导致的稀疏监督，以及阻碍高效批处理的动态序列长度，其表现受限。现有的替代方案（如 PrefixLM）试图弥合这一差距，但仍有许多 token 未被训练，且未能实现编码器 - 解码器架构的全部效率。
方法：基于块的双解码器
作者提出了一种名为基于块的双解码器（Block-Based Double Decoders）的新颖架构，该架构结合了两个解码器堆栈，以实现全损失监督和推理效率。
架构结构：
- 上下文解码器（Context Decoder）： 一个标准的因果仅解码器 Transformer，处理完整输入序列，并为每个 token 输出因果潜在表示（ht）。
- 生成解码器（Generation Decoder）： 第二个解码器堆栈，接收三个输入：来自上下文解码器的因果潜在表示、原始 token 序列以及块划分。
- 双重因果块级掩码（Doubly-Causal Block-Based Masking）： 输入序列被分割为连续的块。在块内，生成解码器使用因果自注意力；在块间，它使用全交叉注意力，关注前序块中上下文解码器的潜在表示。这种结构形成了“双重因果”机制，即注意力流在块内向前流动，并从前序块流向当前块。
训练目标：
- 与掩码 token 的 span corruption 不同，该方法确保输入序列中的每个 token在每次前向传播中恰好贡献一次损失。
- 打包后序列长度保持静态，消除了对动态批处理和 token 填充的需求，从而提高了吞吐量。
- 作者通过分别计算自注意力和交叉注意力 logits，并在后处理阶段通过 log-sum-exp 归一化进行组合（使用 PyTorch 的 FlexAttention）来实现注意力机制，因为目前尚无统一的双重键注意力实现。
推理机制：
- 在推理时，上下文解码器仅对提示词运行一次（预填充）。由于它是标准的因果解码器，其自身输出不需要 KV 缓存（仅生成解码器需要缓存）。
- 生成解码器随后自回归地生成 token，关注缓存的上下文潜在表示及其自身的先前输出。
- 这种分离允许复用仅解码器模型中常见的前缀级 KV 缓存，同时保持编码器 - 解码器架构的内存优势。
主要贡献
- 新颖架构： 提出了基于块的双解码器，利用两个解码器堆栈将上下文处理与 token 生成解耦。
- 双重因果掩码： 一种新的注意力掩码，在保持静态序列长度以实现高效批处理的同时，实现了全 token 监督（100% 损失密度）。
- 推理优化： 与仅解码器模型相比，该架构将 KV 缓存内存和每 token 计算减少了约2/3（假设上下文解码器占 2/3，生成解码器占 1/3），且不牺牲预填充缓存能力。
- 训练效率： 该方法避免了 span-corruption 训练的编码器 - 解码器中固有的稀疏监督和动态批处理问题。
实验结果
作者进行了扩展定律实验，比较了仅解码器模型、标准编码器 - 解码器模型（SED）和基于块的双解码器模型，参数规模涵盖 625 万至 1 亿，token 预算涵盖 6250 万至 10 亿。
- 性能： 双解码器模型显著优于 SED 基线，并紧密跟踪仅解码器模型。在最大规模下（1 亿参数，10 亿 token），双解码器的评估损失仅比匹配的仅解码器模型差约0.2 nats，而 SED 模型比两者都差约0.7 nats。
- 扩展行为： 发现 SED 模型处于“数据受限状态”，由于稀疏监督，增加参数并未提升性能。相比之下，双解码器和仅解码器模型遵循标准扩展定律，在达到容量限制时趋于平缓。
- 计算权衡： 虽然由于额外的 KV 投影，双解码器在特定配置下的训练 FLOPs 略有增加（比仅解码器模型多约 2.4%），但这被推理优势所抵消。
意义与主张
本文主张，基于块的双解码器成功解决了训练效率与推理效率之间的历史二分法。通过将全损失监督恢复到类似编码器 - 解码器的结构中，作者证明了可以实现编码器 - 解码器的推理效率（具体为减少的 KV 缓存和每 token 延迟），同时保留仅解码器模型的训练效率和扩展特性。
作者强调，适度的训练时间损失差距（约 0.2 nats）是实现架构分离的必要代价，该分离带来了显著的推理时间节省。他们认为，这使得该架构特别适用于边缘设备和内存受限的场景，同时也开启了新的研究方向，例如在预训练期间进行潜在链式思维推理，以及将循环 Transformer 专门应用于上下文解码器。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题，再试一次？
无垃圾邮件，随时退订。