SuperAPI AI 工具集

explicit-multi-head-attention-inter-head

通用与其它

查看项目

更新于 10-16 01:34

技能介绍

用途:在 Transformer 注意力层中显式建模头间交互。核心是两项改动:头级线性组合(HLC),用可学习矩阵 W_lc 对 key/value 各头做线性重组;头级 RMSNorm,在输出投影前稳定重组后的表示。

主要能力:提升预训练鲁棒性,允许约 3 倍更大学习率(如 3e-3 对 1e-3);通过 SVD 把 K/V 投影分解为 h' 个组件头加 HLC 重建,推理时只缓存组件头,h'=h/2 时 KV 缓存减半,知识类任务损失很小,奥数级数学约降 3.6%。

使用方式:定位标准 MultiHeadAttention,加入 W_lc_K、W_lc_V(形状 h'×h,初始化为单位阵以保证等价于标准 MHA),在 K/V 上做 einsum 混合,再走缩放点积注意力,输出经头级 RMSNorm 后投影。压缩时对预训练 K/V 权重做 SVD,取前 h'*d_k 个奇异分量作新投影,HLC 矩阵由 Vt 构造,随后做 1–5B token 的恢复微调。

适用场景:自研注意力层、预训练稳定性调优、推理 KV 缓存压缩、继续预训练注入 MEA。

限制:HLC 有少量额外计算;压缩超过 50% 收益递减;数学推理对头容量更敏感;仅面向注意力架构;论文验证到约 7B 参数,更大规模未确认。具体实现细节与超参以原论文为准。