注意力机制演进史：从 Bahdanau 到 Flash Attention 3

1. 注意力机制的起源：Bahdanau Attention (2014)

注意力机制的故事始于机器翻译。2014 年，Bahdanau 等人提出了 Additive Attention，解决了 seq2seq 模型中固定长度编码瓶颈的问题。

核心思想：解码器的每一步不再只依赖一个固定的上下文向量，而是"关注"源序列的不同部分。

# Bahdanau Additive Attention 的 PyTorch 实现
class BahdanauAttention(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.W_query = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_key = nn.Linear(hidden_size, hidden_size, bias=False)
        self.V = nn.Linear(hidden_size, 1, bias=False)
    
    def forward(self, query, keys, values):
        """
        query: (batch, hidden) - 解码器当前状态
        keys: (batch, src_len, hidden) - 编码器所有隐状态
        values: 同 keys
        """
        # 扩展 query 以与 keys 对齐
        query_expanded = query.unsqueeze(1)  # (batch, 1, hidden)
        
        # 加性注意力: score = V^T * tanh(W_q * q + W_k * k)
        scores = self.V(
            torch.tanh(
                self.W_query(query_expanded) + self.W_key(keys)
            )
        )  # (batch, src_len, 1)
        scores = scores.squeeze(-1)  # (batch, src_len)
        
        # 注意力权重
        attn_weights = F.softmax(scores, dim=-1)  # (batch, src_len)
        
        # 加权求和
        context = torch.bmm(attn_weights.unsqueeze(1), values)  # (batch, 1, hidden)
        context = context.squeeze(1)  # (batch, hidden)
        
        return context, attn_weights

1.1 Luong Attention (2015)

Luong 提出了多种变体，其中 Dot-Product Attention 影响最为深远：

变体	公式	特点
Dot	$\text{score}(h_t, h_s) = h_t^\top h_s$	最简单，要求维度相同
General	$\text{score} = h_t^\top W_a h_s$	引入可学习矩阵
Concat	$\text{score} = v_a^\top \tanh(W_a [h_t; h_s])$	同 Bahdanau
Location-based	$\text{score} = W_a h_t$	仅依赖目标位置

2. Self-Attention 与 Transformer (2017)

Vaswani 等人的 “Attention Is All You Need” 将注意力机制推向了极致：抛弃 RNN/CNN，纯靠注意力构建序列模型。

2.1 Scaled Dot-Product Attention

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (batch, heads, seq_len, d_k)
    K: (batch, heads, seq_len, d_k)
    V: (batch, heads, seq_len, d_v)
    """
    d_k = Q.size(-1)
    
    # 计算注意力分数
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 应用掩码（用于因果注意力或padding）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # Softmax 归一化
    attn_weights = F.softmax(scores, dim=-1)
    
    # 加权求和
    output = torch.matmul(attn_weights, V)
    
    return output, attn_weights

为什么需要缩放因子 $\sqrt{d_k}$？ 当 $d_k$ 较大时，点积结果的方差增大，softmax 会进入梯度饱和区域。缩放因子将方差恢复到 1，保证梯度稳定。

2.2 Multi-Head Attention

多头注意力让模型能同时关注不同表示子空间的信息：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Q/K/V 投影矩阵
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        
        # 输出投影
        self.W_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 线性投影 + 分头
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # Scaled Dot-Product Attention
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        
        # 输出投影
        output = self.W_o(attn_output)
        
        return output

2.3 注意力复杂度分析

操作	时间复杂度	空间复杂度
QK^T	O(n²·d)	O(n²)
softmax	O(n²)	O(n²)
AV	O(n²·d)	O(n·d)
总计	O(n²·d)	O(n²)

二次复杂度是长序列的核心瓶颈，后续大量工作围绕这一问题展开。

3. 高效注意力变体

3.1 Sparse Attention

通过限制注意力的感受野来降低复杂度：

全局注意力 (Full):     每个token关注所有token → O(n²)
    
局部窗口 (Local):      每个token只关注邻近w个token → O(n·w)
    
稀疏模式 (Sparse):     组合局部+全局+随机连接 → O(n·log(n))
    
分块注意力 (Block):    将序列分块，块内全注意力+块间稀疏 → O(n·√n)

3.2 Linear Attention

将 softmax 替换为可分解的核函数，实现线性复杂度：

$$\text{Attention}(Q,K,V) = \phi(Q)(\phi(K)^\top V)$$

def linear_attention(Q, K, V, feature_map=None):
    """
    线性注意力: 通过特征映射分解 softmax
    复杂度: O(n·d²) 而非 O(n²·d)
    """
    if feature_map is None:
        # ELU + 1 是常用的特征映射
        feature_map = lambda x: F.elu(x) + 1
    
    Q_feat = feature_map(Q)  # (batch, heads, n, d)
    K_feat = feature_map(K)  # (batch, heads, n, d)
    
    # 关键: 先计算 K^T V (d×d), 再与 Q 相乘
    # 这样避免显式计算 n×n 的注意力矩阵
    KV = torch.matmul(K_feat.transpose(-2, -1), V)  # (batch, heads, d, d)
    output = torch.matmul(Q_feat, KV)  # (batch, heads, n, d)
    
    # 归一化
    K_sum = K_feat.sum(dim=-2, keepdim=True)  # (batch, heads, 1, d)
    normalizer = torch.matmul(Q_feat, K_sum.transpose(-2, -1)) + 1e-6
    output = output / normalizer
    
    return output

方法	复杂度	优点	缺点
标准注意力	O(n²d)	表达力最强	长序列不适用
Linear Attention	O(nd²)	线性序列长度	表达力下降
Sparse Attention	O(n√n)	保留局部精度	实现复杂
Flash Attention	O(n²d)	IO优化、无近似	复杂度不变
RWKV	O(nd)	RNN级效率	并行训练受限

4. Flash Attention：IO 感知的注意力

4.1 核心洞察

Flash Attention (Dao, 2022) 的核心发现是：注意力的瓶颈不在 FLOPs，而在 HBM 读写。

GPU 内存层次:
┌──────────────────┐
│  HBM (显存)       │  ~40-80 GB, 带宽 ~2 TB/s
│  ┌────────────┐  │
│  │ SRAM (L2)  │  │  ~40 MB, 带宽 ~19 TB/s
│  │  ┌──────┐  │  │
│  │  │ Regs │  │  │  ~256 KB, 即时访问
│  │  └──────┘  │  │
│  └────────────┘  │
└──────────────────┘

标准注意力: 需要在 HBM 中读写 n×n 的注意力矩阵
Flash Attention: 在 SRAM 中分块计算，避免物化 n×n 矩阵

4.2 分块算法

# Flash Attention 的分块算法伪代码
def flash_attention(Q, K, V, block_size=128):
    """
    分块注意力: 避免物化完整 n×n 注意力矩阵
    使用 online softmax 技巧
    """
    B, H, N, D = Q.shape
    scale = 1.0 / math.sqrt(D)
    
    # 输出和累积量
    O = torch.zeros_like(Q)
    l = torch.zeros(B, H, N, 1, device=Q.device)  # 行和累积
    m = torch.full((B, H, N, 1), float('-inf'), device=Q.device)  # 行最大值
    
    # 分块遍历 K, V
    for j in range(0, N, block_size):
        K_block = K[:, :, j:j+block_size, :]  # (B, H, block, D)
        V_block = V[:, :, j:j+block_size, :]
        
        # 在 SRAM 中计算注意力分数
        for i in range(0, N, block_size):
            Q_block = Q[:, :, i:i+block_size, :]
            
            # 计算当前块的分数
            S = torch.matmul(Q_block, K_block.transpose(-2, -1)) * scale
            # (B, H, block_i, block_j)
            
            # Online Softmax: 更新行最大值
            m_new = torch.maximum(m[:, :, i:i+block_size], S.max(dim=-1, keepdim=True).values)
            
            # 重新缩放
            alpha = torch.exp(m[:, :, i:i+block_size] - m_new)
            beta = torch.exp(S - m_new)
            
            # 更新累积和
            l[:, :, i:i+block_size] = l[:, :, i:i+block_size] * alpha + beta.sum(dim=-1, keepdim=True)
            
            # 更新输出
            O[:, :, i:i+block_size] = O[:, :, i:i+block_size] * alpha + torch.matmul(beta, V_block)
            
            # 更新最大值
            m[:, :, i:i+block_size] = m_new
    
    # 最终归一化
    O = O / l
    
    return O

4.3 Flash Attention 2 与 3

Flash Attention 2 (Dao, 2023) 的改进：

减少非矩阵乘法 FLOPs（从 2x 降至 1.5x）
优化并行化策略：沿序列维度并行
更好的 GPU 占用率

Flash Attention 3 (2024) 针对 H100 GPU 的优化：

利用 TMA (Tensor Memory Accelerator) 异步加载
WGMMA (Warp Group Matrix Multiply-Accumulate) 指令
FP8 支持，吞吐量翻倍

版本	GPU	序列长度	吞吐量 (TFLOPS)	加速比
标准注意力	A100	4K	12	1.0x
Flash Attn 1	A100	4K	72	6.0x
Flash Attn 2	A100	4K	118	9.8x
Flash Attn 2	H100	8K	235	-
Flash Attn 3	H100	8K	740	3.1x vs FA2

5. 注意力机制的变体设计

5.1 注意力变体分类

                    注意力机制
                   /          \
           稠密注意力          稀疏注意力
          /     \             /        \
    标准注意力  线性注意力  局部稀疏    全局稀疏
         |        |          |          |
    Flash Attn  Performer  Longformer   BigBird
    Sliding      Linformer   Sparse      Routing
    Window       RWKV        Transformer  Transformer

5.2 滑动窗口注意力

Mistral 和 Gemma 采用的滑动窗口注意力：

class SlidingWindowAttention(nn.Module):
    def __init__(self, d_model, num_heads, window_size=4096):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.window_size = window_size
        self.W_qkv = nn.Linear(d_model, 3 * d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        B, N, D = x.shape
        qkv = self.W_qkv(x).reshape(B, N, 3, self.num_heads, self.d_k)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2)  # (B, H, N, d_k)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # 构建滑动窗口掩码
        mask = torch.ones(N, N, dtype=torch.bool, device=x.device)
        for i in range(N):
            start = max(0, i - self.window_size + 1)
            mask[i, start:i+1] = False  # False = 允许注意力
        
        # 应用掩码的标准注意力
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        scores = scores.masked_fill(mask, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)
        
        output = output.transpose(1, 2).reshape(B, N, D)
        return self.W_o(output)

5.3 GQA (Grouped-Query Attention)

Llama 2 70B 引入的 GQA 是 MHA 和 MQA 的折中：

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.num_groups = num_heads // num_kv_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, num_heads * self.d_k)
        self.W_k = nn.Linear(d_model, num_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, num_kv_heads * self.d_k)
        self.W_o = nn.Linear(num_heads * self.d_k, d_model)
    
    def forward(self, x):
        B, N, D = x.shape
        q = self.W_q(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, N, self.num_kv_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, N, self.num_kv_heads, self.d_k).transpose(1, 2)
        
        # 扩展 KV 头以匹配 Q 头
        # 每个 KV 头被 num_groups 个 Q 头共享
        k = k.repeat_interleave(self.num_groups, dim=1)  # (B, num_heads, N, d_k)
        v = v.repeat_interleave(self.num_groups, dim=1)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)
        
        output = output.transpose(1, 2).reshape(B, N, -1)
        return self.W_o(output)

方案	Q头数	KV头数	KV Cache	质量	代表模型
MHA	32	32	100%	最高	GPT-3, BERT
MQA	32	1	3.1%	降低	PaLM, Falcon
GQA	32	8	25%	接近MHA	Llama 2 70B, Mistral

6. 注意力机制的未来

6.1 线性 RNN 的回归

Mamba、RWKV 等线性 RNN 架构在特定场景下展现了替代注意力的潜力：

# Mamba 的选择性状态空间模型核心思想
class SelectiveSSM(nn.Module):
    """简化的 Mamba 选择性机制"""
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_state = d_state
        # 输入相关的参数化 (选择性)
        self.proj_B = nn.Linear(d_model, d_state)
        self.proj_C = nn.Linear(d_model, d_state)
        self.proj_delta = nn.Linear(d_model, 1)
        self.A = nn.Parameter(torch.randn(d_state))
        self.D = nn.Parameter(torch.randn(1))
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        关键创新: B, C, delta 依赖于输入 → 选择性
        """
        B_seq = self.proj_B(x)  # (batch, seq, d_state) - 输入相关
        C_seq = self.proj_C(x)
        delta = F.softplus(self.proj_delta(x))  # (batch, seq, 1)
        
        # 离散化 A: A_bar = exp(delta * A)
        A_bar = torch.exp(delta * self.A)  # (batch, seq, d_state)
        
        # 前向扫描 (实际实现用并行扫描)
        h = torch.zeros(x.size(0), self.d_state, device=x.device)
        outputs = []
        for t in range(x.size(1)):
            h = A_bar[:, t] * h + B_seq[:, t] * x[:, t:t+1]
            y = torch.einsum('bd,bd->b', C_seq[:, t], h)
            outputs.append(y + self.D * x[:, t])
        
        return torch.stack(outputs, dim=1).unsqueeze(-1)

6.2 混合架构

当前趋势是混合注意力与线性模型：

Jamba：每 8 层中 7 层用 Mamba + 1 层 Attention
Zamba：Mamba + 共享 Attention 层交替
Grok-1.5：结合 Sparse Attention 与 Mamba

7. 总结

注意力机制经历了三个大阶段：

2014-2016：萌芽期 — Bahdanau/Luong Attention 解决 seq2seq 瓶颈
2017-2021：Transformer 统治期 — Self-Attention + Multi-Head 成为标准
2022-至今：效率优化期 — Flash Attention 解决 IO 瓶颈，线性注意力/RNN 探索新范式

Flash Attention 证明了"不改数学、只改实现"就能获得数量级加速，这是系统优化对算法研究的深刻启示。未来，混合架构可能成为长上下文建模的主流方案。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

1. 注意力机制的起源：Bahdanau Attention (2014)#

1.1 Luong Attention (2015)#

2. Self-Attention 与 Transformer (2017)#

2.1 Scaled Dot-Product Attention#

2.2 Multi-Head Attention#

2.3 注意力复杂度分析#

3. 高效注意力变体#

3.1 Sparse Attention#

3.2 Linear Attention#

4. Flash Attention：IO 感知的注意力#

4.1 核心洞察#

4.2 分块算法#

4.3 Flash Attention 2 与 3#

5. 注意力机制的变体设计#

5.1 注意力变体分类#

5.2 滑动窗口注意力#

5.3 GQA (Grouped-Query Attention)#

6. 注意力机制的未来#

6.1 线性 RNN 的回归#

6.2 混合架构#

7. 总结#

Flash Attention 证明了"不改数学、只改实现"就能获得数量级加速，这是系统优化对算法研究的深刻启示。未来，混合架构可能成为长上下文建模的主流方案。#

加入讨论#

📌 相关推荐