低秩注意力残差

注意力残差（AttnRes）用 大语言模型（LLM）中先前子层输出的深度注意力替换固定残差和，但使用每个输出作为全维键和值。这将路由与表示结合起来，并使计算深度路由分数的成本与隐藏宽度 $d$ 成比例。我们提出了低秩注意力残差（LR-AttnRes），它在使用 $r$ 维键（其中 $r < d$）进行路由时保留全维残差值。 LR-AttnRes 使用每个值的最后 $r$ 维度作为路由键，减少总的剩余侧 FLOP，同时仍然提高性能。对块数量 ($N$) 和 $r$ 的全面扫描表明，深度路由可以在尺寸远小于模型宽度的情况下有效。在 $1$B 和 $4$B 参数（$r = d/4$）下，LR-AttnRes 比标准 AttnRes 实现了更低的最终验证损失、更高的平均下游精度和更高的测量训练步骤吞吐量。我们还提供支持标准和低秩路由的融合内核。我们发布所有代码、内核和所有经过训练的模型，以方便未来的研究。