Transformer架构完全解析：从注意力机制到GPT原理

# 伪代码展示注意力机制的核心步骤

# 输入：
# Q: query 向量，形状 (seq_len, d_k)
# K: key 向量，形状 (seq_len, d_k)
# V: value 向量，形状 (seq_len, d_v)

# 步骤1：计算 Q 和 K 的相似度（点积）
# 使用 d_k 的平方根进行缩放，防止点积值过大
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
# scores 形状: (seq_len, seq_len)

# 步骤2：对相似度进行 softmax，得到注意力权重
attention_weights = F.softmax(scores, dim=-1)
# 每个 query 对所有 key 的注意力权重加起来 = 1

# 步骤3：用注意力权重对 V 加权求和
output = torch.matmul(attention_weights, V)
# 输出形状: (seq_len, d_v)

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    缩放点积注意力机制的实现
    
    参数:
        Q: Query 张量，形状 (batch_size, num_heads, seq_len_q, d_k)
        K: Key 张量，形状 (batch_size, num_heads, seq_len_k, d_k)
        V: Value 张量，形状 (batch_size, num_heads, seq_len_k, d_v)
        mask: 可选掩码，形状 (batch_size, num_heads, seq_len_q, seq_len_k)
    
    返回:
        output: 注意力输出，形状 (batch_size, num_heads, seq_len_q, d_v)
        attention_weights: 注意力权重，形状 (batch_size, num_heads, seq_len_q, seq_len_k)
    """
    d_k = Q.size(-1)  # 向量维度
    
    # 步骤1：计算 Q 和 K 的点积，然后缩放
    # torch.matmul(Q, K.transpose(-2, -1)) 做矩阵乘法
    # 最后一项是转置：K 的最后两个维度交换
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 步骤2：如果有 mask，加上一个很大的负数，使 softmax 后趋近于 0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 步骤3：计算注意力权重（softmax）
    attention_weights = F.softmax(scores, dim=-1)
    
    # 步骤4：用注意力权重加权 V
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

# 使用示例
batch_size = 2
num_heads = 8
seq_len_q = 10
seq_len_k = 10
d_k = 64
d_v = 64

Q = torch.randn(batch_size, num_heads, seq_len_q, d_k)
K = torch.randn(batch_size, num_heads, seq_len_k, d_k)
V = torch.randn(batch_size, num_heads, seq_len_k, d_v)

output, weights = scaled_dot_product_attention(Q, K, V)
print(f"输出形状: {output.shape}")  # (2, 8, 10, 64)
print(f"注意力权重形状: {weights.shape}")  # (2, 8, 10, 10)

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制
    
    与其只进行一次注意力计算，不如将 Q、K、V 分别投影到
    多个子空间（head），并行计算注意力，最后拼接结果。
    """
    
    def __init__(self, d_model, num_heads):
        """
        参数:
            d_model: 模型的隐藏层维度（如 512）
            num_heads: 注意力头的数量（如 8）
        """
        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  # 每个头的维度
        
        # 定义 Q、K、V 的投影矩阵
        # 输出形状: (d_model, d_model)
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        
        # 最终输出投影
        self.W_o = nn.Linear(d_model, d_model, bias=False)
    
    def split_heads(self, x, batch_size):
        """
        将最后一个维度拆分成 num_heads 个子维度
        输入: (batch_size, seq_len, d_model)
        输出: (batch_size, num_heads, seq_len, d_k)
        """
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # 交换第1和第2维度
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 步骤1：线性投影，然后分头
        Q = self.split_heads(self.W_q(Q), batch_size)   # (B, H, Lq, dk)
        K = self.split_heads(self.W_k(K), batch_size)   # (B, H, Lk, dk)
        V = self.split_heads(self.W_v(V), batch_size)   # (B, H, Lk, dv)
        
        # 步骤2：计算缩放点积注意力
        # attention shape: (batch_size, num_heads, seq_len_q, d_v)
        # attention_weights shape: (batch_size, num_heads, seq_len_q, seq_len_k)
        attention, attention_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # 步骤3：合并多头（还原维度）
        # attention: (B, H, Lq, dk) -> (B, Lq, H*dk) = (B, Lq, d_model)
        attention = attention.transpose(1, 2).contiguous()
        attention = attention.view(batch_size, -1, self.d_model)
        
        # 步骤4：最终线性投影
        output = self.W_o(attention)
        
        return output, attention_weights

# 使用示例
d_model = 512
num_heads = 8
mh_attention = MultiHeadAttention(d_model, num_heads)

