用低维键路由改进大模型残差连接,提升训练效率与性能
Low-Rank Attention Residuals

- 用低维键替代全维键进行深度路由,解耦路由与内容
- 在相同模型宽度下,低维键(r≪d)仍可实现更优验证损失
- 适合追求高效训练的大模型研究者,代码已开源
注意力残差将固定残差求和改为对前序子层输出的深度注意力路由,但使用全维键和值,导致路由与表示耦合,路由得分随隐藏维度$d$增长。我们提出低秩注意力残差(LR-AttnRes),保留全维残差值,仅用$r$维键($r \ll d$)进行路由。投影式LR-AttnRes从现有输出投影中学习低秩键,解耦路由与内容,测试中表现最佳。切片式LR-AttnRes直接使用值的最后$r$维作为路由键,移除额外键投影路径,降低残差侧计算量,同时提升性能。全面消融实验表明,深度路由可在远低于模型宽度的维度下有效工作。代码与模型已公开,便于后续研究。
原文摘要 · Abstract (English)
Attention Residuals replace the fixed residual sum with depthwise attention over previous sub-layer outputs in large language models (LLMs), but use each output as both a full-dimensional key and value. This couples routing with representation and makes depth-routing scores scale with the hidden width $d$. We propose Low-Rank Attention Residuals (LR-AttnRes), which keep full-dimensional residual values while using $r$-dimensional keys, with $r \ll d$, for routing. Projected LR-AttnRes emits learned low-rank keys from existing output projections, decoupling routing from residual content and achieving the best validation loss among the variants tested. Sliced LR-AttnRes uses the last $r$ dimensions of each value as the routing key, removing the auxiliary key-projection path and reducing residual-side FLOPs while still improving performance. Comprehensive sweeps show that depthwise routing can be effective with far fewer dimensions than the model width. We release code and models to facilitate future research.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。