标准 Transformer 会让序列里的每个 Token 走过相同层数。逗号、常见词和决定答案方向的关键词,在每一层都获得同样一次注意力与 MLP 计算。深度混合给部分层设容量,只让路由器选中的 Token 进入这一层的主要计算,其余 Token 沿残差连接直接往后走。
我查看了 Google DeepMind 团队的 Mixture of Depths 论文。它保留固定的总计算预算,又让预算在不同 Token 和不同深度上动态分配。这里的 mixture 指一条路径经过当前 Transformer 块,另一条路径跳过它,和把 Token 送往不同专家的混合专家模型有明显区别。
固定容量让总计算量提前可知
每个带路由的块先给序列中各 Token 算一个分数,再选择得分最高的前 k 个。入选 Token 参加自注意力和 MLP,随后把结果写回残差流;其余位置保留原表示。k 在运行前已经设定,张量大小和总 FLOPs 因此可以预测,入选的是哪些 Token 则会随上下文改变。
论文采用类似专家选择的路由方式,由当前计算块挑选固定数量的 Token。这样每层恰好处理 k 个位置,无需再加负载平衡损失。若让每个 Token 自己选路径,许多 Token 可能同时挤向计算路径,系统还要处理容量溢出与丢弃顺序。
随机跳层行不通。结果差别很大。论文用随机路由做对照,表现明显差于学习得到的路由。作者在一组实验中隔层放置路由块,每个路由块只处理序列的 12.5%,其余完整块仍处理全部 Token。一个 2.2 亿参数的变体达到与同计算预算基线相当的表现,训练步进速度快 66%。继续把容量压得更低,质量开始下降。模型会答错。
路由必须学。完整计算层仍在给序列提供共同的信息交换机会。路由层省下部分注意力与 MLP 工作,模型再通过训练学会哪些位置值得进入。论文观察到,经常进入计算块的 Token 与更高熵的输出预测相关,但作者把它作为初步分析,没有据此宣称路由器已经学会通用的难度判断。
自回归生成还要解决未来信息
训练一整段序列时,top k 可以在全部位置中比较分数。逐词生成看不到未来 Token,直接照搬这套选择会使用尚未出现的信息。论文为生成阶段训练了一个因果预测器,模仿训练时的 top k 决策。其路由预测很快达到 97% 以上准确率,切换后性能下降很小。
论文在相同训练 FLOPs 下找到了一些损失更低、每次前向计算更少的深度混合模型,并报告生成步进速度最高快 50% 以上。实验模型覆盖 6000 万到 30 亿参数,和今天更大规模的生产模型仍有距离。动态路由还会改变内存访问与算子形状,纸面 FLOPs 减少多少,不能直接等同于线上延迟减少多少。
采用这类结构需要从预训练阶段纳入路由,无法给现成模型打开一个开关就得到相同结果。评估时应固定训练总量,比较损失、下游任务、真实步进时间和路由稳定性。若模型在关键词、数字或代码符号上频繁跳过必要计算,平均速度再好也不够用。固定容量解决了预算可控,Token 选得准仍是整套方法的核心。