小红书 算法一面 一
- MHA、MQA 这些核心思想是什么?
## 一、MHA(Multi-Head Attention,多头注意力)核心思想
MHA是Transformer架构(2017年Vaswani等人提出)的**核心组件**,其核心思想可概括为:**通过多个并行的"注意力头",从不同特征子空间同时捕捉输入序列的多维度依赖关系,最后融合各头信息获得更全面的上下文表示** 。
### 1. 核心创新点
- **多视角信息捕捉**:突破单头注意力局限,让模型像人一样从多个"视角"观察输入,每个头专注于不同类型的关系(如语法结构、语义关联、指代关系等)
- **子空间并行建模**:将Q、K、V通过独立线性变换映射到多个低维子空间,每个子空间独立计算注意力,提升特征表达能力
- **信息融合增强**:各头输出拼接后再通过线性变换整合,实现多维度信息的协同捕捉
### 2. 工作流程
1. **线性映射**:输入向量通过三个独立线性层生成Q、K、V矩阵
2. **头部分割**:将Q、K、V各分割为h个并行头部(如8头、16头),每个头部维度为d_k = d_model/h
3. **独立计算**:每个头部独立执行缩放点积注意力计算
4. **结果合并**:拼接所有头部输出,通过最终线性层得到MHA输出
### 3. 数学表达
```
MultiHead(Q,K,V) = Concat(head₁, head₂, ..., headₕ) · W^O
where headᵢ = Attention(Q·W^Qᵢ, K·W^Kᵢ, V·W^Vᵢ)
```
(W^Qᵢ, W^Kᵢ, W^Vᵢ为第i个头的独立投影矩阵,W^O为最终输出投影矩阵)
## 二、MQA(Multi-Query Attention,多查询注意力)核心思想
MQA是2019年提出的MHA变体,核心思想为:**在保持多个查询头(Multi-Query)的同时,让所有查询头共享同一组键(K)和值(V),以牺牲少量表达能力为代价,大幅降低计算与内存开销,提升推理速度** 。
### 1. 核心创新点
- **KV共享机制**:所有Q头共享单组K和V,K和V仅保留1个头,而非h个头
- **计算与内存优化**:将KV缓存大小从O(h·L·d_k)降至O(L·d_k)(L为序列长度),显著减少推理时的内存访问与计算量
- **生成效率提升**:特别适合自回归生成任务(如大模型文本生成),能将decoder生成token速度提升2倍以上
### 2. 工作流程
1. **线性映射**:Q通过h个独立线性层生成h个查询头,K和V仅通过1个线性层生成1组键值对
2. **注意力计算**:每个Q头与同一组K、V计算注意力权重
3. **结果合并**:拼接所有Q头输出,通过线性变换得到最终结果
### 3. 数学表达
```
MultiQueryAttention(Q,K,V) = Concat(head₁, head₂, ..., headₕ) · W^O
where headᵢ = Attention(Q·W^Qᵢ, K·W^K, V·W^V) (∀i)
```
(W^K和W^V为所有头共享的投影矩阵)
## 三、MHA与MQA核心差异对比
| 对比维度 | MHA | MQA |
|---------|-----|-----|
| **核心设计** | Q、K、V均有h个头,相互独立 | Q有h个头,K、V仅1个头,所有Q头共享 |
| **表达能力** | 强,能捕捉更丰富的多维度关系 | 略弱,牺牲部分多子空间表达能力 |
| **计算复杂度** | O(h·L²·d_k) | O(L²·d_k)(降低h倍) |
| **内存开销** | 高,KV缓存随h线性增长 | 低,KV缓存与h无关 |
| **推理速度** | 较慢 | 快,尤其适合长序列生成任务 |
| **适用场景** | 对表达能力要求高的任务(如机器翻译、文本理解) | 对推理速度要求高的生成任务(如大模型部署、实时对话) |
## 四、总结
MHA的核心是**"分而治之"的多视角建模**,通过多个独立注意力头捕捉输入中复杂的多维度关系,是Transformer模型强大表达能力的基础;MQA的核心是**"共享高效"的性能优化**,通过KV共享机制在保持多查询头的同时大幅降低开销,是大模型高效推理的关键技术之一。两者分别代表了注意力机制在表达能力与计算效率两个方向的设计权衡。
- MLA 又是什么?在哪里降维?为什么需要对剩下的部分做 rope?
### 一、核心需求复述
你想了解MLA(多头线性注意力)的核心概念,明确在MLA架构中**降维操作具体发生的位置**,以及为何要对MLA处理后剩余的特征部分应用RoPE(旋转位置编码)。我会从基础概念到具体细节,结合和MHA/MQA的对比,帮你理清这些问题。
### 二、MLA(Multi-Head Linear Attention,多头线性注意力)核心解析
#### 1. MLA的核心思想
MLA是标准MHA的**长序列优化变体**,核心目标是解决MHA的`O(L²)`计算复杂度(L为序列长度)问题——MHA在长序列(如L=10000)场景下计算量会爆炸式增长,而MLA通过**核函数近似**将注意力复杂度降至`O(L)`,同时保留“多头”设计的多视角建模能力。
简单来说:
- MHA是“全局两两比对”(每个token和所有token算注意力),MLA是“全局投影求和”(通过核映射将K/V投影到特征空间后求和,避免两两比对);
- 保留“多头”设计,让每个头独立执行线性注意力计算,兼顾多视角建模和长序列效率。
MLA的核心公式(简化版):
```
# 线性注意力核心:用核映射φ替换MHA的点积,将注意力从两两比对转为求和
LinearAttention(Q,K,V) = φ(Q) · (φ(K)ᵀ · V) / Z
# MLA则是对每个头独立计算后拼接
MLA(Q,K,V) = Concat(head₁, head₂, ..., headₕ) · W^O
where headᵢ = LinearAttention(Q·W^Qᵢ, K·W^Kᵢ, V·W^Vᵢ)
```
#### 2. MLA中降维的具体位置
MLA的降维操作主要发生在**两个关键阶段**,核心目的是控制计算量、聚焦关键特征,且降维逻辑和MHA一脉相承,但适配了线性注意力的特性:
##### 阶段1:多头拆分时的维度投影(核心降维)
这是和MHA一致的基础降维步骤,发生在注意力计算**之前**:
- 输入特征维度为`d_model`(如768、4096),先通过线性层将Q/K/V分别投影到`h × d_k`维度(h为头数,d_k为单头维度,满足`h × d_k = d_model`);
- 例如`d_model=768`、`h=12`时,每个头的`d_k=64`,相当于将768维的全局特征拆分为12个64维的子空间特征——这一步是**最核心的降维**,将高维特征拆分为低维子空间,避免单头计算量过大。
##### 阶段2:线性注意力核映射时的可选降维(优化降维)
这是MLA特有的降维步骤,发生在核函数映射阶段:
- 线性注意力的核函数(如ReLU、Softmax、正余弦核)会将K/V从`d_k`维进一步投影到更低的`d_r`维(`d_r < d_k`,如64→32);
- 目的是进一步降低`φ(K)ᵀ · V`的计算量(从`O(L·d_k²)`降至`O(L·d_k·d_r)`),同时过滤噪声,聚焦核空间的关键特征。
##### 总结降维位置:
```
输入(d_model)
↓ 阶段1:线性投影+拆分(核心降维)
Q/K/V(h × d_k)
↓ 阶段2:核映射(可选降维,MLA特有)
φ(Q)/φ(K)/φ(V)(h × d_r)
↓ 线性注意力计算
各头输出(h × d_r)
↓ 拼接+线性投影
最终输出(d_model)
```
#### 3. 为什么要对MLA剩余部分做RoPE?
先明确:RoPE(Rotary Position Embedding,旋转位置编码)的核心作用是**将位置信息融入注意力计算,且能捕捉token间的相对位置关系**——这是解决Transformer“位置无关”问题的关键。
对MLA剩余部分(通常是注意力计算后的特征、或未参与核映射的Q/K部分)加RoPE,核心原因有3点:
##### 原因1:MLA的线性近似会丢失位置信息
标准MHA通过“两两token比对”天然能隐含位置关系(不同位置的K/V和Q计算的注意力权重不同),但MLA为了降低复杂度,将注意力计算转为“全局求和”(`φ(K)ᵀ · V`是对所有token的K/V求和),**直接丢失了token的位置信息**——如果不加位置编码,MLA无法区分“我吃饭”和“饭吃我”这类语序差异。
##### 原因2:降维会弱化位置表征,RoPE可精准补全
MLA的两次降维(尤其是核映射降维)会压缩特征空间,即使原始输入加了绝对位置编码,降维后位置信息也会被稀释;而RoPE的优势是:
- 不额外增加特征维度(通过旋转矩阵将位置信息融入Q/K的向量中),适配MLA降维后的低维特征空间;
- 捕捉**相对位置**(而非绝对位置),更符合自然语言的语义逻辑(比如“苹果”和“吃”的关系,只和两者的相对距离有关,和绝对位置无关),这对长序列场景(MLA的核心适用场景)至关重要。
##### 原因3:RoPE与MLA的计算逻辑兼容
MLA的核心是“线性计算”,而RoPE的旋转操作也是线性变换(矩阵乘法),不会破坏MLA的`O(L)`复杂度优势;如果用其他位置编码(如绝对位置编码、ALiBi),要么增加计算量,要么无法适配长序列的相对位置需求。
#### 补充:RoPE在MLA中的典型应用方式
RoPE通常加在**MLA降维后的Q/K上**(即阶段1降维后、阶段2核映射前),流程如下:
```
输入(d_model)
↓ 线性投影+拆分(降维到h × d_k)
Q/K/V(h × d_k)
↓ 对Q/K应用RoPE(融入相对位置)
Q_rot/K_rot(h × d_k)
↓ 核映射(可选降维到h × d_r)
φ(Q_rot)/φ(K_rot)/φ(V)(h × d_r)
↓ 线性注意力计算
各头输出(h × d_r)
↓ 拼接+线性投影
最终输出(d_model)
```
### 三、总结
1. **MLA核心**:多头线性注意力,通过核函数近似将MHA的`O(L²)`复杂度降至`O(L)`,兼顾多头多视角建模和长序列计算效率,是长文本处理的核心注意力变体;
2. **MLA降维位置**:核心降维在“多头拆分时的线性投影”(d_model→h×d_k),可选降维在“核映射阶段”(d_k→d_r),均为控制计算量、聚焦关键特征;
3. **RoPE的作用**:MLA的线性求和逻辑会丢失位置信息,降维又会弱化位置表征,RoPE能以线性复杂度融入相对位置信息,且不破坏MLA的效率优势,是MLA适配长序列任务的关键补充。
