手撕 Transformer：从原理到代码，一步步造一个“小型大模型”
作者：Echo_Wish
很多人第一次接触大模型的时候，都会有一种错觉：
Transformer 好像很复杂。
论文几十页、代码几千行、框架一堆组件。
但如果你真的把它拆开来看，其实会发现：
Transformer 的核心只有三样东西：
- Attention（注意力）
- Feed Forward（前馈网络）
- Position Encoding（位置编码）
说白了就是一句话：
Transformer = Attention + MLP + 一点点工程技巧
今天我们就用一种 “手撕源码”的方式，从零写一个最简版 Transformer。
不用任何深度学习框架黑盒。
只用 PyTorch + 基本数学。
你会发现：
大模型的底层其实没有那么神秘。
一、Transformer 到底解决了什么问题？
在 Transformer 出现之前，NLP 主要用的是：
- RNN
- LSTM
- GRU
这些模型有个致命问题：
不能并行计算。
因为它们是这样的：
word1 → word2 → word3 → word4
必须一步一步算。
而 Transformer 的思路很简单：
让每个词都能直接看到所有词。
这就是 Attention 的核心思想。
Transformer 的整体结构
从结构上看，Transformer 主要由两部分组成：
Encoder
Decoder
每一层内部包含：
Multi Head Attention
Feed Forward
Add & Norm
整个模型就是不断堆叠这些模块。
二、Attention 的核心思想
Attention 可以用一句非常接地气的话解释：
每个词在阅读句子的时候，会重点关注几个词。
例如：
The animal didn't cross the street because it was tired
这里的 it 指代谁？
模型需要去关注：
animal
Attention 就是一个计算 “谁和谁更相关” 的机制。
三、Self-Attention 的数学原理
Self-Attention 有三个关键向量：
Q Query
K Key
V Value
计算公式是：
[
Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V
]
简单理解就是三步：
1 计算相关性
2 softmax 归一化
3 加权求和
四、手写 Self-Attention（核心代码）
我们先写最核心的一段代码。
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, embed_size):
super().__init__()
self.embed_size = embed_size
self.Wq = nn.Linear(embed_size, embed_size)
self.Wk = nn.Linear(embed_size, embed_size)
self.Wv = nn.Linear(embed_size, embed_size)
def forward(self, x):
Q = self.Wq(x)
K = self.Wk(x)
V = self.Wv(x)
scores = torch.matmul(Q, K.transpose(-2, -1))
scores = scores / (self.embed_size ** 0.5)
attention = F.softmax(scores, dim=-1)
out = torch.matmul(attention, V)
return out
这段代码就是：
最简版 Self-Attention。
逻辑非常清晰：
输入 embedding
↓
生成 Q K V
↓
计算注意力
↓
加权输出
五、为什么需要 Multi-Head Attention
一个 Attention 其实能力有限。
Transformer 的解决方案是：
多头注意力。
意思是：
同时用多个注意力去理解句子。
例如：
Head1 学语法
Head2 学语义
Head3 学上下文关系
Multi Head Attention 代码实现
class MultiHeadAttention(nn.Module):
def __init__(self, embed_size, heads):
super().__init__()
self.heads = heads
self.embed_size = embed_size
self.head_dim = embed_size // heads
self.q_linear = nn.Linear(embed_size, embed_size)
self.k_linear = nn.Linear(embed_size, embed_size)
self.v_linear = nn.Linear(embed_size, embed_size)
self.fc_out = nn.Linear(embed_size, embed_size)
def forward(self, x):
batch_size = x.shape[0]
Q = self.q_linear(x)
K = self.k_linear(x)
V = self.v_linear(x)
Q = Q.reshape(batch_size, -1, self.heads, self.head_dim)
K = K.reshape(batch_size, -1, self.heads, self.head_dim)
V = V.reshape(batch_size, -1, self.heads, self.head_dim)
scores = torch.einsum("nqhd,nkhd->nhqk", Q, K)
attention = torch.softmax(scores / (self.head_dim ** 0.5), dim=-1)
out = torch.einsum("nhql,nlhd->nqhd", attention, V)
out = out.reshape(batch_size, -1, self.embed_size)
out = self.fc_out(out)
return out
这就是 Transformer 的核心模块。
六、Feed Forward 网络
每个 Transformer Block 还有一个结构：
Feed Forward Network
其实就是一个两层 MLP。
结构非常简单：
Linear
ReLU
Linear
代码如下：
class FeedForward(nn.Module):
def __init__(self, embed_size, hidden_dim):
super().__init__()
self.fc1 = nn.Linear(embed_size, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, embed_size)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
七、Position Encoding（解决顺序问题）
Transformer 最大的问题是：
它不知道词的顺序。
所以论文设计了：
位置编码。
公式：
[
PE(pos,2i)=sin(pos/10000^{2i/d})
]
[
PE(pos,2i+1)=cos(pos/10000^{2i/d})
]
实现代码：
import math
def positional_encoding(seq_len, embed_size):
pe = torch.zeros(seq_len, embed_size)
for pos in range(seq_len):
for i in range(0, embed_size, 2):
pe[pos, i] = math.sin(pos / (10000 ** (i / embed_size)))
pe[pos, i + 1] = math.cos(pos / (10000 ** (i / embed_size)))
return pe
八、组合成 Transformer Block
现在我们把组件组合起来：
Attention
+
FeedForward
+
LayerNorm
代码如下：
class TransformerBlock(nn.Module):
def __init__(self, embed_size, heads):
super().__init__()
self.attention = MultiHeadAttention(embed_size, heads)
self.norm1 = nn.LayerNorm(embed_size)
self.ff = FeedForward(embed_size, embed_size*4)
self.norm2 = nn.LayerNorm(embed_size)
def forward(self, x):
attn = self.attention(x)
x = self.norm1(x + attn)
forward = self.ff(x)
out = self.norm2(x + forward)
return out
这就是 Transformer 的一层。
真实模型就是：
Stack 12 layers
Stack 24 layers
Stack 96 layers
九、完整 Transformer（简化版）
class Transformer(nn.Module):
def __init__(self, vocab_size, embed_size, heads, num_layers):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_size)
self.layers = nn.ModuleList(
[TransformerBlock(embed_size, heads) for _ in range(num_layers)]
)
self.fc = nn.Linear(embed_size, vocab_size)
def forward(self, x):
x = self.embedding(x)
for layer in self.layers:
x = layer(x)
out = self.fc(x)
return out
一个简化版 Transformer 就完成了。
十、我第一次读 Transformer 的感受
说实话，当年第一次看论文：
完全看不懂。
公式一堆、图一堆。
但后来自己手写了一遍代码之后，突然就明白了：
Transformer 本质其实很简单：
Embedding
↓
Attention
↓
MLP
↓
重复很多层
真正复杂的不是模型结构。
而是：
- 训练数据
- 参数规模
- 工程优化
这也是为什么现在的大模型动不动就是：
70B 参数
100B 参数
1T 参数
结构没变。
只是规模变了。
十一、最后聊点真实感受
很多人把 Transformer 看成一个非常神秘的 AI 技术。
但在我看来，它更像一个：
工程奇迹。
2017 年那篇论文其实只有一句核心思想：
Attention Is All You Need
但这句话却改变了整个 AI 行业。
从：
- GPT
- BERT
- Stable Diffusion
- ChatGPT
几乎所有大模型，
背后都是 Transformer。
有时候技术革命真的很神奇：
改变世界的，可能只是一条公式。
