混合深度架构(Mixture-of-Depths, MoD)是一种让Transformer在深度维度上动态分配计算的范式。每个MoD层内置路由器,根据token重要性决定其通过该层还是跳过,从而实现每个token经历不同数量的层。简单token走捷径快速通过,复杂token走完整路径深度处理,在保持质量的同时显著减少计算量。
动态深度是MoD的核心特征。通过早退机制,简单token可以在浅层就获得足够好的表示直接输出;跳层机制让token绕过不需要的层;可变深度使不同token根据自身复杂度经历不同层数。这种按需分配深度的策略让计算资源精准匹配每个token的实际需求。
路由机制是MoD的决策核心。每层包含一个轻量路由网络(小MLP),对每个token输出一个标量权重,经Softmax归一化后通过Top-K选择决定哪些token通过该层。负载均衡损失确保token不会过度集中在某些层,维持训练稳定性和计算效率的平衡。
MoD与MoE(混合专家)是互补的稀疏计算范式。MoE在宽度维度路由——每个token选择不同的专家子网络;MoD在深度维度路由——每个token选择是否通过某一层。MoE增加宽度选择,MoD增加深度选择。两者可以组合使用,在宽度和深度上同时实现稀疏化。
MoD通过跳过不重要的token计算实现显著的FLOPs节省。每层仅处理Top-k个token(如50%),其余直接跳过。对于N层Transformer,总FLOPs约为标准的50-80%取决于路由比例,同时保持相当的模型质量。实际推理速度还受益于更少的内存访问和更短的序列处理路径。
MoD的训练需要特殊策略来保证路由器学到有意义的分配。辅助负载均衡损失防止token坍缩到少数路径;温度调度从高温探索逐渐降温到确定性路由;课程学习先训练密集模型再引入路由;梯度稳定性措施确保路由决策的可微性和训练收敛。
MoD在精度-计算权衡上展现出色特性。在相同FLOPs预算下,MoD比标准Transformer获得更高精度;在相同精度下,MoD所需计算更少,位于Pareto前沿。消融实验表明路由比例、层数选择和路由器设计是关键超参数,合理配置可在不损失质量的前提下节省大量计算。
MoD的动态计算特性使其在多种实际场景中大放异彩。高效推理服务通过按需计算降低延迟和成本;边缘部署利用稀疏计算在受限硬件上运行大模型;实时系统借助早退机制满足严格延迟要求;大模型加速在保持质量的同时扩展到更大规模。