揭秘大模型的魔法：实现带有可训练权重的多头自注意力机制

原创已于 2025-08-11 13:44:06 修改 · 602 阅读

19 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #机器学习 #深度学习

于 2025-08-11 13:42:11 首次发布

今天我想把这块魔术板拆开来给你看个究竟：如何把单头注意力改成多头注意力，让每个头能学会自己的注意力分布。

大家好，我是写代码的中年人！

自注意力(Self-Attention)是大模型里最常让人“眼花”的魔术道具：看起来只是一堆矩阵乘法和 softmax，可是组合起来就能学到“句子里谁重要、谁次要”的规则，甚至能学到某些头只盯标点、某些头专盯主谓关系。

今天我想把这块魔术板拆开来给你看个究竟：如何把单头注意力改成多头注意力，让每个头能学会自己的注意力分布。

01、回顾单头自注意力机制

假设你在开会，桌上有一堆文件，你想找跟“项目进度”相关的内容。

你心里有个问题（Query）：“项目进度在哪儿？

”每份文件上有个标签（Key），写着它的主题，比如“预算”“进度”“人员”。

你会先挑出标签里跟“进度”相关的文件（匹配），然后重点看这些文件的内容（Value），最后把这些内容总结成你的理解。

自注意力就像是给每个词都做了一次这样的“信息筛选和总结”，让每个词都能根据上下文更好地表达自己。

02、理解多头自注意力机制

继续用开会的场景：

桌上还是那堆文件（代表句子里的词），但现在你不是一个人干活，而是找了3个助手（假设3头注意力）。每个助手都有自己的“专长”，他们会从不同的角度问问题、匹配标签和提取内容。

每个头独立工作（多视角筛选）：

头1（进度专家）：他的问题（Query）是“进度怎么样？”他只关注标签里跟“进度”“时间表”相关的文件，忽略其他。挑出匹配的文件后，他总结出一份“进度报告”。

头2（预算专家）：他的问题是“预算超支了吗？”他匹配标签里的“预算”“开销”，然后从那些文件的内容里提炼“预算分析”。

头3（风险专家）：问题是“有什么隐患？”他找“风险”“问题”相关的标签，输出一份“风险评估”。

每个头都像单头注意力一样：生成自己的问题、钥匙和内容，计算匹配度，加权总结。但他们用的“眼镜”不同（在机器里，这通过不同的线性变换实现），所以捕捉的信息侧重点不一样。

把多头结果合起来（综合决策）：

一旦每个头都给出自己的总结，你就把这些报告拼在一起（或简单平均一下），形成一份完整的“项目概览”。现在，你的理解不只是“进度”，而是进度+预算+风险的全方位视图。万一某个头漏了什么，其他头能补上，确保没死角。

03、用代码实现多头自注意力机制

# ONE

我们使用水浒传的内容进行演示，使用前三回各 100 字的文本，并按“字”切分成模型可用的格式。

复制

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np


plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# ======  准备水浒传真实语料 ======
raw_texts = [
    "話說大宋仁宗天子在位，嘉祐三年三月三日五更三點，天子駕坐紫宸殿，受百官朝賀。但見：祥雲迷鳳閣，瑞氣罩龍樓。含煙御柳拂旌旗，帶露宮花迎劍戟。天香影裏，玉簪珠履聚丹墀。仙樂聲中，繡襖錦衣扶御駕。珍珠廉卷，黃金殿上現金輿。鳳尾扇開，白玉階前停寶輦。隱隱凈鞭三下響，層層文武兩班齊。",
    "那高俅在臨淮州，因得了赦宥罪犯，思鄉要回東京。這柳世權卻和東京城里金梁橋下開生藥鋪的董將士是親戚，寫了一封書札，收拾些人事盤纏，赍發高俅回東京，投奔董將士家過活。",
    "話說當時史進道：「卻怎生是好？」朱武等三個頭領跪下答道：「哥哥，你是乾淨的人，休為我等連累了大郎。可把索來綁縛我三個，出去請賞，免得負累了你不好看。」"
]


# ======  按字切分 ======
def char_tokenize(text):
    return [ch for ch in text if ch.strip()]  # 去掉空格、换行


sentences = [char_tokenize(t) for t in raw_texts]


# 构建词表
vocab = {}
for sent in sentences:
    for ch in sent:
        if ch not in vocab:
            vocab[ch] = len(vocab)


# ======  转成索引形式并做 padding ======
max_len = max(len(s) for s in sentences)
PAD_TOKEN = "<PAD>"
vocab[PAD_TOKEN] = len(vocab)


input_ids = []
for sent in sentences:
    ids = [vocab[ch] for ch in sent]
    # padding
    ids += [vocab[PAD_TOKEN]] * (max_len - len(ids))
    input_ids.append(ids)


input_ids = torch.tensor(input_ids)  # (batch_size, seq_len)


# ======  多头自注意力模块 ======
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim 必须能整除 num_heads"
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads


        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)


        self.dropout = dropout
        self.last_attn_weights = None  # 保存最后一次注意力权重 (batch, heads, seq, seq)


    def forward(self, x):
        B, T, C = x.size()


        Q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)


        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)


        self.last_attn_weights = attn_weights.detach()  # (B, heads, T, T)


        out = torch.matmul(attn_weights, V)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        out = self.out_proj(out)
        return out


# ======  模型训练 ======
embed_dim = 32
num_heads = 4
vocab_size = len(vocab)


embedding = nn.Embedding(vocab_size, embed_dim)
model = MultiHeadSelfAttention(embed_dim, num_heads)
criterion = nn.MSELoss()
optimizer = optim.Adam(list(model.parameters()) + list(embedding.parameters()), lr=1e-3)


epochs = 200
for epoch in range(epochs):
    model.train()
    x = embedding(input_ids)
    target = x.clone()


    out = model(x)
    loss = criterion(out, target)


    optimizer.zero_grad()
    loss.backward()
    optimizer.step()


    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1:3d}, Loss: {loss.item():.6f}")


# ======  可视化注意力热图 ======
for idx, sent in enumerate(sentences):
    attn = model.last_attn_weights[idx]  # (heads, seq, seq)
    sent_len = len(sent)
    for head in range(num_heads):
        plt.figure(figsize=(8, 6))
        plt.imshow(attn[head, :sent_len, :sent_len].numpy(), cmap='viridis')
        plt.title(f"第{idx+1}句 第{head+1}头 注意力矩阵")
        plt.xticks(ticks=np.arange(sent_len), labels=sent, rotatinotallow=90)
        plt.yticks(ticks=np.arange(sent_len), labels=sent)
        plt.xlabel("Key (字)")
        plt.ylabel("Query (字)")
        plt.colorbar(label="Attention Strength")
        for i in range(sent_len):
            for j in range(sent_len):
                plt.text(j, i, f"{attn[head, i, j]:.2f}", ha="center", va="center", color="white", fnotallow=6)
        plt.tight_layout()
        plt.savefig(f"attention_sentence{idx+1}_head{head+1}.png")
        plt.close()


print("注意力热图已保存。")