统一视角看 Attention 与 MoE

注: 本文 原始 idea 来自于香港大学(HKU) 黄毅老师 (主页: Ngai Wong), 与本人共同讨论, 形成此文.

---

Update: 简易版报告在我主页, 可以访问

https://wutaiqiang.github.io/pdf/Unified_view_for_Attention_and_MoE.pdf​wutaiqiang.github.io/pdf/Unified_view_for_Attention_and_MoE.pdf

---

LLM 几乎已经成为了 NLP 社区的唯一关键词, 同时也牢牢占据了 CV 和 Robotic 的核心词. LLM 的主流(近乎唯一)架构是Transformer 结构, 主要包含 Attention 和 FFN 两个模块. 其中 Attention 模块主要是依赖 QKV 三个部分, 实现近似于 (QK^T)V 的结果. 与此同时, MoE 也是当下主流的 LLM 架构之一, 核心思路是增大参数量的同时, 维持计算量不变. Deepseek 公司基于 MoE 架构有很多落地的模型.

不过, 细看Attention 和 MoE,其实有很多的相似之处.

基础定义

记输入为 𝑥1,𝑥2,...,𝑥𝑀 , 我们可以看看这两种结构的相似之处.

Attention 结构主要是:

核心思路是, 经过线性投影以后获得 QKV, 然后基于 QK + softmax 来计算 Attention, 最后与 V 相乘, 获得最终的输出. 某个角度来看, Attention 结构的作用其实就是使用加权和的形式来聚合不同 token 的信息.

类似地,我们可以写出 MoE 地结构表达式:

简单来说, MoE 其实是通过线性投影来学习一个 Router, 基于 Router 来聚合不同 expert 的信息. 本质上来说, router 的输出经过 top-k+softmax 来学习一个加权求和的权重, 最终聚合不同 token 的信息.

统一视角

看到这里, 想必你也发现了两者的相似之处, 那就是都是学习一个权重来聚合信息.

公式角度来看, 形式就是:

对于 Attention 和 MoE 二者来说, 区别在于 f g 具体定义不同, 但是本质都是 f是权重, g 是学习到的新特征.

两者的相似之处: 1) 二者都使用了 softmax 来归一化权重f, 本质上也是希望学习稳定. 2) 二者本质都是对新特征的加权和. 3) 二者都是动态权重, 根据输入来自适应地聚合信息.

二者的区别在于: 1) Attention 是聚合不同 token 的信息, MoE 是聚合不同地 Expert 2) Attention 一定会用上所有的 token (Softmax 地结果为正数,故而所有的 token 都会用上), MoE 是一种稀疏的选择,只会选择部分结果.

进一步研究

既然能把 Attention 和 MoE 用统一的视角来看, 那么我们可以用一个新的角度来看待二者的变体.

为了处理长文本或者长输入, Sparse Attention 是一种常用的设计, 如对 attention 的范围做一个加窗的处理. 可以参考这个回答

有哪些令你印象深刻的魔改transformer？118 赞同 · 4 评论回答

其实从 MoE 的角度来看,这种改进都特别自然. 其实就是使用上了 MoE 的稀疏化聚合这种形式, 应用到 Attention 之中, 不过现有的 sparse attention 主要是基于位置关系或者语义相似度. 当然,我们也可以学习 MoE 的形式来经过一个线性层学习这个稀疏的权重分布.

类似地, MoE 的路由仅仅使用当前 Token 的预测,来计算所有 expert 上的权重分布. 类似于 Attention 的操作, 这个 router 也可以使用其他的 token 信息来加强, 这也用到了 Yuan 2.0-M32 这个大模型上.

Linsight：从Yuan2.0到Yuan2.0-M325 赞同 · 0 评论文章