# 导入 PyTorch 核心库
import torch # 深度学习框架：提供张量运算与自动求导
import torch.nn as nn # 神经网络模块：Linear、Embedding、LayerNorm、ModuleList 等
import torch.nn.functional as F # 函数式接口：提供 softmax 等无状态操作
import torch.optim as optim # 优化器模块：提供 Adam 等参数更新算法
import math # 数学库：用于 sqrt 计算缩放因子
# ============================================================
# 0. vocab — 词表定义
# ============================================================
# 示例仅使用 4 个汉字作为最小词表，便于理解注意力机制每一步的计算
chars = ["你", "好", "世", "界"] # 词表：4 个字符
vocab_size = len(chars) # 词表大小 = 4
# 正向映射：字符 → 整数索引，供模型输入使用
char2idx = {c: i for i, c in enumerate(chars)}
# 反向映射：整数索引 → 字符，供推理阶段还原汉字显示
idx2char = {i: c for c, i in char2idx.items()}
# ============================================================
# 1. Transformer 超参数（标准 Transformer 风格）
# ============================================================
d_model = 512 # 模型隐藏维度（嵌入维度），Transformer 经典配置
n_heads = 8 # 多头注意力头数
n_layers = 6 # Transformer 块堆叠层数
d_ff = 2048 # 前馈网络中间层维度（通常为 d_model 的 4 倍）
# 确保 d_model 可被 n_heads 整除，每头维度为整数
assert d_model % n_heads == 0
d_head = d_model // n_heads # 每个注意力头的维度 = 512 / 8 = 64
# ============================================================
# 2. 训练数据 — 自回归序列构造
# ============================================================
# Transformer 自动学习 token → embedding → attention → output logits 的完整流程，
# 无需手动指定 target_vectors。
#
# 输入序列: [你, 好, 世, 界] → 索引 [0, 1, 2, 3]
# 目标序列: [好, 世, 界, 你] → 索引 [1, 2, 3, 0]
# 这就是“自回归”: 给定前 t 个字符，预测第 t+1 个字符
data = torch.tensor([[0, 1, 2, 3]]) # 输入张量，形状 (batch=1, seq_len=4)
target = torch.tensor([[1, 2, 3, 0]]) # 目标张量，形状 (batch=1, seq_len=4)
# ============================================================
# 3. Multi-Head Attention — 多头注意力层
# ============================================================
class MultiHeadAttention(nn.Module):
"""
标准缩放点积多头自注意力（Scaled Dot-Product Multi-Head Self-Attention）。
数学定义:
Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V
流程: 输入 x → 线性投影得到 Q/K/V → 拆分为多头 → 计算注意力得分 →
缩放 → 因果掩码 → softmax → 加权求和 V → 合并多头 → 输出投影
"""
def __init__(self, d_model, n_heads):
"""
参数:
d_model: 模型总维度，同时作为 Q/K/V 的投影维度
n_heads: 注意力头数
"""
super().__init__()
self.d_model = d_model # 模型维度，例如 512
self.n_heads = n_heads # 注意力头数，例如 8
self.d_head = d_model // n_heads # 每头维度 = 512 / 8 = 64
# 四个无偏置线性投影矩阵（形状均为 d_model × d_model）:
self.W_Q = nn.Linear(d_model, d_model, bias=False) # 查询投影: x → Q
self.W_K = nn.Linear(d_model, d_model, bias=False) # 键投影: x → K
self.W_V = nn.Linear(d_model, d_model, bias=False) # 值投影: x → V
self.W_O = nn.Linear(d_model, d_model, bias=False) # 输出投影: 拼接后融合
def forward(self, x, return_attention=False):
"""
前向传播。
参数:
x: 输入张量，形状 (B, T, D)
return_attention: 是否返回中间计算结果（Q/K/V/scores/alpha）供可解释性分析
返回:
if return_attention=False: out, 形状 (B, T, D)
if return_attention=True: (out, Q, K, V, scores, scores_scaled, alpha)
"""
B, T, D = x.shape # B=批次大小, T=序列长度, D=d_model
# ---------- 第 1 步：线性投影 x → Q, K, V ----------
Q = self.W_Q(x) # (B, T, D) → (B, T, D)，每个 token 的查询向量
K = self.W_K(x) # (B, T, D) → (B, T, D)，每个 token 的键向量
V = self.W_V(x) # (B, T, D) → (B, T, D)，每个 token 的值向量
# ---------- 第 2 步：拆分为多头 ----------
# view: (B, T, D) → (B, T, n_heads, d_head)
# transpose: 交换维度 1 和 2 → (B, n_heads, T, d_head)
# 此后每个头独立拥有 T × d_head 的 Q/K/V 子空间
Q = Q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
K = K.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
V = V.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
# ---------- 第 3 步：计算注意力得分 S = Q · Kᵀ ----------
# Q: (B, n_heads, T, d_head), Kᵀ: (B, n_heads, d_head, T)
# scores: (B, n_heads, T, T) — 位置 i 和位置 j 之间的原始相似度
scores = Q @ K.transpose(-2, -1)
# ---------- 第 4 步：缩放 S / √d_k ----------
# 除以 √d_k 防止点积值过大，避免 softmax 进入梯度饱和区
# 在标准配置下 d_head = 64，因此缩放因子 = 8
scores_scaled = scores / math.sqrt(self.d_head)
# ---------- 第 5 步：因果掩码（Causal Mask）----------
# 自回归语言模型要求位置 i 只能看到 ≤ i 的 token（不可预先获取未来信息）
# torch.triu(..., diagonal=1) 生成上三角矩阵（对角线以上为 True）
# 例如 T=4 时:
# [[F, T, T, T],
# [F, F, T, T],
# [F, F, F, T],
# [F, F, F, F]]
mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
# 掩码位置填充 -inf，经 softmax 后权重≈0，即“禁止关注”
scores_scaled = scores_scaled.masked_fill(mask, float("-inf"))
# ---------- 第 6 步：Softmax 归一化 ----------
# 沿最后一维（key 方向）做 softmax，得到注意力权重分布 α
# 每行权重之和 = 1
alpha = F.softmax(scores_scaled, dim=-1)
# ---------- 第 7 步：加权求和 output = α · V ----------
# α: (B, n_heads, T, T) — 注意力权重
# V: (B, n_heads, T, d_head) — 值向量
# out: (B, n_heads, T, d_head) — 加权后的上下文表示
out = alpha @ V
# ---------- 第 8 步：合并多头 ----------
# transpose: (B, n_heads, T, d_head) → (B, T, n_heads, d_head)
# contiguous + view: 展平为 (B, T, n_heads * d_head) = (B, T, D)
out = out.transpose(1, 2).contiguous()
out = out.view(B, T, D)
# ---------- 第 9 步：输出投影 ----------
out = self.W_O(out) # W_O 融合来自不同头的信息
# 根据 return_attention 标志决定返回内容
if return_attention:
return out, Q, K, V, scores, scores_scaled, alpha
return out
# ============================================================
# 4. FeedForward — 前馈网络
# ============================================================
class FeedForward(nn.Module):
"""
位置式前馈网络（Position-wise Feed-Forward Network）。
对每个位置的表示独立应用两层全连接 + ReLU 激活:
FFN(x) = ReLU(x·W₁ + b₁)·W₂ + b₂
中间维度 d_ff 通常为 d_model 的 4 倍 (512 → 2048)，扩展后压缩回原维度。
"""
def __init__(self, d_model, d_ff):
"""
参数:
d_model: 输入/输出维度
d_ff: 中间隐藏层维度（通常是 d_model 的 4 倍）
"""
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model, d_ff), # 升维：d_model → d_ff (512 → 2048)
nn.ReLU(), # 非线性激活函数，引入非线性
nn.Linear(d_ff, d_model) # 降维：d_ff → d_model (2048 → 512)
)
def forward(self, x):
return self.net(x)
# ============================================================
# 5. TransformerBlock — Transformer 块
# ============================================================
class TransformerBlock(nn.Module):
"""
一个完整的 Transformer 块，采用 Pre-Norm 残差结构：
x → LayerNorm → MultiHeadAttention → + → x'
x' → LayerNorm → FeedForward → + → x"
Pre-Norm（先归一化再子层）相比 Post-Norm 训练更稳定。
"""
def __init__(self, d_model, n_heads, d_ff):
"""
参数:
d_model: 模型隐藏维度
n_heads: 注意力头数
d_ff: 前馈网络中间层维度
"""
super().__init__()
# 第一个 Pre-Norm 子层: LayerNorm + MultiHeadAttention
self.ln1 = nn.LayerNorm(d_model) # 层归一化，沿最后一维标准化
self.attn = MultiHeadAttention(d_model, n_heads)
# 第二个 Pre-Norm 子层: LayerNorm + FeedForward
self.ln2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model, d_ff)
def forward(self, x, return_attention=False):
"""
前向传播。
参数:
x: 输入张量 (B, T, D)
return_attention: 是否传递至注意力层以捕获中间值
返回:
普通模式: 输出张量 (B, T, D)
注意力模式: (输出, Q, K, V, scores, scores_scaled, alpha)
"""
if return_attention:
# 需要捕获注意力中间值:
# 先 LayerNorm，再注意力（同时返回注意力中间值），最后残差连接
attn_out, Q, K, V, scores, scores_scaled, alpha = \
self.attn(self.ln1(x), return_attention=True)
x = x + attn_out # 残差连接 1: x + Attention(LayerNorm(x))
x = x + self.ffn(self.ln2(x)) # 残差连接 2: x + FFN(LayerNorm(x))
return x, Q, K, V, scores, scores_scaled, alpha
else:
# 普通前向: 不捕获中间值
x = x + self.attn(self.ln1(x)) # 残差连接 1
x = x + self.ffn(self.ln2(x)) # 残差连接 2
return x
# ============================================================
# 6. GPT — 自回归语言模型
# ============================================================
class GPT(nn.Module):
"""
微型 GPT（Generative Pre-trained Transformer）模型。
架构组成:
Token Embedding → 与 Position Embedding 相加
→ N 层 TransformerBlock
→ 最终 LayerNorm
→ 线性投影头（输出词表维度 logits）
通过 return_attention=True 可以捕获第一层 TransformerBlock 的注意力中间值，
用于可视化分析和教学演示。
"""
def __init__(
self,
vocab_size, # 词表大小，此处为 4
d_model=512, # 隐藏维度
n_heads=8, # 注意力头数
n_layers=6, # Transformer 块数
d_ff=2048, # FFN 中间维度
max_len=128 # 支持的最大序列长度
):
super().__init__()
# Token Embedding: 将字符索引映射为 d_model 维稠密向量
self.token_embedding = nn.Embedding(vocab_size, d_model)
# Position Embedding: 为位置 0~max_len-1 各分配一个可学习的嵌入向量
# 使模型感知 token 的相对/绝对位置（Transformer 本身无序列感知能力）
self.position_embedding = nn.Embedding(max_len, d_model)
# 用 ModuleList 而非 Sequential，以便按索引遍历并取出特定层输出
self.blocks = nn.ModuleList([
TransformerBlock(d_model, n_heads, d_ff)
for _ in range(n_layers) # 堆叠 n_layers=6 个 Transformer 块
])
# 最终层归一化：在所有 Transformer 块之后、输出投影之前稳定分布
self.ln_f = nn.LayerNorm(d_model)
# 输出投影头：将 d_model 维隐藏状态映射回 vocab_size 维，得到每个 token 的 logits
self.head = nn.Linear(d_model, vocab_size)
def forward(self, idx, return_attention=False):
"""
前向传播。
参数:
idx: 输入 token 索引，形状 (B, T)
return_attention: 是否返回第一层 TransformerBlock 的注意力中间值
返回:
if return_attention=False: logits, 形状 (B, T, vocab_size)
if return_attention=True: (logits, (Q, K, V, scores, scores_scaled, alpha))
"""
B, T = idx.shape # B=批次大小, T=序列长度
# 生成位置索引 [0, 1, 2, ..., T-1]，形状 (1, T)，移至与 idx 相同设备
positions = torch.arange(T).unsqueeze(0).to(idx.device)
# Token 嵌入 + 位置嵌入（逐元素相加）
token_emb = self.token_embedding(idx) # (B, T) → (B, T, d_model)
pos_emb = self.position_embedding(positions) # (1, T) → (1, T, d_model) → 广播
x = token_emb + pos_emb # 嵌入融合: (B, T, d_model)
saved = None # 用于保存第一层注意力的中间结果
for i, block in enumerate(self.blocks):
# 遍历每一层 TransformerBlock
if return_attention and i == 0:
# 仅在第一层 (i==0) 捕获注意力中间值供分析
x, Q, K, V, scores, scores_scaled, alpha = \
block(x, return_attention=True)
saved = (Q, K, V, scores, scores_scaled, alpha)
else:
x = block(x) # 其余层正常前向
x = self.ln_f(x) # 最终 LayerNorm
logits = self.head(x) # (B, T, d_model) → (B, T, vocab_size)
if return_attention:
return logits, saved # 返回预测值 + 第一层注意力数据
return logits
# ============================================================
# 7. 模型实例化与优化器/损失函数配置
# ============================================================
model = GPT( # 实例化 GPT 模型
vocab_size=vocab_size, # 词表大小 = 4
d_model=d_model, # 维度 = 512
n_heads=n_heads, # 8 个注意力头
n_layers=n_layers, # 6 层 Transformer
d_ff=d_ff # FFN 中间维度 = 2048
)
optimizer = optim.Adam(model.parameters(), lr=1e-4) # Adam 优化器，学习率 0.0001
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失：衡量 logits 与目标分布的差距
# ============================================================
# 8. 训练循环
# ============================================================
for epoch in range(1000): # 训练 1000 个 epoch
optimizer.zero_grad() # 清空上一轮的梯度缓存
logits = model(data) # 前向传播：输入 [0,1,2,3] → logits (1,4,4)
# 计算交叉熵损失
# logits.view(-1, vocab_size): (1*4, 4) = (4, 4) — 展平为 4 个样本，每个 4 类
# target.view(-1): (1*4,) = (4,) — 展平为 4 个目标标签
loss = loss_fn(
logits.view(-1, vocab_size),
target.view(-1)
)
loss.backward() # 反向传播：计算所有参数的梯度
optimizer.step() # 参数更新：沿梯度方向优化
if epoch % 100 == 0: # 每 100 轮输出一次当前损失
print(f"\nEpoch {epoch}, loss={loss.item():.6f}")
# ============================================================
# 9. 推理与注意力可视化
# ============================================================
with torch.no_grad(): # 推理阶段禁用梯度计算，节省显存
# 推理时设置 return_attention=True，捕获第一层注意力的中间值
logits, saved = model(data, return_attention=True)
# argmax 取每个位置 logits 最大的索引，即模型预测的下一个字符
pred = torch.argmax(logits, dim=-1) # (1, 4, 4) → (1, 4)
# 解包保存的注意力中间值
Q, K, V, scores, scores_scaled, alpha = saved
# ---------- 打印输入/输出 ----------
print("\n==============================")
print("输入")
print("==============================")
print([idx2char[i.item()] for i in data[0]]) # 索引还原为汉字
print("\n==============================")
print("预测")
print("==============================")
print([idx2char[i.item()] for i in pred[0]])
# ---------- 可视化第一层注意力 ----------
# 以下取 batch=0, head=0 的数据展示注意力计算全过程
print("\n==============================")
print("Q")
print("==============================")
print(Q[0, 0].detach().numpy().round(3)) # 第 0 个 batch 第 0 个头: (T, d_head)
print("\n==============================")
print("K")
print("==============================")
print(K[0, 0].detach().numpy().round(3))
print("\n==============================")
print("V")
print("==============================")
print(V[0, 0].detach().numpy().round(3))
print("\n==============================")
print("scores = QK^T") # S = Q · Kᵀ，原始相似度
print("==============================")
print(scores[0, 0].detach().numpy().round(3)) # (T, T) 矩阵，每行是查询对各键的得分
print("\n==============================")
print("scores_scaled") # S / √d_k，缩放后
print("==============================")
print(scores_scaled[0, 0].detach().numpy().round(3))
print("\n==============================")
print("softmax alpha") # α = softmax(S_scaled)
print("==============================")
print(alpha[0, 0].detach().numpy().round(3)) # 每行和≈1，因果掩码使 α[i][j>i]=0
