2026年5月12日约 4 分钟阅读

Multi-Head Attention: An Intuitive Guide

从直觉到实现，理解为什么 Transformer 需要多个 Attention Head，以及它们各自在学习什么。

Multi-Head Attention 是 Transformer 的核心组件。但为什么要用多个 Head？每个 Head 在学什么？

Single-Head Attention 的局限

单个 Attention Head 只能学习一种关系模式：

"The cat sat on the mat"

这个句子有多种关系：

语法关系：cat 是 sat 的主语
语义关系：cat 和 mat 都是动物/物体
位置关系：sat 在 cat 之后

一个 Head 很难同时捕捉所有这些关系。

Multi-Head Attention 的直觉

想象你在阅读一个句子，你会同时关注：

语法 Head：谁对谁做了什么？
语义 Head：哪些词意思相近？
位置 Head：哪些词靠得近？
指代 Head：代词指代谁？

The Formula

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$

其中每个 Head：

$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

def multi_head_attention(Q, K, V, n_heads=8):
    d_k = Q.shape[-1] // n_heads
    
    # 分割成多个 Head
    heads = []
    for i in range(n_heads):
        q_i = Q[:, :, i*d_k:(i+1)*d_k]
        k_i = K[:, :, i*d_k:(i+1)*d_k]
        v_i = V[:, :, i*d_k:(i+1)*d_k]
        
        # 每个 Head 独立计算 Attention
        head_i = scaled_dot_product_attention(q_i, k_i, v_i)
        heads.append(head_i)
    
    # 拼接所有 Head
    return torch.cat(heads, dim=-1)

What Do Attention Heads Learn?

研究发现，不同层的 Attention Head 确实学到了不同的模式：

Bottom Layers (Near Input)

位置模式：关注相邻的词
形态学：词缀、词根关系

Middle Layers

语法关系：主谓宾、修饰关系
短距离依赖：介词短语、从句

Top Layers (Near Output)

语义关系：同义词、反义词
长距离依赖：指代消解、篇章结构

Visualizing Attention

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(tokens, attention_weights, head_idx):
    """可视化某个 Head 的 Attention 权重"""
    fig, ax = plt.subplots(figsize=(10, 10))
    sns.heatmap(
        attention_weights[head_idx],
        xticklabels=tokens,
        yticklabels=tokens,
        ax=ax,
        cmap='Blues'
    )
    ax.set_title(f'Head {head_idx} Attention')
    plt.show()

Why 8 Heads?

原始 Transformer 使用 8 个 Head， $d_{\text{model}} = 512$ ，所以每个 Head 的维度是 64。

Rule of Thumb：

Head 数太少：表达能力不足
Head 数太多：计算开销大，可能过拟合

Modern Choices：

GPT-3：96 Heads
BERT-large：16 Heads
LLaMA-7B：32 Heads

Attention Head Redundancy

研究发现，很多 Attention Head 是冗余的：

Paper: “Are Sixteen Heads Really Better than One?”

Finding: 删除大部分 Head 后，性能下降很小

这说明：

不是所有 Head 都同等重要
可以进行 Head Pruning 来压缩模型
某些 Head 可能是”备份”

Complete Implementation

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        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)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # Linear Transformations
        Q = self.W_q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # Compute Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attention = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention, V)
        
        # Concatenate Heads
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        return self.W_o(output), attention

Experiment: Observing Attention Heads

# 使用 BERT 可视化 Attention
from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased', output_attentions=True)

inputs = tokenizer("The cat sat on the mat", return_tensors="pt")
outputs = model(**inputs)

# outputs.attentions: (n_layers, batch, n_heads, seq_len, seq_len)
for layer_idx, layer_attention in enumerate(outputs.attentions):
    print(f"Layer {layer_idx}:")
    for head_idx in range(layer_attention.size(1)):
        # 找到每个 Head 最关注的词对
        max_attention = layer_attention[0, head_idx].max()
        print(f"  Head {head_idx}: max attention = {max_attention:.3f}")

Summary

Multi-Head Attention 让 Transformer 能够同时学习多种关系模式。每个 Head 专注于不同的方面，共同构成丰富的表示。理解这一点，有助于理解为什么 Transformer 如此强大，以及如何设计更好的模型。

Further Reading：