注意力机制：从 Seq2Seq 到 Transformer

import torch
import torch.nn as nn

class SimpleEncoder(nn.Module):
    """基于 LSTM 的编码器"""
    def __init__(self, vocab_size, embed_dim, hidden_dim, n_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, n_layers,
                           dropout=0.3, batch_first=True)

    def forward(self, src):
        # src: (batch, seq_len)
        embedded = self.embedding(src)  # (batch, seq_len, embed_dim)
        outputs, (hidden, cell) = self.lstm(embedded)
        # outputs: (batch, seq_len, hidden_dim) — 所有时间步的隐藏状态
        # hidden: (n_layers, batch, hidden_dim) — 最终隐藏状态
        return outputs, hidden, cell

class Seq2Seq(nn.Module):
    """完整的 Seq2Seq 模型"""
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        # src: (batch, src_len), trg: (batch, trg_len)
        encoder_outputs, hidden, cell = self.encoder(src)
        outputs = self.decoder(trg, hidden, cell, encoder_outputs,
                              teacher_forcing_ratio)
        return outputs

# 使用示例
enc = SimpleEncoder(vocab_size=10000, embed_dim=256, hidden_dim=512)
dec = SimpleDecoder(vocab_size=10000, embed_dim=256, hidden_dim=512)
model = Seq2Seq(enc, dec)

src = torch.randint(1, 10000, (32, 20))   # batch=32, seq_len=20
trg = torch.randint(1, 10000, (32, 25))   # batch=32, seq_len=25
output = model(src, trg)
print(f"输出形状: {output.shape}")  # (32, 25, 10000)

import torch
import torch.nn as nn
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """缩放点积注意力的纯 PyTorch 实现
    
    Args:
        Q: (batch, n_q, d_k) 查询矩阵
        K: (batch, n_k, d_k) 键矩阵
        V: (batch, n_k, d_v) 值矩阵
        mask: 可选的掩码矩阵，屏蔽不需要关注的位置
    
    Returns:
        output: (batch, n_q, d_v) 注意力输出
        weights: (batch, n_q, n_k) 注意力权重（可用于可视化）
    """
    d_k = Q.size(-1)
    
    # 1. 计算点积相似度: Q · K^T
    scores = torch.matmul(Q, K.transpose(-2, -1))
    # scores shape: (batch, n_q, n_k)
    
    # 2. 缩放: 除以 sqrt(d_k)
    scores = scores / math.sqrt(d_k)
    
    # 3. 应用掩码（如果有）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 4. softmax 归一化
    weights = torch.softmax(scores, dim=-1)
    # weights shape: (batch, n_q, n_k)
    
    # 5. 加权求和
    output = torch.matmul(weights, V)
    # output shape: (batch, n_q, d_v)
    
    return output, weights

# 缩放因子的必要性演示
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

d_k_values = [16, 64, 256, 1024]
for d_k in d_k_values:
    Q = torch.randn(1, 1, d_k)
    K = torch.randn(1, 10, d_k)
    
    # 未缩放的点积
    raw_scores = torch.matmul(Q, K.transpose(-2, -1))
    print(f"d_k={d_k:4d}: 未缩放 scores 均值={raw_scores.mean():.2f}, "
          f"标准差={raw_scores.std():.2f}")
    
    # 缩放后的点积
    scaled_scores = raw_scores / math.sqrt(d_k)
    print(f"         缩放后 scores 均值={scaled_scores.mean():.2f}, "
          f"标准差={scaled_scores.std():.2f}")
    
    # softmax 后的熵（衡量分布的均匀程度）
    raw_weights = torch.softmax(raw_scores, dim=-1)
    scaled_weights = torch.softmax(scaled_scores, dim=-1)
    
    raw_entropy = -(raw_weights * (raw_weights + 1e-10).log()).sum().item()
    scaled_entropy = -(scaled_weights * (scaled_weights + 1e-10).log()).sum().item()
    
    print(f"         未缩放 softmax 熵: {raw_entropy:.4f}")
    print(f"         缩放后 softmax 熵: {scaled_entropy:.4f}")
    print()

import torch
import torch.nn as nn
import math

class MultiHeadSelfAttention(nn.Module):
    """多头自注意力的完整实现"""
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度
        
        # Q, K, V 的线性投影
        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.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # 1. 线性投影 + 多头拆分
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        
        # transpose: (batch, n_heads, seq_len, d_k)
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        
        # 2. 缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        weights = torch.softmax(scores, dim=-1)
        weights = self.dropout(weights)
        
        # 3. 加权求和
        context = torch.matmul(weights, V)
        
        # 4. 多头拼接 + 输出投影
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_len, self.d_model)
        output = self.W_O(context)
        
        return output, weights

