交错头注意力又是啥? 幼儿园教程来啦!

Meta 搞得新注意力机制很有意思, 揭示了注意力头之间产生联系可能会提升模型的推理能力!

赶紧给大家带来新论文交错头注意力 (Interleaved Head Attention) 的解读. 其实本身很简单啊, 我们先来看传统 transformer 的架构, 每个注意力头都是彼此隔离的, 完全没有联系.

这就会遇到【表达能力上限】问题, 即: 传统模型的一个注意力头往往只能专注于一种主要的语义或逻辑模式, 因为模型的头数是固定且有限的 (比如LLaMA-3-8B只有32个头), 如果一道复杂的数学推理题或代码逻辑, 需要模型同时捕捉 100 种错综复杂的交叉关系, 传统模型怎么办？它只有 32 个头, 分配不过来! 所以做难题的时候能力就下降了. (这也是为啥更大的模型能力强, 有一部分原因就是它的注意力头也更多, DeepSeek-R1 有 128 个头).

那为啥不直接干他一个小目标的注意力头? 因为搞不起哈. 注意力头增加KV Cache就增加, 需要显存就更多, 同时需要计算的算力也更多.

那么有没有成年人的选择呢? 于是 Meta 这个论文灵机一动了下:

三个臭皮匠之所以能PK诸葛亮关键是他们凑到一起商量了!

于是在正式开始计算注意力之前, 模型会先把所有头的数据揉在一起, 形成"伪头"(Pseudo-heads), 然后再把这些伪头在整条数据流上搅拌(交替排布 Interleaving)在一起!

这相当于以前 1 个头只能表达 1 种逻辑关系, 现在经过交叉混合, 1 个头内部能激发出【指数级爆炸的组合模式】, 臭皮匠的鬼点子都混合到一起了!

等会, 混到一起排列组合那是不是鬼点子变多了? 这不还是要显存爆炸?

于是 Meta 又极其聪明地引入了滑动窗口! 还记得咱们上面说的交替排布吗？正是因为混合得足够均匀, 所以就算我们套上滑动窗口(只看局部的词, 不看全图), 窗口也能一眼把各种维度的特征一网打尽. 牺牲了一点点远距离视野, 换来了极深度的逻辑表达能力, 简直血赚, 完美解决了算力膨胀的问题.

那这种方法就没别人能想到吗? 不是的, 之前也有人搞过, 只不过是算完注意力之后再混合的, 会破坏标准算子的数学结构, 用不了 Flash Attention 相当于自废武功. 而 Meta 这个方法好就好在可以完全复用 Flash Attention (因为根本还没开始计算注意力矩阵呢). 所以这个新方法对现有架构特别友好.

最后, Meta团队做了验证, 用了这个新技术后, 增加的参数成本微乎其微, 但带来的收益却是惊人的：
在完全同等算力的情况下, 模型在数学竞赛评测 MATH-500 上提高了 2.8%, 在逻辑推理榜单 GSM8K 上涨了 5.8%! 多键信息检索能力上直接飙升了 10%~20%!
这相当于凭空薅羊毛了一大波推理智商!

最后盲猜一波会不会新的 Muse Spark 大模型就应用了这个架构?
#HOW I AI##MuseSpark##meta#