# 模拟输入: batch_size=2, seq_len=10, d_model=512
Q = torch.randn(2, 10, 512)
K = torch.randn(2, 10, 512)
V = torch.randn(2, 10, 512)

output, weights = mh_attention(Q, K, K)
print(f"输出形状: {output.shape}")  # (2, 10, 512)
print(f"注意力权重形状: {weights.shape}")  # (2, 8, 10, 10)

import torch
import math

def positional_encoding(seq_len, d_model):
    """
    生成位置编码矩阵
    
    公式：
    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    
    其中 pos 是位置，i 是维度索引
    
    参数:
        seq_len: 序列长度
        d_model: 模型维度（必须为偶数）
    
    返回:
        pe: 位置编码矩阵，形状 (seq_len, d_model)
    """
    pe = torch.zeros(seq_len, d_model)
    
    # 生成位置索引 [0, 1, 2, ..., seq_len-1]
    position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
    # 形状: (seq_len, 1)
    
    # 计算除数项: 10000^(2i/d_model)
    # 使用 log 避免指数溢出
    div_term = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
    )
    # 形状: (d_model/2,)
    
    # 偶数维度用 sin
    pe[:, 0::2] = torch.sin(position * div_term)
    # 奇数维度用 cos
    pe[:, 1::2] = torch.cos(position * div_term)
    
    return pe

# 示例
seq_len = 50  # 最长序列 50
d_model = 512  # 模型维度

pe = positional_encoding(seq_len, d_model)
print(f"位置编码形状: {pe.shape}")  # (50, 512)

# 可视化：绘制第 0、1、2 号位置编码
import matplotlib.pyplot as plt

plt.figure(figsize=(14, 5))
for i in [0, 1, 2]:
    plt.plot(range(50), pe[:, i].numpy(), label=f"dim {i}")
plt.xlabel("Position")
plt.ylabel("Encoding value")
plt.title("Positional Encoding (first 3 dimensions)")
plt.legend()
plt.grid(True)
plt.show()

import torch
import torch.nn as nn
import math