import torch
import torch.nn as nn

class EncoderLayer(nn.Module):
    """Transformer 编码器层"""
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 子层1: 多头自注意力 + 残差 + LayerNorm
        attn_out, _ = self.self_attn(x, mask)
        x = self.layer_norm1(x + self.dropout(attn_out))
        # 子层2: 前馈网络 + 残差 + LayerNorm
        ff_out = self.feed_forward(x)
        x = self.layer_norm2(x + self.dropout(ff_out))
        return x


class DecoderLayer(nn.Module):
    """Transformer 解码器层"""
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.masked_self_attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.layer_norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # 子层1: 掩码自注意力
        attn_out, _ = self.masked_self_attn(x, tgt_mask)
        x = self.layer_norm1(x + self.dropout(attn_out))
        # 子层2: 交叉注意力
        cross_out, _ = self.cross_attn(x, enc_output, src_mask)
        x = self.layer_norm2(x + self.dropout(cross_out))
        # 子层3: 前馈网络
        ff_out = self.feed_forward(x)
        x = self.layer_norm3(x + self.dropout(ff_out))
        return x

class Transformer(nn.Module):
    """完整的 Transformer 模型"""
    def __init__(self, src_vocab, tgt_vocab, d_model=512, n_heads=8,
                 num_layers=6, d_ff=2048, max_len=5000, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        
        self.src_embedding = nn.Embedding(src_vocab, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        self.output_linear = nn.Linear(d_model, tgt_vocab)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        # 嵌入 + 位置编码
        src_emb = self.dropout(self.positional_encoding(
            self.src_embedding(src) * math.sqrt(self.d_model)))
        tgt_emb = self.dropout(self.positional_encoding(
            self.tgt_embedding(tgt) * math.sqrt(self.d_model)))
        
        # 编码
        enc_output = src_emb
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)
        
        # 解码
        dec_output = tgt_emb
        for layer in self.decoder_layers:
            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
        
        # 输出投影
        logits = self.output_linear(dec_output)
        return logits

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """正弦/余弦位置编码（原始 Transformer 方案）"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(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, dtype=torch.float)
                            * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度：sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度：cos
        
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

# RoPE 旋转位置编码的简化实现
import torch

def rotate_half(x):
    """将向量对半拆分并旋转"""
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(q, k, position_ids):
    """应用 RoPE 到 Query 和 Key
    
    Args:
        q: (batch, heads, seq_len, head_dim)
        k: (batch, heads, seq_len, head_dim)
        position_ids: (batch, seq_len) 位置索引
    """
    head_dim = q.shape[-1]
    half_dim = head_dim // 2
    
    # 计算旋转角度
    inv_freq = 1.0 / (10000 ** (torch.arange(0, half_dim, 2, 
                       dtype=torch.float32) / half_dim))
    inv_freq = inv_freq.to(q.device)
    
    # (batch, seq_len, half_dim)
    freqs = position_ids.float().unsqueeze(-1) * inv_freq
    freqs = torch.cat([freqs, freqs], dim=-1)  # 扩展到 head_dim
    
    # (1, 1, seq_len, head_dim)
    freqs = freqs.unsqueeze(0).unsqueeze(0)
    
    # 旋转变换: q → q·cos + rotate_half(q)·sin
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    
    return q_embed, k_embed

# 可视化多头注意力的不同模式
import torch
import torch.nn as nn
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def visualize_attention_heads(attention_weights, sentence_tokens):
    """可视化多头注意力的权重分布
    
    Args:
        attention_weights: (n_heads, seq_len, seq_len)
        sentence_tokens: list of str
    """
    n_heads = attention_weights.shape[0]
    fig, axes = plt.subplots(2, 4, figsize=(20, 8))
    axes = axes.flatten()
    
    for h in range(n_heads):
        im = axes[h].imshow(attention_weights[h], cmap='viridis',
                          aspect='auto', vmin=0, vmax=1)
        axes[h].set_title(f'Head {h+1}', fontsize=10)
        axes[h].set_xticks(range(len(sentence_tokens)))
        axes[h].set_xticklabels(sentence_tokens, rotation=45, fontsize=7)
        axes[h].set_yticks(range(len(sentence_tokens)))
        axes[h].set_yticklabels(sentence_tokens, fontsize=7)
    
    plt.colorbar(im, ax=axes, fraction=0.02)
    plt.tight_layout()
    plt.savefig('multi_head_attention.png', dpi=150)
    plt.close()


# 示例：分析一个训练好的模型
model = load_trained_transformer()
attention_weights = get_attention_weights(model, "The cat sat on the mat")
visualize_attention_heads(attention_weights[0],  # 第一层
                         ["The", "cat", "sat", "on", "the", "mat"])

# 计算复杂度对比
def compute_complexity(seq_len, d_model, n_heads, num_layers):
    """计算 Transformer 各组件的复杂度"""
    d_k = d_model // n_heads
    
    # 自注意力: O(n² · d)
    self_attn = 4 * seq_len**2 * d_model + 2 * seq_len * d_model**2
    
    # FFN: O(n · d²)
    ffn = 8 * seq_len * d_model**2  # d_ff = 4 * d_model, 2 次矩阵乘法
    
    # 总复杂度（单层）
    per_layer = self_attn + ffn
    total = per_layer * num_layers
    
    print(f"序列长度: {seq_len}")
    print(f"d_model: {d_model}, 头数: {n_heads}, 层数: {num_layers}")
    print(f"自注意力: {self_attn/1e9:.2f} G FLOPs")
    print(f"FFN: {ffn/1e9:.2f} G FLOPs")
    print(f"单层总计: {per_layer/1e9:.2f} G FLOPs")
    print(f"模型总计: {total/1e9:.2f} G FLOPs")
    print()

# 对比不同配置
compute_complexity(512, 512, 8, 6)    # 原始 Transformer
compute_complexity(512, 768, 12, 12)  # BERT-base
compute_complexity(4096, 4096, 32, 32)  # GPT-3 小配置

# BERT 风格的编码器-only 模型
import torch
import torch.nn as nn

class BERTLikeModel(nn.Module):
    """类 BERT 的预训练模型"""
    def __init__(self, vocab_size, d_model=768, n_heads=12, 
                 num_layers=12, max_len=512):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.segment_embedding = nn.Embedding(2, d_model)
        
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(d_model, n_heads) for _ in range(num_layers)
        ])
        
        # MLM 预测头
        self.mlm_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.LayerNorm(d_model),
            nn.Linear(d_model, vocab_size)
        )
        # NSP (Next Sentence Prediction) 头
        self.nsp_head = nn.Linear(d_model, 2)
        self.pooler = nn.Linear(d_model, d_model)
    
    def forward(self, input_ids, token_type_ids=None, attention_mask=None):
        batch_size, seq_len = input_ids.shape
        positions = torch.arange(seq_len, device=input_ids.device)
        positions = positions.unsqueeze(0).expand(batch_size, -1)
        
        x = (self.token_embedding(input_ids) +
             self.position_embedding(positions))
        if token_type_ids is not None:
            x += self.segment_embedding(token_type_ids)
        
        for layer in self.encoder_layers:
            x = layer(x, attention_mask)
        
        # MLM 输出
        mlm_logits = self.mlm_head(x)
        # NSP 输出（使用 [CLS] token 的表示）
        pooled = torch.tanh(self.pooler(x[:, 0, :]))
        nsp_logits = self.nsp_head(pooled)
        
        return mlm_logits, nsp_logits

# 稀疏注意力：局部窗口 + 全局 token
import torch

class SparseSelfAttention(nn.Module):
    """Longformer 风格的稀疏自注意力
    
    每个 token 关注:
    1. 局部窗口内的所有 token (w=512)
    2. 全局 token (如 [CLS], 段落标题等)
    """
    def __init__(self, d_model=768, n_heads=12, window_size=512):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.window_size = window_size
        
        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)
    
    def forward(self, x, global_mask=None):
        batch_size, seq_len, _ = x.size()
        half_w = self.window_size // 2
        
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k)
        
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        
        # 局部窗口注意力 (简化版)
        # 实际实现中使用 sliding window 避免 O(n²)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k  0.5)
        
        # 创建滑动窗口掩码
        window_mask = torch.ones(seq_len, seq_len, device=x.device)
        for i in range(seq_len):
            window_mask[i, :max(0, i-half_w)] = 0
            window_mask[i, i+half_w+1:] = 0
        
        # 全局 token 可以看到所有位置
        if global_mask is not None:
            window_mask[global_mask, :] = 1
            window_mask[:, global_mask] = 1
        
        scores = scores.masked_fill(window_mask == 0, -1e9)
        weights = torch.softmax(scores, dim=-1)
        
        context = torch.matmul(weights, V)
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_len, self.d_model)
        
        return self.W_O(context)

# KV Cache 的推理优化示例
import torch

class KVCacheTransformer:
    """带 KV Cache 的 Transformer 推理"""
    def __init__(self, model):
        self.model = model
        self.kv_cache = None
    
    def generate(self, input_ids, max_new_tokens=100, temperature=1.0):
        """自回归生成，使用 KV Cache 加速"""
        # 预填充阶段：一次性处理所有输入 token
        outputs = self.model(input_ids, use_cache=True)
        logits = outputs.logits[:, -1, :] / temperature
        self.kv_cache = outputs.past_key_values
        
        generated = input_ids
        for _ in range(max_new_tokens):
            # 采样下一个 token
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            generated = torch.cat([generated, next_token], dim=-1)
            
            if next_token.item() == self.model.config.eos_token_id:
                break
            
            # 使用 KV Cache：只传递新生成的 token
            outputs = self.model(next_token, use_cache=True,
                                past_key_values=self.kv_cache)
            logits = outputs.logits[:, -1, :] / temperature
            self.kv_cache = outputs.past_key_values
        
        return generated
    
    def get_cache_size_mb(self):
        """估算 KV Cache 的内存占用"""
        if self.kv_cache is None:
            return 0
        total_bytes = 0
        for layer_kv in self.kv_cache:
            for tensor in layer_kv:  # (key, value)
                total_bytes += tensor.numel() * tensor.element_size()
        return total_bytes / (1024 * 1024)

# 注意力权重可解释性分析
import torch
import numpy as np

def analyze_attention_patterns(attention_weights, n_layers, n_heads):
    """分析注意力权重的统计特性
    
    Args:
        attention_weights: (batch, n_layers, n_heads, seq_len, seq_len)
    """
    results = []
    
    for layer in range(n_layers):
        for head in range(n_heads):
            weights = attention_weights[0, layer, head]
            
            # 熵：衡量注意力分布的均匀程度
            entropy = -(weights * (weights + 1e-10).log()).sum(dim=-1).mean()
            max_entropy = np.log(weights.size(-1))
            normalized_entropy = entropy / max_entropy
            
            # 集中度：有多少比例的注意力集中在 top-k 位置上
            top5_ratio = (weights.topk(5, dim=-1).values.sum(dim=-1)
                        / weights.sum(dim=-1)).mean()
            
            # 对角线倾向（自注意力权重）
            diag = torch.diag(weights).mean()
            
            results.append({
                'layer': layer, 'head': head,
                'entropy': normalized_entropy.item(),
                'top5_ratio': top5_ratio.item(),
                'diag_weight': diag.item(),
            })
    
    import pandas as pd
    df = pd.DataFrame(results)
    print("各层各头注意力统计:")
    print(df.to_string(index=False))
    return df

# 从零实现完整的注意力层（练习框架）
import torch
import torch.nn as nn
import math

class AttentionFromScratch(nn.Module):
    """练习：从零实现自注意力层
    
    TODO: 补全以下方法的实现
    """
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        # TODO: 创建 Q, K, V 和输出投影层
        self.W_Q = None
        self.W_K = None
        self.W_V = None
        self.W_O = None
        
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # TODO: 1. 线性投影 Q, K, V
        # TODO: 2. 多头拆分 (view + transpose)
        # TODO: 3. 缩放点积 Q·K^T / sqrt(d_k)
        # TODO: 4. 应用掩码（如果有）
        # TODO: 5. softmax 归一化
        # TODO: 6. 加权求和 weights·V
        # TODO: 7. 多头拼接 (transpose + view)
        # TODO: 8. 输出投影
        # TODO: 9. 残差连接 + LayerNorm
        
        pass  # 替换为你的实现


# 验证你的实现是否正确
def test_attention_implementation():
    """对比从零实现与 PyTorch 内置实现"""
    torch.manual_seed(42)
    d_model, n_heads = 64, 4
    batch_size, seq_len = 2, 10
    
    x = torch.randn(batch_size, seq_len, d_model)
    
    # PyTorch 内置
    pytorch_attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
    pytorch_out, _ = pytorch_attn(x, x, x)
    
    # 你的实现
    your_attn = AttentionFromScratch(d_model, n_heads)
    # TODO: 将 pytorch_attn 的参数复制给你的实现
    # your_out = your_attn(x)
    
    # TODO: 验证输出是否一致 (误差 < 1e-5)
    # diff = (pytorch_out - your_out).abs().max()
    # assert diff < 1e-5, f"误差过大: {diff}"
    
    print("测试框架已就绪，请完成 AttentionFromScratch 的实现")

test_attention_implementation()