手撕 Transformer：从原理到代码，一步步造一个“小型大模型”

作者：Echo_Wish

很多人第一次接触大模型的时候，都会有一种错觉：

Transformer 好像很复杂。

论文几十页、代码几千行、框架一堆组件。

但如果你真的把它拆开来看，其实会发现：

Transformer 的核心只有三样东西：

1. Attention（注意力）
2. Feed Forward（前馈网络）
3. 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。

有时候技术革命真的很神奇：

改变世界的，可能只是一条公式。