注意力机制与 Transformer 架构

import torch
import torch.nn as nn
import math

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k: int, dropout: float = 0.1):
        super().__init__()
        self.d_k = d_k
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        """
        Args:
            Q, K, V: (batch, heads, seq_len, d_k)
            mask: (batch, 1, 1, seq_len) 或 (batch, 1, seq_len, seq_len)
        Returns:
            output: (batch, heads, seq_len, d_k)
            weights: (batch, heads, seq_len, seq_len)
        """
        scores = torch.matmul(Q, K.transpose(-2, -1))
        scores = scores / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        output = torch.matmul(attention_weights, V)
        return output, attention_weights

# 演示：注意力权重的含义
d_k = 64
attention = ScaledDotProductAttention(d_k)
Q = torch.randn(1, 1, 3, d_k)  # 3个位置的Query
K = Q.clone()  # 自注意力：Q=K
V = torch.randn(1, 1, 3, d_k)
output, weights = attention(Q, K, V)
print(f"注意力权重形状: {weights.shape}")
print(f"(batch=1, heads=1, 3个位置×3个位置)")
# 每行是一个位置对所有位置的注意力分布
print(f"位置0的注意力分布: {weights[0,0,0,:].detach().numpy().round(3)}")

# 缩放因子的重要性演示
import torch
import math

def softmax_without_scale(d_k=64):
    """演示不缩放时 softmax 的梯度消失问题"""
    for d in [8, 32, 64, 128, 512]:
        # 模拟点积结果（均值为0，方差为d的高斯分布）
        scores = torch.randn(1000, d)
        weights = torch.softmax(scores, dim=-1)
        # 检查最大权重（接近1表示sharp）
        max_w = weights.max(dim=-1).values.mean().item()
        # 检查梯度（softmax的Jacobian范数）
        weights.requires_grad_(True)
        loss = weights.sum()
        loss.backward()
        grad_norm = weights.grad.norm().item()
        print(f"  d_k={d:4d} | max_weight={max_w:.4f} | grad_norm={grad_norm:.6f}")

print("=== 无缩放时 softmax 的行为 ===")
softmax_without_scale()

print(" === 有缩放 (除以 √d_k) ===")
for d in [8, 32, 64, 128, 512]:
    scores = torch.randn(1000, d) / math.sqrt(d)
    weights = torch.softmax(scores, dim=-1)
    max_w = weights.max(dim=-1).values.mean().item()
    print(f"  d_k={d:4d} | max_weight={max_w:.4f}")

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_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.attention = ScaledDotProductAttention(self.d_k, dropout)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 线性变换并分头: (batch, seq_len, d_model) -> (batch, heads, seq_len, d_k)
        Q = self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 多头注意力
        attn_output, weights = self.attention(Q, K, V, mask)
        
        # 拼接所有头: (batch, heads, seq_len, d_k) -> (batch, seq_len, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous() \
            .view(batch_size, -1, self.d_model)
        
        output = self.W_o(attn_output)
        return self.dropout(output), weights

# 参数量对比
d_model, num_heads = 512, 8
single_head_params = 3 * (d_model * d_model) + d_model  # Q, K, V 投影 + 输出
multi_head_params = 3 * (d_model * d_model) + d_model   # 相同！
print(f"单头参数量: {single_head_params:,}")
print(f"多头参数量: {multi_head_params:,} (相同，但表达能力更强)")

class TransformerEncoderLayer(nn.Module):
    """Transformer 编码器单层"""
    
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 子层1: 自注意力 + 残差 + LayerNorm
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        # 子层2: FFN + 残差 + LayerNorm
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        return x

class TransformerDecoderLayer(nn.Module):
    """Transformer 解码器单层"""
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        # 子层1: 掩码自注意力
        attn_out, _ = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_out))
        # 子层2: 交叉注意力
        cross_out, _ = self.cross_attn(x, memory, memory, src_mask)
        x = self.norm2(x + self.dropout(cross_out))
        # 子层3: FFN
        ffn_out = self.ffn(x)
        x = self.norm3(x + self.dropout(ffn_out))
        return x

class PositionalEncoding(nn.Module):
    """原始 Transformer 的正弦位置编码"""
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).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 shape: (batch, seq_len, d_model)
        return x + self.pe[:, :x.size(1)]

# 可视化位置编码
pe = PositionalEncoding(d_model=64)
positions = pe.pe[0, :20, :].numpy()
print(f"位置编码形状: {positions.shape}")
print(f"(20个位置 × 64维)")
# 不同维度使用不同频率
print(f" 位置 0 的前8维: {positions[0, :8].round(3)}")
print(f"位置 5 的前8维: {positions[5, :8].round(3)}")
print(f"位置 10 的前8维: {positions[10, :8].round(3)}")

class FFN(nn.Module):
    """Transformer 前馈网络"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        return self.linear2(self.dropout(self.activation(self.linear1(x))))

# FFN 的扩展比
d_model = 512
d_ff = 2048  # 4倍扩展
print(f"d_model={d_model}, d_ff={d_ff}, 扩展比={d_ff/d_model}x")
ffn = FFN(d_model, d_ff)
x = torch.randn(2, 10, d_model)
out = ffn(x)
print(f"输入: {x.shape} -> 中间: {d_ff}维 -> 输出: {out.shape}")
print(f"FFN 参数量: {sum(p.numel() for p in ffn.parameters()):,}")