)
1. 从「读得懂公式」到「跑得通代码」Transformer 最小复现的卡点在哪Attention is all you need 这篇论文很多人第一次读都能看懂「自注意力就是把 Q 和 K 做点积再 softmax」但真正动手写代码时问题就来了Q、K、V 到底从哪来多头注意力的张量形状怎么切位置编码为什么是 sin/cos 而不是直接加个序号编码器-解码器骨架里哪些地方要加残差和 LayerNorm我自己第一次复现时卡在多头注意力的 reshape 上整整一个下午。论文里写的是「把 d_model 投影到 h 个 d_k 维度」但代码里到底是先 reshape 再 transpose还是先 transpose 再 reshape顺序错了结果就完全不对。更麻烦的是这些形状错误不会报错只会让 loss 不下降你根本不知道哪里出了问题。这篇是「上篇」目标很明确不追求训练出能翻译的模型而是先把 Transformer 的编码器-解码器骨架搭起来用一次前向传播做形状校验确保每个张量的维度都和论文对得上。同时我会用 TaoToken 的统一 Key 来跑通模型对话验证环节这样你不需要在多个平台之间切换 API Key一个 Key 就能完成从代码调试到模型行为验证的闭环。适合谁看已经读过论文但没动手写过 Transformer 的人写过但形状总是对不上的人想用统一 API 通道做论文复现验证的人。下面从环境准备开始一步步把骨架跑通。2. TaoToken 前置统一 Key 在论文复现里扮演什么角色论文复现和普通调库最大的区别是你需要频繁验证「模型输出是否符合预期」。比如位置编码加进去之后模型对同一句话在不同位置的注意力权重是否合理多头注意力切分后每个头的输出是否真的在不同子空间里。这些验证如果每次都靠人眼看 tensor效率很低。更实际的做法是把中间结果整理成自然语言描述丢给一个对话模型帮你判断「这个形状和论文描述是否一致」。TaoToken 在这里的作用就是提供一个统一的 API 通道你不需要分别去申请多个平台的 Key一个 Key 就能覆盖模型对话、代码生成辅助、文档查询等场景。具体来说TaoToken 的 API 地址是https://taotoken.net/api兼容 OpenAI 风格的接口格式。你可以在本地 Python 脚本里直接调用也可以配合 Cursor、VS Code 插件等工具使用。对于论文复现这种需要反复调试的场景统一 Key 的好处是你不需要在多个配置文件里维护不同的 base_url 和 api_key改一个地方就行。如果你还没拿 Key可以去官网看一下https://taotoken.net/?utm_sourcetaotoken_aicg_blog_endutm_mediumcsdnutm_campaignrewriteutm_content 。拿到 Key 之后下面直接进入配置环节。3. 可复制配置config.toml 与 settings.json 骨架论文复现项目建议用config.toml管理模型超参数用settings.json管理 API 通道和本地环境。这样做的原因是超参数需要频繁调整比如 d_model、n_heads、n_layers而 API Key 这类敏感信息不应该硬编码在代码里。先建项目目录mkdir transformer-repro cd transformer-repro mkdir config src checkpoints touch config/config.toml config/settings.json src/model.py src/shape_check.pyconfig/config.toml的内容如下参数直接对照论文的 base 配置[model] d_model 512 n_heads 8 n_layers 6 d_ff 2048 dropout 0.1 max_seq_len 128 [training] batch_size 2 lr 0.0001 epochs 10 [data] src_vocab_size 1000 tgt_vocab_size 1000这里解释几个关键参数d_model512是论文里所有子层和嵌入层的输出维度n_heads8表示多头注意力的头数每个头的维度是d_model / n_heads 64d_ff2048是前馈网络中间层的维度是d_model的 4 倍。这些数字不是随便定的论文在 3.1 节和 3.3 节有明确说明。config/settings.json管理 API 通道{ api_base: https://taotoken.net/api, api_key: sk-your-key-here, model: gpt-4o-mini, timeout: 30 }注意api_base后面不要加/v1TaoToken 的接口路径已经内置了兼容层。如果你用的是其他工具比如 Cursor 或 Continue在设置里填https://taotoken.net/api作为 OpenAI Base URL 即可。提示settings.json建议加入.gitignore避免 Key 泄露。如果你在团队里共享代码可以提供一个settings.example.json作为模板。4. 逐段对照论文自注意力、多头注意力、位置编码的代码清单这一节是核心。我会按论文的公式顺序把每个模块的代码写出来并在注释里标注对应的论文段落和形状变化。4.1 缩放点积注意力公式 (1) 的代码实现论文公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) Vimport torch import torch.nn as nn import math class ScaledDotProductAttention(nn.Module): def __init__(self, d_k): super().__init__() self.d_k d_k def forward(self, q, k, v, maskNone): # q: (batch, n_heads, seq_len, d_k) # k: (batch, n_heads, seq_len, d_k) # v: (batch, n_heads, seq_len, d_k) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) # scores: (batch, n_heads, seq_len, seq_len) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) output torch.matmul(attn, v) # output: (batch, n_heads, seq_len, d_k) return output, attn这里的关键是k.transpose(-2, -1)把最后两个维度交换让(seq_len, d_k)变成(d_k, seq_len)这样矩阵乘法才能得到(seq_len, seq_len)的注意力分数矩阵。除以sqrt(d_k)是论文 3.2.1 节强调的缩放操作防止点积过大导致 softmax 梯度消失。4.2 多头注意力论文 3.2.2 节的投影与拼接多头注意力的核心是把 Q、K、V 分别投影到 h 个低维子空间并行做注意力再拼接回来。class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 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) self.dropout nn.Dropout(dropout) self.attention ScaledDotProductAttention(self.d_k) def forward(self, q, k, v, maskNone): batch_size q.size(0) # 线性投影并切分多头 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) # q/k/v: (batch, n_heads, seq_len, d_k) out, attn self.attention(q, k, v, mask) # out: (batch, n_heads, seq_len, d_k) out out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # out: (batch, seq_len, d_model) out self.w_o(out) return out, attn形状变化是这里最容易出错的地方。view(batch_size, -1, n_heads, d_k)把d_model拆成n_heads * d_k然后transpose(1, 2)把n_heads换到前面变成(batch, n_heads, seq_len, d_k)。注意力算完之后transpose(1, 2)换回来再用contiguous().view()合并成(batch, seq_len, d_model)。如果你发现输出形状不对大概率是transpose和view的顺序搞反了。4.3 位置编码论文 3.5 节的 sin/cos 实现位置编码的公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))class PositionalEncoding(nn.Module): def __init__(self, d_model, max_seq_len128, dropout0.1): super().__init__() self.dropout nn.Dropout(dropout) pe torch.zeros(max_seq_len, d_model) position torch.arange(0, max_seq_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) # (1, max_seq_len, d_model) self.register_buffer(pe, pe) def forward(self, x): # x: (batch, seq_len, d_model) x x self.pe[:, :x.size(1), :] return self.dropout(x)register_buffer的作用是把pe注册为模型的一部分但不参与梯度更新。这样保存模型时pe会一起保存加载时也能恢复。pe[:, 0::2]和pe[:, 1::2]分别填充偶数和奇数维度对应公式里的2i和2i1。4.4 编码器层与解码器层残差 LayerNorm 的组合论文 3.1 节说每个子层的输出是LayerNorm(x Sublayer(x))。编码器层包含两个子层多头自注意力和前馈网络。解码器层多一个交叉注意力子层。class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): attn_out, _ self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_out)) ffn_out self.ffn(x) x self.norm2(x self.dropout(ffn_out)) return x解码器层多了一个交叉注意力Q 来自解码器上一层K 和 V 来自编码器输出class DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.cross_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, enc_out, src_maskNone, tgt_maskNone): attn_out, _ self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_out)) cross_out, _ self.cross_attn(x, enc_out, enc_out, src_mask) x self.norm2(x self.dropout(cross_out)) ffn_out self.ffn(x) x self.norm3(x self.dropout(ffn_out)) return x注意cross_attn的 Q 是xK 和 V 都是enc_out。这就是论文里说的「解码器每个位置可以关注输入序列的所有位置」。5. 验证请求一次前向传播的形状校验与 API 通道验证代码写完了但怎么确认形状是对的最直接的方法就是跑一次前向传播把每个中间张量的形状打印出来和论文描述对照。import torch from model import EncoderLayer, DecoderLayer, PositionalEncoding def shape_check(): batch_size 2 seq_len 10 d_model 512 n_heads 8 d_ff 2048 x torch.randn(batch_size, seq_len, d_model) pos_enc PositionalEncoding(d_model, max_seq_len128) x pos_enc(x) print(fafter positional encoding: {x.shape}) # 期望: (2, 10, 512) encoder_layer EncoderLayer(d_model, n_heads, d_ff) enc_out encoder_layer(x) print(fafter encoder layer: {enc_out.shape}) # 期望: (2, 10, 512) decoder_layer DecoderLayer(d_model, n_heads, d_ff) dec_out decoder_layer(x, enc_out) print(fafter decoder layer: {dec_out.shape}) # 期望: (2, 10, 512) # 检查多头注意力的中间形状 mha encoder_layer.self_attn q mha.w_q(x).view(batch_size, -1, n_heads, d_model // n_heads).transpose(1, 2) print(fq after multi-head split: {q.shape}) # 期望: (2, 8, 10, 64) if __name__ __main__: shape_check()运行结果应该是after positional encoding: torch.Size([2, 10, 512]) after encoder layer: torch.Size([2, 10, 512]) after decoder layer: torch.Size([2, 10, 512]) q after multi-head split: torch.Size([2, 8, 10, 64])如果某个形状对不上比如q变成了(2, 10, 8, 64)说明transpose(1, 2)漏掉了或者顺序错了。这一步看起来简单但能帮你省下大量调试时间。接下来用 TaoToken 的 API 通道做一次模型对话验证。你可以把形状检查的结果整理成一段描述让模型帮你判断是否符合论文import json import requests with open(config/settings.json) as f: settings json.load(f) headers { Authorization: fBearer {settings[api_key]}, Content-Type: application/json } payload { model: settings[model], messages: [ {role: system, content: 你是一个 Transformer 论文复现助手擅长检查张量形状是否符合论文描述。}, {role: user, content: 我跑了一次前向传播输入是 (2, 10, 512)经过位置编码后是 (2, 10, 512)经过编码器层后是 (2, 10, 512)经过解码器层后是 (2, 10, 512)。多头注意力切分后 q 的形状是 (2, 8, 10, 64)。请判断这些形状是否符合 Attention is all you need 论文的 base 配置。} ] } resp requests.post( f{settings[api_base]}/chat/completions, headersheaders, jsonpayload, timeoutsettings[timeout] ) print(resp.json()[choices][0][message][content])如果返回的内容确认形状正确说明你的骨架和论文是对齐的。如果模型指出某个地方不对你可以回到代码里检查对应的 reshape 操作。注意TaoToken 的 API 兼容 OpenAI 格式所以requests.post的路径是/chat/completions。如果你用的是 OpenAI SDK直接把base_url设成https://taotoken.net/api就行。6. 本篇常见错排查形状对不上、位置编码加错、mask 方向反了复现 Transformer 时下面这几个错误几乎每个人都会遇到。我按出现频率从高到低排列你可以对照自己的情况排查。错误一多头注意力 reshape 顺序错误。最常见的写法是x.view(batch, seq_len, n_heads, d_k).transpose(1, 2)但有人会写成x.view(batch, n_heads, seq_len, d_k)这样直接把seq_len和n_heads搞混了。正确的做法是先view再transpose因为view是按内存顺序切分transpose才是交换维度。如果你发现注意力权重矩阵的形状是(batch, seq_len, seq_len, n_heads)而不是(batch, n_heads, seq_len, seq_len)就是这里错了。错误二位置编码的维度对不上。位置编码的pe形状是(1, max_seq_len, d_model)加到输入x上时x的形状是(batch, seq_len, d_model)。如果seq_len max_seq_len切片self.pe[:, :x.size(1), :]会返回比x短的张量广播时就会报错。解决办法是把max_seq_len设得比实际序列长度大或者在forward里动态扩展pe。错误三mask 的方向反了。解码器的自注意力需要 mask 掉未来位置论文里说「把 softmax 输入中对应非法连接的值设为 -∞」。代码里通常是scores.masked_fill(mask 0, -1e9)但 mask 的 0 和 1 含义容易搞反。如果你发现模型在训练时 loss 不下降或者生成时总是重复同一个词检查一下 mask 是不是把应该保留的位置也 mask 掉了。错误四LayerNorm 的位置放错。论文写的是LayerNorm(x Sublayer(x))也就是先残差再归一化。但有些实现会写成先归一化再残差这样训练动态会不一样。如果你发现训练初期 loss 震荡很大检查一下norm和dropout的顺序。错误五API 调用时 base_url 多了/v1。TaoToken 的接口地址是https://taotoken.net/api如果你在 OpenAI SDK 里写成https://taotoken.net/api/v1可能会 404。正确的做法是直接用https://taotoken.net/apiSDK 会自动拼接/chat/completions。如果你在排查过程中需要查论文原文的某个公式可以直接用 TaoToken 的模型对话功能把问题描述清楚让它帮你定位。模型对话入口在https://taotoken.net/models?utm_sourcetaotoken_aicg_blog_endutm_contentmodel_chatutm_campaignrewrite 。如果你更习惯在编辑器里直接调试可以看接入文档https://taotoken.net/doc?utm_sourcetaotoken_aicg_blog_endutm_contentdocutm_campaignrewrite 。7. 下一步从骨架到可训练模型以及 Coding Plan 的用法这篇「上篇」完成了三件事拆解了自注意力、多头注意力、位置编码的公式与形状给出了可复制的config.toml和settings.json骨架跑通了一次前向传播的形状校验。你现在应该有一个能跑通但不一定能训练的 Transformer 骨架。「下篇」会做这几件事补全编码器和解码器的堆叠逻辑实现完整的 mask 生成函数用一个小规模数据集跑通训练循环验证模型是否真的能学到东西。如果你等不及可以先把这篇的代码整理好确保每个模块单独测试通过。对于长期做论文复现和 Agent 开发的人TaoToken 的 Coding Plan 可能更适合你。它提供更稳定的调用配额和更低的延迟适合需要反复调试的场景。具体可以看https://taotoken.net/coding-plan?utm_sourcetaotoken_aicg_blog_endutm_contentcoding_planutm_campaignrewrite 。如果你只是想先拿个 Key 跑通这篇的验证代码直接去 API Keys 页面https://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite 。最后说一个我踩过的坑复现论文时不要一上来就追求和论文完全一致的超参数。先把d_model改成 64、n_heads改成 2、n_layers改成 1用一个小模型快速验证代码逻辑。等形状和梯度都对了再放大到论文的 base 配置。这样调试效率会高很多。