class EncoderLayer(nn.Module):
    """
    Transformer Encoder 层
    
    包含：
    1. Multi-Head Self-Attention（自注意力）
    2. Feed-Forward Network（前馈网络）
    每层都使用残差连接和层归一化
    """
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),      # 扩张层
            nn.ReLU(),                      # 激活函数
            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, mask=None):
        # 自注意力 + 残差连接
        # Q=K=V=x，即"自己注意自己"
        attn_output, _ = self.self_attention(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # 前馈网络 + 残差连接
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        
        return x

class Encoder(nn.Module):
    """
    完整的 Transformer Encoder
    由 N 个 EncoderLayer 堆叠而成
    """
    
    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.num_layers = num_layers
    
    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return x

# 参数配置示例（以 BERT-base 为例）
num_layers = 12        # 12 层 encoder
d_model = 768          # 隐藏层维度
num_heads = 12          # 12 个注意力头
d_ff = 3072            # 前馈网络中间层维度
dropout = 0.1

encoder = Encoder(num_layers, d_model, num_heads, d_ff, dropout)
print(f"模型参数量估计: {sum(p.numel() for p in encoder.parameters()) / 1e6:.1f}M")

def create_mask(seq_len):
    """
    创建下三角掩码，用于 Masked Self-Attention
    确保 Decoder 在预测第 i 个 token 时，只能看到 0..i-1 的内容
    """
    # 生成下三角矩阵，1 表示可以关注，0 表示需要遮住
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask  # 形状 (seq_len, seq_len)

def create_padding_mask(seq):
    """
    创建填充掩码，遮住 padding 部分（通常是 0）
    """
    return (seq != 0).unsqueeze(1).unsqueeze(2)  # (B, 1, 1, seq_len)

# 使用示例
seq_len = 10
mask = create_mask(seq_len)
print("掩码形状:", mask.shape)
print("掩码内容（下三角=1，上三角=0）:")
print(mask.int())

import torch
import torch.nn as nn

class GPTLMHead(nn.Module):
    """
    GPT 语言模型头：预测下一个 token
    
    结构：
    Transformer Decoder → LayerNorm → 线性映射到词表大小 → softmax
    """
    
    def __init__(self, d_model, vocab_size):
        super().__init__()
        self.transformer = Decoder(num_layers=12, d_model=d_model, ...)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        # 权重绑定：Embedding 和 LM Head 共享权重（节省参数）
        self.lm_head.weight = self.transformer.embed_tokens.weight
    
    def forward(self, input_ids):
        """
        参数:
            input_ids: token ID 序列，形状 (batch_size, seq_len)
        返回:
            logits: 下一个 token 的预测 logits，形状 (batch_size, seq_len, vocab_size)
        """
        # 通过 Transformer Decoder
        hidden_states = self.transformer(input_ids)
        # 投影到词表
        logits = self.lm_head(hidden_states)
        return logits

# 训练时的损失计算
def compute_lm_loss(logits, labels):
    """
    计算语言模型损失（Cross Entropy）
    
    参数:
        logits: 模型输出，形状 (B, seq_len, vocab_size)
        labels: 目标 token ID，形状 (B, seq_len)
        标签是 "下一个 token"，即 input_ids 右移一位
    """
    # 将 logits 和 labels 转成正确的形状
    # 预测第 i 个位置的 token 时，labels 应该是 input_ids[i+1]
    # 所以我们取 logits[:, :-1] 和 labels[:, 1:]
    shift_logits = logits[:, :-1, :].contiguous()  # (B, seq_len-1, V)
    shift_labels = labels[:, 1:].contiguous()        # (B, seq_len-1)
    
    # 计算交叉熵损失
    loss = nn.functional.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    return loss

# 示例
batch_size = 4
seq_len = 128
vocab_size = 50000
d_model = 768

model = GPTLMHead(d_model, vocab_size)
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

logits = model(input_ids)
loss = compute_lm_loss(logits, input_ids)

print(f"Logits 形状: {logits.shape}")   # (4, 128, 50000)
print(f"损失值: {loss.item():.4f}")      # 约 4-8 左右（取决于模型训练程度）

def generate_gpt(model, start_tokens, max_new_tokens=50, temperature=1.0, top_k=50):
    """
    GPT 自回归生成
    
    参数:
        model: 训练好的 GPT 模型
        start_tokens: 起始 token 序列，形状 (1, seq_len)
        max_new_tokens: 最多生成多少个新 token
        temperature: 采样温度，>1 增加随机性，<1 减少随机性
        top_k: 只从概率最高的 k 个 token 中采样
    返回:
        generated: 生成的完整序列
    """
    model.eval()
    
    # 如果还有 GPU 就用 GPU
    device = next(model.parameters()).device
    input_ids = start_tokens.to(device)
    
    with torch.no_grad():
        for _ in range(max_new_tokens):
            # 取最后 block_size 个 token（防止超出训练时的上下文长度）
            logits = model(input_ids)
            
            # 只关心最后一个位置的预测
            logits = logits[:, -1, :] / temperature  # (1, vocab_size)
            
            # Top-k 采样
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            
            # 转为概率分布并采样
            probs = torch.softmax(logits, dim=-1)  # (1, vocab_size)
            next_token = torch.multinomial(probs, num_samples=1)  # (1, 1)
            
            # 追加到序列
            input_ids = torch.cat([input_ids, next_token], dim=1)
            
            # 遇到 EOS（End of Sentence）就停止
            if next_token.item() == tokenizer.eos_token_id:
                break
    
    return input_ids

# 使用示例
# 假设 start = "今天天气" 的 tokenized 结果
start_tokens = torch.tensor([[101, 2023, 3152, 5445]])  # 简化示例

generated = generate_gpt(model, start_tokens, max_new_tokens=100, temperature=0.8)
print(f"生成结果 ID: {generated}")
print(f"解码后文字: {tokenizer.decode(generated[0])}")

"""
完整的 Transformer 编码器实现（可运行）
"""

import torch
import torch.nn as nn
import math
import torch.nn.functional as F


class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-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_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


class TransformerEncoder(nn.Module):
    """完整的 Transformer 编码器"""
    
    def __init__(
        self,
        vocab_size=50000,
        d_model=512,
        num_heads=8,
        num_layers=6,
        d_ff=2048,
        dropout=0.1,
        max_len=5000
    ):
        super().__init__()
        
        # 词嵌入 + 位置编码
        self.embed_tokens = nn.Embedding(vocab_size, d_model)
        self.embed_positions = PositionalEncoding(d_model, max_len, dropout)
        
        # N 层 Encoder
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, input_ids, mask=None):
        # 词嵌入 + 位置编码
        x = self.embed_tokens(input_ids)
        x = self.embed_positions(x)
        
        # N 层编码
        for layer in self.layers:
            x = layer(x, mask)
        
        return self.norm(x)


# 参数统计
model = TransformerEncoder()
total_params = sum(p.numel() for p in model.parameters())
total_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"总参数量: {total_params:,}")
print(f"可训练参数量: {total_trainable:,}")