论文

使用与范数无关的残差网络缩放自适应深度

Scaling Adaptive Depth with Norm-Agnostic Residual Networks

模型架构Transformer

摘要

残差架构在深度学习中无处不在,但它们受到微妙的结构限制:残差流的范数可以随着深度而快速增长。结果,来自后面层的更新相对于累积的残余状态变得很小。这减少了它们对表示的影响,并限制了深度缩放模型的好处。为了解决这个问题,我们引入了NAG,一种与范数无关的残差架构,它将残差流中的幅度与方向信息分开,在整个深度中保留有意义的层贡献,并防止后续更新被残差范数增长系统性地抑制。重要的是,NAG 仅引入了少量的附加参数,并且依赖于易于内核融合的简单操作,从而在实践中保持了训练效率。我们表明,该架构优于基线 Transformer,其增益随着深度的增长而大幅增加,从而能够有效训练更深的模型。与规范无关的公式还导致了一种可解释的深度混合 (MoD) 机制,该机制自适应地跳过注意力层和 MLP 层。除了用作 后训练 精度计算权衡之外,该机制还可以用作预训练时间扩展策略:在 iso-FLOP 训练下,通过减少每个词元前向传递成本节省的计算可以重新投入到更多词元的训练中,同时保持总参数计数和 KV 缓存预算固定。在我们的实验中,大约 20%-25% 的适度深度混合率与同等训练计算下的全深度基线性能相匹配,同时大大减少了执行层参数和前向传递 FLOP 的数量。这些结果将深度稀疏性确定为固定计算训练的新缩放轴,从而实现非常深且 FLOP 高效的模型。