1. 从论文到代码Self-Attention 到底在算什么《Attention Is All You Need》这篇论文我翻过很多遍每次重读都会发现之前漏掉的细节。它提出的 Transformer 架构现在几乎成了大模型的地基但论文本身写得相当紧凑很多工程上的坑只有真正动手复现才会暴露出来。这篇笔记不打算逐段翻译论文而是聚焦一件事把 Self-Attention 和 Multi-Head Attention 这两个核心模块的参数配置拆开让你能在本地环境里跑起来、打印张量维度、验证注意力权重是否符合预期。先说清楚这套东西是什么、能做什么、适合谁。Self-Attention 是一种让序列中每个位置都能直接看到其他所有位置的机制它把输入序列映射成 Query、Key、Value 三组向量通过点积计算位置之间的相关性权重再对 Value 加权求和。Multi-Head Attention 则是把这件事并行做 h 次每次用不同的线性投影让模型同时关注不同子空间的信息。它适合正在复现 Transformer、调试注意力机制、或者想搞清楚大模型内部张量流转的开发者。如果你只是调 API 用现成模型这篇可能偏底层但如果你想自己搭一个小 Transformer 或者排查注意力相关的 bug下面的配置表可以直接抄。论文里给出的核心公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这个公式看着简单但 d_k 取多少、h 取多少、d_model 怎么分配、mask 怎么加每一个都影响最终结果。论文的设定是 d_model512、h8、d_kd_v64编码器和解码器各堆叠 N6 层前馈网络中间层 d_ff2048。这些数字不是随便定的d_model 必须能被 h 整除否则每个头的维度就对不上。我见过不少人直接把 h 改成 12 却忘了改 d_model结果 reshape 的时候维度报错排查半天。位置编码这块也容易被忽略。因为 Transformer 没有循环也没有卷积序列的顺序信息完全靠位置编码注入。论文用的是不同频率的 sin 和 cos 函数偶数维用 sin、奇数维用 cos波长从 2π 到 10000·2π 构成等比级数。这个设计的好处是对于任意固定偏移 kPE(posk) 可以表示成 PE(pos) 的线性函数模型更容易学到相对位置关系。工程上你要验证的就是位置编码矩阵的 shape 是不是 (max_len, d_model)第 0 行是不是全 0 附近的值相邻行的差异是否符合预期。下面我会先讲清楚复现这套机制需要准备什么环境然后给出可直接复制的配置片段接着用一段完整代码验证张量维度和注意力权重最后把常见的报错逐个拆解。整个过程你可以在本地 Python 环境里跟着做不需要 GPU 也能跑通小规模验证。2. 复现前的环境与 TaoToken 接入准备在动手写注意力代码之前先把运行环境和模型调用通道理清楚。纯 PyTorch 复现 Self-Attention 其实不需要任何外部服务装好 torch 就行。但如果你想把复现出来的模块和真实大模型的注意力行为做对照或者想用 API 快速验证不同参数下的输出差异那就需要一个稳定的模型接入通道。我平时用 TaoToken 来做这类对照实验它的接口兼容 OpenAI 格式切换模型只要改一个 model 字段比较适合边写代码边验证。TaoToken 的官网入口是 https://taotoken.net/?utm_sourcetaotoken_aicg_blog_endutm_mediumcsdnutm_campaignrewriteutm_content API 基础地址是 https://taotoken.net/api 。注意 API 地址后面不加 UTM 参数直接拼 /v1/chat/completions 就是标准的对话补全端点。如果你要长期做编码类实验可以看看 Coding Plan 页面如果只是想快速对话验证模型对话入口更直接需要管理密钥就去 API Keys 页面。这里要强调一点TaoToken 是合规的模型接入服务不是那种来路不明的中转。你在配置的时候把 Base URL 填成 https://taotoken.net/api Key 填自己生成的Model ID 按文档里支持的模型名填三件套齐了就能跑。我试过在 Cline、Continue 这类插件里配也是同样的三件套逻辑Base URL、API Key、Model ID 一个都不能少。环境准备的具体步骤第一步确认 Python 版本。建议 3.9 以上太老的版本有些类型注解会报错。用python --version看一眼。第二步装 PyTorch。如果你只是做 CPU 上的维度验证装 CPU 版就够pip install torch --index-url https://download.pytorch.org/whl/cpu第三步装 numpy 和 matplotlib后面可视化注意力权重会用到pip install numpy matplotlib第四步如果你要用 API 做对照装 openai 的 SDKpip install openai装完之后建一个工作目录比如transformer_notes后面所有代码文件都放这里。我习惯把配置单独放一个config.py把模型定义放attention.py验证脚本放verify.py这样排查问题的时候定位快。关于 API Key 的获取去 TaoToken 的 API Keys 页面生成一个复制出来存到环境变量里别硬编码在代码里。Linux/macOS 下export TAOTOKEN_API_KEY你的keyWindows PowerShell$env:TAOTOKEN_API_KEY你的key这样后面代码里用os.environ.get(TAOTOKEN_API_KEY)读取就行。把 Key 和 Base URL 分开管理换环境的时候只改环境变量代码不用动。还有一点如果你打算在 Claude Code 或者类似的编码工具里接入配置方式略有不同。Claude Code 走的是 Anthropic 兼容格式Base URL 要填对应的端点Key 和 Model ID 同样不能少。具体路径参考接入文档别自己猜。文档里把每个端点的用途写得很清楚照着填就行。环境搭好之后先跑一个最小验证python -c import torch; print(torch.__version__)能打印出版本号就说明基础环境没问题。接下来就可以进入注意力模块的配置环节了。3. 可复制的 Self-Attention 与 Multi-Head Attention 配置这一节是重点我把论文里的参数和工程实现对照着列出来你可以直接复制到自己的项目里。先看核心配置表参数论文取值说明工程注意点d_model512模型隐藏维度必须能被 h 整除h8注意力头数改这个要同步改 d_kd_k64每个头的 Query/Key 维度d_model / hd_v64每个头的 Value 维度通常等于 d_kd_ff2048前馈网络中间层一般是 d_model 的 4 倍N6编码器/解码器层数堆叠层数dropout0.1丢弃率训练时生效max_len5000位置编码最大长度按任务调整把这些参数写成一个 JSON 配置方便不同实验之间切换{ d_model: 512, num_heads: 8, d_k: 64, d_v: 64, d_ff: 2048, num_layers: 6, dropout: 0.1, max_len: 5000, vocab_size: 32000 }如果你用 TOML 管理配置比如配合某些训练框架等价写法是[model] d_model 512 num_heads 8 d_k 64 d_v 64 d_ff 2048 num_layers 6 dropout 0.1 max_len 5000 vocab_size 32000接下来是 Self-Attention 的实现。核心就是把输入线性投影成 Q、K、V然后按公式算。注意缩放因子是 sqrt(d_k)不是 sqrt(d_model)这是论文里明确写的很多人第一次实现会搞错。import torch import torch.nn as nn import math class ScaledDotProductAttention(nn.Module): def __init__(self, d_k, dropout0.1): super().__init__() self.d_k d_k self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): # q: (batch, heads, seq_len, d_k) # k: (batch, heads, seq_len, d_k) # v: (batch, heads, seq_len, d_v) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) attn self.dropout(attn) output torch.matmul(attn, v) return output, attn这段代码里 mask 的处理很关键。解码器的自注意力需要屏蔽未来位置做法是把 softmax 输入里对应位置设成负无穷这样 softmax 之后权重就是 0。mask 的 shape 要能广播到 scores 的 shape通常是 (batch, 1, seq_len, seq_len) 或者 (batch, 1, 1, seq_len)。然后是 Multi-Head Attention。思路是先做一次大的线性投影把 d_model 拆成 h 份每份 d_k 维并行算注意力最后拼接再投影回 d_model。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.d_v d_model // num_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) self.attention ScaledDotProductAttention(self.d_k, dropout) def forward(self, q, k, v, maskNone): batch_size q.size(0) # 线性投影并拆头 q self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v self.w_v(v).view(batch_size, -1, self.num_heads, self.d_v).transpose(1, 2) # 并行注意力 out, attn self.attention(q, k, v, mask) # 拼接多头 out out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) out self.w_o(out) return out, attn这里view和transpose的顺序不能乱。先 view 成 (batch, seq_len, heads, d_k)再 transpose 成 (batch, heads, seq_len, d_k)这样每个头的数据是连续的。如果顺序反了注意力计算就会串头结果完全错乱。contiguous()在 transpose 之后调用是为了让内存连续否则 view 会报错。位置编码的实现也要贴出来因为它是验证环节的重要对照class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(dropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return self.dropout(x)register_buffer的作用是把 pe 注册成不参与梯度更新的缓冲区这样保存和加载模型的时候它会跟着走但不会被优化器更新。div_term那个指数计算对应论文里的 10000^(2i/d_model)用 exp 和 log 组合是为了数值稳定。把这三个模块拼起来一个最小的 Transformer 编码层就有了。你可以先不接前馈网络单独验证注意力部分。配置片段和代码都齐了下一节我们跑起来看结果。4. 验证请求打印张量维度与注意力权重代码写完不验证等于没写。这一节我带你跑一遍完整的验证流程确认张量维度对得上、注意力权重符合预期、位置编码的数值范围正常。先写一个验证脚本构造随机输入走一遍 Multi-Head Attentionimport torch from attention import MultiHeadAttention, PositionalEncoding torch.manual_seed(42) batch_size 2 seq_len 10 d_model 512 num_heads 8 mha MultiHeadAttention(d_model, num_heads) x torch.randn(batch_size, seq_len, d_model) out, attn mha(x, x, x) print(输入 shape:, x.shape) print(输出 shape:, out.shape) print(注意力权重 shape:, attn.shape) print(注意力权重每行和:, attn.sum(dim-1)[0, 0])预期输出输入 shape: torch.Size([2, 10, 512]) 输出 shape: torch.Size([2, 10, 512]) 注意力权重 shape: torch.Size([2, 8, 10, 10]) 注意力权重每行和: tensor([1.0000, 1.0000, ...])注意力权重的 shape 是 (batch, heads, seq_len, seq_len)最后两维是 query 位置和 key 位置。每一行经过 softmax 之后和应该为 1这是验证 softmax 有没有算对的最直接方法。如果和不是 1检查 dim 参数是不是设成了 -1。接着验证位置编码pe PositionalEncoding(d_model, max_len100) pos_enc pe.pe[0, :5, :8] print(位置编码前5行前8维:) print(pos_enc) print(位置编码 shape:, pe.pe.shape)预期看到第 0 行接近全 0sin(0)0cos(0)1但偶数维是 sin 所以是 0奇数维是 cos 所以是 1。第 1 行的值会随维度变化波长越长的维度变化越慢。你可以画个热力图横轴是维度、纵轴是位置会看到明显的条纹图案这就是不同频率的正弦波叠加的效果。再验证一下 mask 是否生效。构造一个下三角 mask模拟解码器的自回归屏蔽seq_len 5 mask torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0) print(mask shape:, mask.shape) x torch.randn(1, seq_len, d_model) out, attn mha(x, x, x, maskmask) print(带 mask 的注意力权重:) print(attn[0, 0])预期看到注意力矩阵是下三角的上三角全是 0。因为未来位置被设成了负无穷softmax 之后权重为 0。如果上三角还有值说明 mask 的 shape 或者 masked_fill 的条件写反了。如果你想用 API 做对照验证可以发一个请求让模型解释注意力机制确认接入通道正常import os from openai import OpenAI client OpenAI( api_keyos.environ.get(TAOTOKEN_API_KEY), base_urlhttps://taotoken.net/api ) resp client.chat.completions.create( modelgpt-4o-mini, messages[{role: user, content: 用一句话解释 scaled dot-product attention 里为什么要除以 sqrt(d_k)}] ) print(resp.choices[0].message.content)如果返回正常文本说明 Base URL、Key、Model ID 三件套配置正确。如果报 401检查 Key 有没有过期或者复制的时候带了空格。如果报 model not found检查 Model ID 拼写。验证环节的核心就是三看看 shape、看权重和、看 mask 效果。这三样都对了说明你的注意力模块实现基本正确。接下来把常见的报错整理一下方便你遇到问题时快速定位。5. 常见报错排查从 401 到维度不匹配复现过程中踩的坑我基本都踩过一遍这里按报错类型整理你对着查就行。报错一401 Unauthorized / invalid api key这个通常出现在 API 调用环节。原因无非几个Key 没设置到环境变量、Key 复制时带了换行或空格、Key 已过期。排查步骤先echo $TAOTOKEN_API_KEY确认环境变量有值再检查代码里读取的变量名和设置的是否一致。如果用的是配置文件确认没有把 Key 写错行。TaoToken 的 Key 在 API Keys 页面可以重新生成生成后旧的立即失效注意更新。报错二local proxy failed / connection refused这个报错说明请求根本没发出去卡在本地网络层。检查 Base URL 是不是写成了https://taotoken.net/api注意结尾没有多余的斜杠。如果你在代码里手动拼了/v1/chat/completions确认拼接后是https://taotoken.net/api/v1/chat/completions。另外检查有没有设置HTTP_PROXY之类的环境变量干扰有的话临时 unset 掉再试。报错三reading choices of undefined这个报错说明响应体结构和你预期的不一样通常是请求本身失败了但代码直接去读resp.choices。加一层判断if resp and hasattr(resp, choices) and len(resp.choices) 0: print(resp.choices[0].message.content) else: print(响应异常:, resp)这样能看到真实的错误信息而不是被 None 掩盖。报错四RuntimeError: shape [2, 10, 8, 64] is invalid for input of size ...这是维度不匹配的典型报错几乎都出在 Multi-Head Attention 的 reshape 环节。检查 d_model 是否等于 num_heads * d_k。如果你改了 num_heads 但没改 d_model或者反过来就会对不上。另外检查 view 之前的张量元素总数是否等于 batch * seq_len * d_model。transpose 之后记得加contiguous()否则 view 会报内存不连续的错。报错五softmax 输出全为 nan这个通常是 mask 处理有问题。如果你用float(-inf)填充而某一行全是负无穷比如 mask 把整行都屏蔽了softmax 就会产生 nan。检查 mask 是否至少保留了一个有效位置。解码器自注意力的下三角 mask 第一行只有第一个位置有效这是正常的不会全屏蔽。报错六OAuth / authentication failed编码工具场景如果你在 Claude Code 或类似工具里接入报 OAuth 相关错误说明认证方式选错了。这类工具走的是 Anthropic 兼容格式不是 OpenAI 格式。Base URL、Key、Model ID 三件套要按接入文档里的说明填别混用。文档里对每个工具的配置路径都有截图照着填就行。报错七位置编码数值异常如果位置编码出现很大的值或者 nan检查div_term的计算。torch.arange(0, d_model, 2)生成的是偶数索引长度是 d_model/2。如果 d_model 是奇数这里会少一个导致pe[:, 0::2]和pe[:, 1::2]长度不一致。所以 d_model 建议用偶数论文里 512 就是偶数。排查的时候有个通用技巧在报错行前面打印相关张量的 shape 和 dtype大部分问题看一眼 shape 就能定位。维度问题永远是深度学习调试的第一嫌疑人。6. 把注意力机制用起来下一步怎么走代码跑通、维度验证通过之后你可以做几件事把理解再推进一步。第一件把单层注意力扩展成完整的编码器层。加上前馈网络和残差连接残差连接的形式是LayerNorm(x Sublayer(x))注意 LayerNorm 在残差之后。前馈网络是两层线性变换加 ReLU中间维度 d_ff2048。堆叠 6 层就是一个完整的编码器。第二件可视化注意力权重。把 attn 矩阵用 matplotlib 画成热力图你会看到不同的头关注的位置模式不一样。有的头关注相邻位置有的头关注句法结构相关的远距离位置。论文附录里展示的就是这种可视化自己画一遍印象会深很多。第三件用真实文本跑一遍。把一句话 tokenize 成 id过 embedding 和位置编码再进注意力层观察不同 token 之间的注意力权重。你会发现一些有意思的现象比如动词和它的宾语之间权重往往较高。如果你想把复现的模块和真实大模型做对照可以用 TaoToken 的模型对话入口发一些需要理解长距离依赖的句子观察模型的回答再回头看你自己的注意力权重两边对照着理解。长期做这类实验的话Coding Plan 的额度更划算适合反复调试。接入文档里有各个端点的详细说明遇到配置问题先翻文档大部分坑里面都写了。API Keys 页面管理你的密钥建议定期轮换。模型对话适合快速验证Coding Plan 适合长期编码实验按需选择。最后留一个实用技巧调试注意力的时候把 batch_size 和 seq_len 都设成很小的值比如 2 和 5这样打印出来的矩阵你能一眼看完。等逻辑对了再放大规模。维度验证永远从小规模开始这是省时间的关键。