论文

多头注意力残差(MHAR)

Multi-Head Attention Residuals

模型架构Transformer

摘要

Transformer通过单一的加性残差流在深度方向传播信息:每个子层只能读取最近的状态。注意力残差(attention residuals)放宽了这一限制,让每个子层可以通过一个可学习的softmax执行注意力读取。然而,这种读取使用的是在整个宽度上共享的单一查询,因此每个特征子空间都必须通过同一个分布来读取深度历史。这种被迫妥协的代价随各子空间在“读取哪些层”上的分歧程度而增长,而分歧又随模型宽度增长。我们提出多头注意力残差(Multi-Head Attention Residuals, MHAR):将路由查询重塑为H个按子空间划分的头,每个头在深度历史上拥有各自的softmax。读取由此变为块对角形式,重塑不增加任何参数、计算量可忽略不计,且H = 1时精确还原注意力残差。在经过质量过滤、偏重STEM与代码的去重Nemotron退火语料上从头训练,MHAR在100M、350M和1B规模上将验证损失较标准Transformer分别改善-0.061、-0.149和-0.140。它在所有设置中都是四种方法中的最优,且增益从100M到更大规模不断增大。头数是一个真实的设计维度而非可随意扭动的旋钮:验证损失关于H呈U形,各规模下的平坦最优区在H = 4或H = 8。大规模模型我们采用H = 8;超过该点继续过度切分(H = 16)会稳定地吐回一部分增益。对训练后查询的直接探测证实,习得的子空间分歧是底层驱动因素。融合的Triton路由内核将注意力残差训练吞吐量从基线的0.2-0.5倍提升至0.55-0.88倍,同时保持接近基线的峰值显存。基于delta注意力残差的保恒等转换支持8B模型的中训练(mid-training),在GSM8K和GPQA上分别带来+3.2和+3.1的提升。

多头注意力残差(MHAR):论文配图
图 8:100M 处 MHAR 的验证损失作为路由头 HH (x) 和 KV 头 (y) 的函数。越低越好。金星标志着每 KV 的最佳 HH;红色虚线是头匹配对角线H=H{=} KV。在 100M 处,KV= 4\,{=}\,4 和 KV= 8\,{=}\,8 最佳值位于 H=H{=} KV 对角线上。表 1 的 100M 架构,以 5×10−45\times 10^{-4} 进行训练。