输入（法语）："Je t'aime"
输出（英语）："I love you"

┌─────────────────────────────────────────────────┐
│               Transformer架构                    │
├─────────────────────────────────────────────────┤
│                                                 │
│  输入: "Je t'aime"                              │
│      ↓                                          │
│  ┌──────────────┐                               │
│  │   编码器     │  ← 理解输入，提取语义          │
│  │  (Encoder)   │                               │
│  └──────────────┘                               │
│      ↓                                          │
│  [语义表示向量]                                 │
│      ↓                                          │
│  ┌──────────────┐                               │
│  │   解码器     │  ← 基于语义，生成翻译          │
│  │  (Decoder)   │                               │
│  └──────────────┘                               │
│      ↓                                          │
│  输出: "I love you"                              │
└─────────────────────────────────────────────────┘

输入嵌入 → 位置编码
    ↓
┌──────────────────┐
│ 编码器层 × N     │  （通常N=6或12）
│                  │
│  ┌────────────┐  │
│  │ 自注意力   │  │  ← 捕获全局依赖
│  └────────────┘  │
│       ↓          │
│  ┌────────────┐  │
│  │ 前馈网络   │  │  ← 非线性变换
│  └────────────┘  │
└──────────────────┘
    ↓
输出：每个位置的语义向量

目标嵌入 → 位置编码
    ↓
┌──────────────────┐
│ 解码器层 × N     │
│                  │
│  ┌────────────┐  │
│  │ 自注意力   │  │  ← 只能看到左边（因果掩码）
│  └────────────┘  │
│       ↓          │
│  ┌────────────┐  │
│  │ 交叉注意力 │  │  ← 关注编码器输出
│  └────────────┘  │
│       ↓          │
│  ┌────────────┐  │
│  │ 前馈网络   │  │
│  └────────────┘  │
└──────────────────┘
    ↓
输出：预测下一个词的概率分布

步骤1: 编码器处理输入
输入: "Je t'aime"
  → 分词: [Je, t', aime]
  → 嵌入: [[e₁], [e₂], [e₃]]
  → 编码器: [[h₁], [h₂], [h₃]]  ← 语义表示

步骤2: 解码器生成输出（自回归）
初始化: [<BOS>]  （Begin of Sequence）

第1步生成:
  输入: [<BOS>]
  查询编码器: [h₁, h₂, h₃]
  预测: "I"

第2步生成:
  输入: [<BOS>, I]
  查询编码器: [h₁, h₂, h₃]
  预测: "love"

第3步生成:
  输入: [<BOS>, I, love]
  查询编码器: [h₁, h₂, h₃]
  预测: "you"

第4步生成:
  输入: [<BOS>, I, love, you]
  查询编码器: [h₁, h₂, h₃]
  预测: <EOS>  ← 结束

最终输出: "I love you"

from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

# 加载T5模型（编码器-解码器架构）
model_name = "t5-small"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

# T5使用任务前缀
text = "translate English to German: The house is wonderful."
inputs = tokenizer(text, return_tensors="pt")

print("输入Token IDs:", inputs.input_ids)
print("输入Tokens:", tokenizer.convert_ids_to_tokens(inputs.input_ids[0]))

# 生成翻译
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_length=50,
        num_beams=4,  # Beam Search
        early_stopping=True
    )

translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("\n翻译结果:", translated)

# 查看模型内部结构
print("\n模型结构:")
print(f"编码器层数: {len(model.encoder.block)}")
print(f"解码器层数: {len(model.decoder.block)}")
print(f"隐藏维度: {model.config.d_model}")
print(f"注意力头数: {model.config.num_heads}")

输入Token IDs: tensor([[13959,  1566,    12,  2968,    10,    37,   629,    19,  1627,     5,      1]])
输入Tokens: ['▁translate', '▁English', '▁to', '▁German', ':', '▁The', '▁house', '▁is', '▁wonderful', '.', '</s>']

翻译结果: Das Haus ist wunderbar.

模型结构:
编码器层数: 6
解码器层数: 6
隐藏维度: 512
注意力头数: 8

输入: "The cat sat on the mat"

RNN处理过程:
t=1: 输入"The"    → 隐状态h₁
t=2: 输入"cat"    → 隐状态h₂  （依赖h₁）
t=3: 输入"sat"    → 隐状态h₃  （依赖h₂）
t=4: 输入"on"     → 隐状态h₄  （依赖h₃）
t=5: 输入"the"    → 隐状态h₅  （依赖h₄）
t=6: 输入"mat"    → 隐状态h₆  （依赖h₅）

输入: "The cat sat on the mat"

自注意力:
"mat" 可以直接关注:
  - "The" ✓  （距离=5，但注意力权重可以很高）
  - "cat" ✓  （语义相关）
  - "sat" ✓
  - "on"  ✓
  - "the" ✓  （"the mat"是一个短语）

所有计算并行进行！

句子1中"银行"的注意力分布:
  - "河边" ← 高权重  （上下文线索）
  - "散步" ← 中等权重
  - "的"   ← 低权重
  → 模型推断："银行"指"河岸"

句子2中"银行"的注意力分布:
  - "取钱" ← 高权重  （上下文线索）
  - "去"   ← 中等权重
  → 模型推断："银行"指"金融机构"

你的需求（Query）: "深度学习教程"
书架上的书：
  - 书1（Key）: "深度学习入门"  → 相关度高 → 你会仔细阅读（Value权重高）
  - 书2（Key）: "Python编程"     → 相关度中 → 简单翻翻（Value权重中）
  - 书3（Key）: "古诗词鉴赏"     → 相关度低 → 不看（Value权重低）

句子: "The cat sat"

当处理"cat"时:
  Query_cat: "我是'cat'，我想知道哪些词与我相关"

  计算与所有词的相关性:
    相关性(Query_cat, Key_The) = 0.2
    相关性(Query_cat, Key_cat) = 1.0
    相关性(Query_cat, Key_sat) = 0.7  （主语和谓语相关）

  加权融合Value:
    Output_cat = 0.2 * Value_The + 1.0 * Value_cat + 0.7 * Value_sat

搜索引擎场景：
- 用户输入（Query）："好吃的川菜"
- 餐馆标签（Key）："火锅"、"串串"、"麻辣烫"
- 餐馆详情（Value）：地址、菜单、评分

如果Query = Key：
用户必须输入"火锅"才能找到"火锅"
→ 无法语义匹配（"好吃的川菜"匹配不到"火锅"）

# 输入嵌入（同一个词"bank"）
X_bank = [0.2, 0.5, 0.8, ...]  # 768维

# 场景1："river bank"
# Query空间（查询上下文）
Q_bank = X_bank @ W_Q  # → [位置信息, 地理特征, ...]
# Key空间（提供位置信息）
K_river = X_river @ W_K  # → [水体特征, 地理相关, ...]
# 注意力：Q_bank · K_river 高分 → 关注"river"

# 场景2："bank account"
# Query空间（查询金融信息）
Q_bank = X_bank @ W_Q  # → [金融特征, 账户相关, ...]
# Key空间（提供金融信息）
K_account = X_account @ W_K  # → [金融特征, 数字相关, ...]
# 注意力：Q_bank · K_account 高分 → 关注"account"

Key匹配阶段（Q·K）：
  判断"cat"和"猫"语义相关（高分）

Value提取阶段（Attention·V）：
  提取"猫"的【翻译】信息：
    - V可能编码：发音"māo"、字形、语法属性
    - 而K只编码：语义相似度特征

如果V=K：
  V被迫同时承担"匹配"和"内容"双重职责
  → 表达能力受限

单矩阵（Q=K=V=XW）：
  所有信息压缩到同一个64维子空间
  [←────────64维────────→]

三矩阵（独立）：
  信息分布在三个可能正交的子空间
  Q: [←────64维────→]
  K:          [←────64维────→]
  V:                   [←────64维────→]
  总容量: 最多192维

输入X（512维）
  ↓
分流到三个独立空间：
  ├─ W^Q → 查询特征（64维）
  ├─ W^K → 键特征（64维）
  └─ W^V → 值特征（64维）
  ↓
注意力机制组合（Query·Key匹配 + Value提取）
  ↓
输出（512维，包含X的多视角信息）

场景：在图书馆找书
Query：我的目标是什么？
  → "找一本关于深度学习的书"

Key：书架上每本书的"标签"
  → "Python编程"（不相关）
  → "深度学习入门"（高度相关！）
  → "机器学习基础"（中度相关）

Value：不是书的"标签"，而是书的"内容"
  → 提取："反向传播算法"、"神经网络架构"等知识

# 配置1：标准三矩阵（基线）
class StandardAttention(nn.Module):
    def __init__(self, d_model, d_k):
        self.W_q = nn.Linear(d_model, d_k)  # 独立
        self.W_k = nn.Linear(d_model, d_k)  # 独立
        self.W_v = nn.Linear(d_model, d_k)  # 独立

# 配置2：V=K（共享值和键）
class SharedKV(nn.Module):
    def __init__(self, d_model, d_k):
        self.W_q = nn.Linear(d_model, d_k)
        self.W_kv = nn.Linear(d_model, d_k)  # 共享
    def forward(self, x):
        q = self.W_q(x)
        k = v = self.W_kv(x)  # K和V相同

# 配置3：Q=K（共享查询和键）
class SharedQK(nn.Module):
    def __init__(self, d_model, d_k):
        self.W_qk = nn.Linear(d_model, d_k)  # 共享
        self.W_v = nn.Linear(d_model, d_k)
    def forward(self, x):
        q = k = self.W_qk(x)  # Q和K相同
        v = self.W_v(x)

# 配置4：Q=K=V=X（无变换）
class NoProjection(nn.Module):
    def forward(self, x):
        q = k = v = x  # 全部相同，无学习参数

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

class SelfAttention(nn.Module):
    """
    自注意力模块
    """
    def __init__(self, d_model, d_k):
        """
        Args:
            d_model: 输入嵌入维度
            d_k: Query和Key的维度
        """
        super().__init__()
        self.d_k = d_k

        # Q、K、V的线性变换
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_k, bias=False)

    def forward(self, x, mask=None):
        """
        Args:
            x: [batch_size, seq_len, d_model]
            mask: [batch_size, seq_len, seq_len] 可选掩码

        Returns:
            output: [batch_size, seq_len, d_k]
            attention_weights: [batch_size, seq_len, seq_len]
        """
        # 步骤1: 计算Q、K、V
        Q = self.W_q(x)  # [batch, seq_len, d_k]
        K = self.W_k(x)  # [batch, seq_len, d_k]
        V = self.W_v(x)  # [batch, seq_len, d_k]

        # 步骤2: 计算注意力分数（QK^T）
        scores = torch.matmul(Q, K.transpose(-2, -1))  # [batch, seq_len, seq_len]

        # 步骤3: 缩放
        scores = scores / math.sqrt(self.d_k)

        # 步骤4: 应用掩码（可选）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 步骤5: Softmax
        attention_weights = F.softmax(scores, dim=-1)  # [batch, seq_len, seq_len]

        # 步骤6: 加权求和Value
        output = torch.matmul(attention_weights, V)  # [batch, seq_len, d_k]

        return output, attention_weights


# 测试
batch_size = 2
seq_len = 5
d_model = 512
d_k = 64

# 随机输入
x = torch.randn(batch_size, seq_len, d_model)

# 创建模块
attention = SelfAttention(d_model, d_k)

# 前向传播
output, weights = attention(x)

print(f"输入形状: {x.shape}")
print(f"输出形状: {output.shape}")
print(f"注意力权重形状: {weights.shape}")

# 查看第一个样本的注意力权重
print("\n第一个样本的注意力权重矩阵:")
print(weights[0])
print("\n每行的和（应该都是1.0）:")
print(weights[0].sum(dim=-1))

输入形状: torch.Size([2, 5, 512])
输出形状: torch.Size([2, 5, 64])
注意力权重形状: torch.Size([2, 5, 5])

第一个样本的注意力权重矩阵:
tensor([[0.1823, 0.2154, 0.1932, 0.2011, 0.2080],
        [0.2234, 0.1876, 0.1943, 0.2001, 0.1946],
        [0.1987, 0.2123, 0.1854, 0.2067, 0.1969],
        [0.2056, 0.1932, 0.2098, 0.1876, 0.2038],
        [0.1943, 0.2011, 0.2087, 0.1989, 0.1970]], grad_fn=<SelectBackward0>)

每行的和（应该都是1.0）:
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000], grad_fn=<SumBackward1>)

样本1: "Hello world"         → 长度=2
样本2: "I love AI"            → 长度=3
样本3: "Transformers are great" → 长度=3

样本1: "Hello world <PAD>"
样本2: "I love AI"
样本3: "Transformers are great"

生成"The cat sat":
  - "The" 只能看 "The"
  - "cat" 只能看 "The", "cat"
  - "sat" 只能看 "The", "cat", "sat"

import torch
import torch.nn.functional as F

def create_padding_mask(seq_len, valid_len):
    """
    创建填充掩码

    Args:
        seq_len: 序列总长度
        valid_len: 有效长度（非填充部分）

    Returns:
        mask: [seq_len, seq_len]，有效位置为1，填充位置为0
    """
    # 创建位置索引
    positions = torch.arange(seq_len).unsqueeze(0)  # [1, seq_len]

    # 创建掩码：位置 < valid_len 的为True
    mask = positions < valid_len  # [1, seq_len]

    # 扩展到 [seq_len, seq_len]（每行相同）
    mask = mask.unsqueeze(0).expand(seq_len, -1)

    return mask.float()


# 示例：序列长度=5，有效长度=3
mask = create_padding_mask(seq_len=5, valid_len=3)
print("填充掩码:")
print(mask)

填充掩码:
tensor([[1., 1., 1., 0., 0.],
        [1., 1., 1., 0., 0.],
        [1., 1., 1., 0., 0.],
        [1., 1., 1., 0., 0.],
        [1., 1., 1., 0., 0.]])

def apply_mask(scores, mask):
    """
    应用掩码到注意力分数

    Args:
        scores: [batch, seq_len, seq_len] 注意力分数
        mask: [seq_len, seq_len] 掩码

    Returns:
        masked_scores: 掩码后的分数
    """
    # 将mask=0的位置设为-1e9（近似-∞）
    return scores.masked_fill(mask == 0, -1e9)


# 示例
scores = torch.randn(1, 5, 5) * 2  # 随机注意力分数
print("原始分数:\n", scores[0])

masked_scores = apply_mask(scores, mask.unsqueeze(0))
print("\n掩码后分数:\n", masked_scores[0])

# Softmax后
attn_weights = F.softmax(masked_scores, dim=-1)
print("\nSoftmax后注意力权重:\n", attn_weights[0])

原始分数:
tensor([[ 1.2, -0.5,  0.8,  1.1, -0.3],
        [ 0.6,  1.3, -0.7,  0.9,  1.5],
        ...])

掩码后分数:
tensor([[ 1.2000e+00, -5.0000e-01,  8.0000e-01, -1.0000e+09, -1.0000e+09],
        [ 6.0000e-01,  1.3000e+00, -7.0000e-01, -1.0000e+09, -1.0000e+09],
        ...])

Softmax后注意力权重:
tensor([[0.4234, 0.0781, 0.2985, 0.0000, 0.0000],  ← 填充位置权重=0
        [0.2123, 0.4234, 0.0643, 0.0000, 0.0000],
        ...])

def create_causal_mask(seq_len):
    """
    创建因果掩码（下三角矩阵）

    Args:
        seq_len: 序列长度

    Returns:
        mask: [seq_len, seq_len]
    """
    # 创建下三角矩阵
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask


# 示例
causal_mask = create_causal_mask(5)
print("因果掩码（下三角）:")
print(causal_mask)

因果掩码（下三角）:
tensor([[1., 0., 0., 0., 0.],  ← 位置0只能看自己
        [1., 1., 0., 0., 0.],  ← 位置1能看0和1
        [1., 1., 1., 0., 0.],  ← 位置2能看0、1、2
        [1., 1., 1., 1., 0.],
        [1., 1., 1., 1., 1.]]) ← 位置4能看所有

import matplotlib.pyplot as plt
import seaborn as sns

# 模拟注意力分数
scores = torch.randn(5, 5)

# 应用因果掩码
masked_scores = scores.masked_fill(causal_mask == 0, -1e9)
attn_weights = F.softmax(masked_scores, dim=-1)

# 可视化
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# 左图：原始分数
sns.heatmap(scores.numpy(), annot=True, fmt=".2f", cmap="RdBu",
            center=0, ax=axes[0], cbar_kws={'label': '分数'})
axes[0].set_title("原始注意力分数")
axes[0].set_xlabel("Key位置")
axes[0].set_ylabel("Query位置")

# 右图：掩码后的注意力权重
sns.heatmap(attn_weights.numpy(), annot=True, fmt=".2f", cmap="YlOrRd",
            ax=axes[1], cbar_kws={'label': '权重'})
axes[1].set_title("应用因果掩码后的注意力权重")
axes[1].set_xlabel("Key位置")
axes[1].set_ylabel("Query位置")

plt.tight_layout()
plt.savefig('causal_mask_effect.png', dpi=300)
plt.show()

Encoder = 阅读理解：拿到完整文章，理解每个词的含义
Decoder = 写作文：只能看到已写的内容，预测下一个字

# 训练样本："我 爱 NLP"
# 目标：预测下一个词

# 位置0预测"爱"时
# 如果用双向注意力，模型能看到:
输入: [我, 爱, NLP]  # 完整句子
目标: 预测 "爱"

# 问题：模型已经看到答案"爱"了！
# 相当于开卷考试，模型会学会"抄答案"而不是真正学习语言模式

# 推理时生成句子
# 第1步：只有 [<BOS>]
# 第2步：只有 [<BOS>, 我]
# 第3步：只有 [<BOS>, 我, 爱]

# 如果训练时模型习惯看到完整句子（双向）
# 推理时只有部分句子 → 分布不匹配 → 性能崩溃

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 实验：双向 vs 单向 Attention
class BidirectionalGPT2(nn.Module):
    """错误示范：双向Decoder"""
    def __init__(self, config):
        super().__init__()
        self.transformer = GPT2LMHeadModel(config)

    def forward(self, input_ids):
        # 移除因果掩码（允许双向）
        # 注意：这是错误的！
        outputs = self.transformer(
            input_ids,
            use_cache=False,
            # 不使用 causal mask
        )
        return outputs


# 正确的单向Decoder
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model_causal = GPT2LMHeadModel.from_pretrained('gpt2')

# 测试句子
text = "I love natural language"
inputs = tokenizer(text, return_tensors='pt')

# 单向生成（正确）
with torch.no_grad():
    outputs_causal = model_causal.generate(
        inputs['input_ids'],
        max_length=10,
        do_sample=False
    )

print("单向Decoder生成:", tokenizer.decode(outputs_causal[0]))
# 输出: "I love natural language processing and machine learning"

# 如果用双向（训练-推理不匹配）
# 生成质量会严重下降，出现：
# - 重复token
# - 语义不连贯
# - 困惑度飙升

# 可视化每个位置的有效上下文长度
def analyze_causal_context(seq_len=10):
    """分析因果掩码下每个位置的信息量"""
    positions = list(range(1, seq_len + 1))
    context_sizes = positions  # 位置i能看到i个token

    import matplotlib.pyplot as plt
    plt.figure(figsize=(10, 6))
    plt.bar(positions, context_sizes, color='skyblue', edgecolor='black')
    plt.xlabel('位置', fontsize=12)
    plt.ylabel('可见上下文大小', fontsize=12)
    plt.title('因果掩码下各位置的信息量', fontsize=14)
    plt.axhline(y=seq_len/2, color='r', linestyle='--',
                label=f'平均上下文={seq_len/2}')
    plt.legend()
    plt.grid(axis='y', alpha=0.3)
    plt.savefig('causal_context_distribution.png', dpi=300)
    plt.show()

    # 统计
    avg_context = sum(context_sizes) / len(context_sizes)
    print(f"平均上下文大小: {avg_context:.1f} tokens")
    print(f"最小上下文: {min(context_sizes)} (位置1)")
    print(f"最大上下文: {max(context_sizes)} (位置{seq_len})")

analyze_causal_context(seq_len=10)

平均上下文大小: 5.5 tokens
最小上下文: 1 (位置1)
最大上下文: 10 (位置10)

def create_xlnet_mask(seq_len, perm):
    """
    XLNet的排列掩码

    Args:
        seq_len: 序列长度
        perm: 排列顺序，如 [2, 0, 3, 1]

    Returns:
        mask: [seq_len, seq_len]
    """
    mask = torch.zeros(seq_len, seq_len)
    for i, pos in enumerate(perm):
        # 位置pos能看到排列中它之前的所有位置
        for j in range(i):
            prev_pos = perm[j]
            mask[pos, prev_pos] = 1
    return mask

# 示例：序列长度4，排列 [2, 0, 3, 1]
perm = [2, 0, 3, 1]
xlnet_mask = create_xlnet_mask(4, perm)
print("XLNet排列掩码:")
print(xlnet_mask)
# 输出：
# tensor([[0., 0., 1., 0.],  ← 位置0能看位置2（排列中的前驱）
#         [1., 0., 1., 1.],  ← 位置1能看2, 0, 3（排列中的前驱）
#         [0., 0., 0., 0.],  ← 位置2第一个，看不到任何位置
#         [0., 0., 1., 1.]]) ← 位置3能看2, 0（排列中的前驱）

class DecoderLayer(nn.Module):
    def forward(self, x, memory, tgt_mask, memory_mask):
        # 1. 自注意力：因果掩码（单向）
        x = self.self_attn(
            query=x, key=x, value=x,
            attn_mask=tgt_mask  # 因果掩码
        )

        # 2. 交叉注意力：无掩码（双向）
        x = self.cross_attn(
            query=x,           # Decoder的隐状态
            key=memory,        # Encoder的输出（完整source）
            value=memory,
            attn_mask=None     # 无因果限制！
        )

        # 3. FFN
        x = self.ffn(x)
        return x

def create_combined_mask(seq_len, valid_len):
    """
    创建组合掩码（Padding + Causal）

    Args:
        seq_len: 序列总长度
        valid_len: 有效长度

    Returns:
        mask: [seq_len, seq_len]
    """
    # 因果掩码
    causal = create_causal_mask(seq_len)

    # 填充掩码
    padding = create_padding_mask(seq_len, valid_len)

    # 两者取交集（都为1才为1）
    combined = causal * padding

    return combined


# 示例：序列长度=5，有效长度=3
combined_mask = create_combined_mask(seq_len=5, valid_len=3)
print("组合掩码:")
print(combined_mask)

组合掩码:
tensor([[1., 0., 0., 0., 0.],  ← 位置0：只看自己，且自己有效
        [1., 1., 0., 0., 0.],  ← 位置1：能看0、1，且都有效
        [1., 1., 1., 0., 0.],  ← 位置2：能看0、1、2，且都有效
        [1., 1., 1., 0., 0.],  ← 位置3：因果允许看0-3，但3是填充
        [1., 1., 1., 0., 0.]]) ← 位置4：因果允许看0-4，但4是填充

# 测试掩码对梯度的影响
x = torch.randn(1, 5, 64, requires_grad=True)
attention = SelfAttention(d_model=64, d_k=64)

# 不使用掩码
output1, _ = attention(x, mask=None)
loss1 = output1.sum()
loss1.backward()
grad1 = x.grad.clone()
x.grad.zero_()

# 使用掩码
mask = create_causal_mask(5).unsqueeze(0)
output2, _ = attention(x, mask=mask)
loss2 = output2.sum()
loss2.backward()
grad2 = x.grad.clone()

print("梯度差异:")
print(f"不使用掩码的梯度范数: {grad1.norm():.4f}")
print(f"使用掩码的梯度范数: {grad2.norm():.4f}")
print(f"梯度是否相同: {torch.allclose(grad1, grad2)}")

from transformers import AutoTokenizer, AutoModel
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# 加载BERT模型
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True)

# 测试句子
sentence = "The cat sat on the mat"
inputs = tokenizer(sentence, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0])

print("Tokens:", tokens)

# 前向传播，获取注意力权重
with torch.no_grad():
    outputs = model(**inputs)
    # outputs.attentions: 12层，每层的注意力权重
    # 取第6层、第1个头的注意力
    attention = outputs.attentions[5][0, 0].numpy()  # [seq_len, seq_len]

# 可视化
plt.figure(figsize=(10, 8))
sns.heatmap(
    attention,
    xticklabels=tokens,
    yticklabels=tokens,
    cmap="YlOrRd",
    annot=True,
    fmt=".2f",
    cbar_kws={'label': '注意力权重'}
)
plt.xlabel("被关注的Token")
plt.ylabel("当前Token")
plt.title("BERT第6层第1头的注意力权重")
plt.tight_layout()
plt.savefig('attention_heatmap.png', dpi=300)
plt.show()

位置5的编码 → 通过线性变换 → 得到"位置5比位置2远3个位置"

# 不同维度的频率
dims = [0, 64, 128, 192, 255]
positions = range(100)

for dim in dims:
    freq = 1 / (10000 ** (dim / 256))
    values = [np.sin(pos * freq) for pos in positions]

    plt.plot(positions, values, label=f'维度{dim}')

plt.legend()
plt.title('不同维度的位置编码频率')

训练长度: 512
测试长度: 2048  → 性能下降(外推失败)

import torch
import numpy as np
import matplotlib.pyplot as plt

def get_positional_encoding(seq_len, d_model):
    """
    生成正弦余弦位置编码

    Args:
        seq_len: 序列长度
        d_model: 嵌入维度

    Returns:
        pos_encoding: [seq_len, d_model]
    """
    # 创建位置和维度的索引
    position = torch.arange(seq_len).unsqueeze(1)  # [seq_len, 1]
    div_term = torch.exp(
        torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model)
    )  # [d_model/2]

    # 初始化位置编码矩阵
    pos_encoding = torch.zeros(seq_len, d_model)

    # 偶数维度用sin
    pos_encoding[:, 0::2] = torch.sin(position * div_term)

    # 奇数维度用cos
    pos_encoding[:, 1::2] = torch.cos(position * div_term)

    return pos_encoding


# 生成位置编码
seq_len = 100
d_model = 512
pe = get_positional_encoding(seq_len, d_model)

print(f"位置编码形状: {pe.shape}")
print(f"位置0的编码（前10维）:\n{pe[0, :10]}")
print(f"位置1的编码（前10维）:\n{pe[1, :10]}")

# 可视化
plt.figure(figsize=(15, 5))

# 子图1：位置编码热力图
plt.subplot(1, 2, 1)
plt.imshow(pe.numpy(), cmap='RdBu', aspect='auto')
plt.xlabel('维度')
plt.ylabel('位置')
plt.title('位置编码可视化')
plt.colorbar()

# 子图2：几个位置的编码曲线
plt.subplot(1, 2, 2)
positions_to_plot = [0, 10, 20, 50]
for pos in positions_to_plot:
    plt.plot(pe[pos, :128].numpy(), label=f'位置 {pos}')
plt.xlabel('维度')
plt.ylabel('编码值')
plt.title('不同位置的编码曲线（前128维）')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.savefig('positional_encoding.png', dpi=300)
plt.show()

import torch

def compute_theta(dim: int, base: float = 10000.0) -> torch.Tensor:
    """计算角频率

    Args:
        dim: 注意力头维度（必须是偶数）
        base: 基数，通常为10000

    Returns:
        theta: [dim/2] 角频率向量
    """
    # θⱼ = 1 / (base^{2j/d})
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    return inv_freq

# 示例：64维注意力头
theta = compute_theta(64)
print(f"θ₀ = {theta[0]:.6f}")  # 高频：θ₀ = 1.000000
print(f"θ₃₁ = {theta[31]:.6f}") # 低频：θ₃₁ = 0.000100

class RotaryEmbedding(nn.Module):
    """RoPE位置编码（LLaMA/Qwen实现）"""

    def __init__(self, dim: int, base: float = 10000.0, max_seq_len: int = 2048):
        super().__init__()
        # 计算逆频率：1 / (base^{2i/d})
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # 预计算缓存（优化性能）
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len: int):
        """预计算cos和sin值"""
        # 位置索引：[0, 1, 2, ..., seq_len-1]
        t = torch.arange(seq_len, device=self.inv_freq.device).float()

        # 计算 m*θⱼ：[seq_len, dim/2]
        freqs = torch.outer(t, self.inv_freq)

        # 重复拼接（对应特征对的x和y分量使用相同角度）
        emb = torch.cat((freqs, freqs), dim=-1)  # [seq_len, dim]

        # 缓存cos和sin
        self.cos_cached = emb.cos()
        self.sin_cached = emb.sin()

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
        """
        Args:
            x: [batch, seq_len, num_heads, head_dim]
            position_ids: [batch, seq_len]

        Returns:
            cos, sin: [batch, seq_len, head_dim]
        """
        # 动态扩展缓存
        seq_len = position_ids.max() + 1
        if seq_len > self.cos_cached.shape[0]:
            self._build_cache(seq_len)

        # 根据position_ids索引
        cos = self.cos_cached[position_ids]
        sin = self.sin_cached[position_ids]

        return cos, sin


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """将后半部分移到前面并取负：[-x_{d/2:}, x_{:d/2}]

    对应复数乘法的虚部：(a+bi)*(cosθ+i·sinθ) 的交叉项
    """
    x1 = x[..., :x.shape[-1]//2]
    x2 = x[..., x.shape[-1]//2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor,
                         cos: torch.Tensor, sin: torch.Tensor):
    """应用RoPE旋转

    数学等价于：x * e^{imθ} = x * (cos(mθ) + i*sin(mθ))

    Args:
        q, k: [batch, seq_len, num_heads, head_dim]
        cos, sin: [batch, seq_len, head_dim]

    Returns:
        q_embed, k_embed: 旋转后的查询和键
    """
    # 广播维度匹配
    cos = cos.unsqueeze(2)  # [batch, seq_len, 1, head_dim]
    sin = sin.unsqueeze(2)

    # 公式：x*cos(mθ) + rotate_half(x)*sin(mθ)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)

    return q_embed, k_embed

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """预计算频率的复数指数形式（cis = cos + i*sin）

    Returns:
        freqs_cis: [end, dim/2] 复数张量
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[:(dim//2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()  # [end, dim/2]

    # 生成复数：e^{i*mθ} = cos(mθ) + i*sin(mθ)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis


def apply_rotary_emb(xq, xk, freqs_cis):
    """使用复数乘法应用旋转（更简洁但需要复数支持）"""
    # 重塑为复数形式：[..., d] -> [..., d/2] complex
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # 复数乘法实现旋转
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

def position_interpolation(position_ids, max_train_len, current_len):
    """位置插值

    Args:
        position_ids: [batch, seq_len] 原始位置索引
        max_train_len: 训练时最大长度（如2048）
        current_len: 当前序列长度（如4096）

    Returns:
        插值后的位置索引
    """
    scale = max_train_len / current_len
    return (position_ids.float() * scale).long()

def ntk_scaled_rope(base, scale_factor, dim):
    """NTK-aware缩放"""
    return base * (scale_factor ** (dim / (dim - 2)))

# 示例：扩展2倍长度
base_new = ntk_scaled_rope(10000, 2.0, 128)  # ~40000

emb = torch.cat((freqs, freqs), dim=-1)

位置:  0    1  2  3  4 5  6  7
Token: 小明  很 聪明 ， 他 考了 满 分

头5的注意力权重:
"他"(位置4) 对各位置的注意力:
  小明: 0.85  ← 强关联！
  很:   0.02
  聪明: 0.05
  ，:   0.01
  他:   0.03
  考了: 0.02
  满:   0.01
  分:   0.01

原始空间(512维)
        ↓
头1: W₁^Q投影 → 子空间1(64维)  [学语法]
头2: W₂^Q投影 → 子空间2(64维)  [学语义]
头3: W₃^Q投影 → 子空间3(64维)  [学位置]
...

     小狗  追逐  小猫
小狗  0.1  0.8   0.1   ← "小狗"强关注"追逐"(主谓关系)
追逐  0.4  0.1   0.5   ← "追逐"关注主语和宾语
小猫  0.1  0.8   0.1   ← "小猫"强关注"追逐"(动宾关系)

     小狗  追逐  小猫
小狗  0.2  0.1   0.7   ← "小狗"关注"小猫"(语义相关:都是动物)
追逐  0.3  0.4   0.3
小猫  0.7  0.1   0.2   ← "小猫"关注"小狗"

# 单头注意力矩阵示例
A_single = [
    [0.7, 0.2, 0.05, 0.05],  # 第1个token
    [0.1, 0.8, 0.05, 0.05],  # 第2个token
    [0.1, 0.1, 0.7,  0.1 ],  # 第3个token
    [0.1, 0.1, 0.1,  0.7 ]   # 第4个token
]
# 每行和=1（Softmax约束）
# 实际秩：rank(A) ≈ 2-3（远小于理论上限4）

单头注意力（低秩）：
所有信息压缩到一个低维流形
[██████░░░░░░░░] rank ≈ 8-12 (远小于序列长度)

多头注意力（高秩）：
不同头覆盖不同子空间，总体接近满秩
头1: [██████░░░░░░░░] 语法子空间
头2: [░░░░██████░░░░] 语义子空间
头3: [░░░░░░░░██████] 位置子空间
...
总计: [██████████████] rank ≈ 60-80 (接近满秩)

import torch
import torch.nn.functional as F

def compute_attention_rank(n_tokens=128, d_k=64, n_heads=1):
    """计算注意力矩阵的实际秩"""
    # 模拟Q, K
    Q = torch.randn(1, n_heads, n_tokens, d_k)
    K = torch.randn(1, n_heads, n_tokens, d_k)

    # 计算注意力权重
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    attn = F.softmax(scores, dim=-1)  # [1, n_heads, n_tokens, n_tokens]

    # 计算每个头的秩（使用SVD）
    ranks = []
    for i in range(n_heads):
        A = attn[0, i].detach()
        # 计算秩（奇异值>1e-5的数量）
        s = torch.linalg.svdvals(A)
        rank = (s > 1e-5).sum().item()
        ranks.append(rank)

    return ranks, attn

# 实验1：单头
ranks_1, _ = compute_attention_rank(n_tokens=128, d_k=64, n_heads=1)
print(f"单头秩: {ranks_1[0]}/128")  # 输出: 约40-60 (远小于128)

# 实验2：8头
ranks_8, _ = compute_attention_rank(n_tokens=128, d_k=64, n_heads=8)
print(f"8头秩: {sum(ranks_8)}/128")  # 输出: 约100-120 (接近128)

# 理论验证
print(f"\n理论上限:")
print(f"  单头: min(n-1, d_k) = min(127, 64) = 64")
print(f"  8头: min(n, 8×平均秩) ≈ min(128, 8×50) = 128")

单头秩: 54/128  ← Softmax瓶颈导致低秩
8头秩: 115/128  ← Multi-Head恢复接近满秩

理论上限:
  单头: min(n-1, d_k) = min(127, 64) = 64
  8头: min(n, 8×平均秩) ≈ min(128, 8×50) = 128

头1, 头2, 头3, 头4 → 共享 K₁, V₁
头5, 头6, 头7, 头8 → 共享 K₂, V₂

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

class GroupedQueryAttention(nn.Module):
    """
    分组查询注意力（GQA）
    """
    def __init__(self, d_model, num_heads, num_kv_groups):
        """
        Args:
            d_model: 模型维度
            num_heads: Query头数
            num_kv_groups: KV分组数（GQA的核心参数）
                           - num_kv_groups=num_heads → 标准MHA
                           - num_kv_groups=1 → MQA
                           - 1 < num_kv_groups < num_heads → GQA
        """
        super().__init__()
        assert num_heads % num_kv_groups == 0, "num_heads必须能被num_kv_groups整除"

        self.d_model = d_model
        self.num_heads = num_heads
        self.num_kv_groups = num_kv_groups
        self.num_heads_per_group = num_heads // num_kv_groups
        self.head_dim = d_model // num_heads

        # Q投影：每个头都有独立的Q
        self.W_q = nn.Linear(d_model, num_heads * self.head_dim, bias=False)

        # K、V投影：每个组共享K和V
        self.W_k = nn.Linear(d_model, num_kv_groups * self.head_dim, bias=False)
        self.W_v = nn.Linear(d_model, num_kv_groups * self.head_dim, bias=False)

        # 输出投影
        self.W_o = nn.Linear(num_heads * self.head_dim, d_model, bias=False)

    def forward(self, x, mask=None):
        """
        Args:
            x: [batch_size, seq_len, d_model]
            mask: [batch_size, seq_len, seq_len]

        Returns:
            output: [batch_size, seq_len, d_model]
        """
        batch_size, seq_len, _ = x.shape

        # 计算Q、K、V
        Q = self.W_q(x)  # [batch, seq_len, num_heads * head_dim]
        K = self.W_k(x)  # [batch, seq_len, num_kv_groups * head_dim]
        V = self.W_v(x)  # [batch, seq_len, num_kv_groups * head_dim]

        # 重塑Q: [batch, num_heads, seq_len, head_dim]
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 重塑K、V: [batch, num_kv_groups, seq_len, head_dim]
        K = K.view(batch_size, seq_len, self.num_kv_groups, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_kv_groups, self.head_dim).transpose(1, 2)

        # 扩展K、V，让每组的K和V被多个Q头共享
        # [batch, num_kv_groups, seq_len, head_dim] → [batch, num_heads, seq_len, head_dim]
        K = K.repeat_interleave(self.num_heads_per_group, dim=1)
        V = V.repeat_interleave(self.num_heads_per_group, dim=1)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # 应用掩码
        if mask is not None:
            scores = scores.masked_fill(mask.unsqueeze(1) == 0, -1e9)

        # Softmax
        attn_weights = F.softmax(scores, dim=-1)

        # 加权求和
        attn_output = torch.matmul(attn_weights, V)  # [batch, num_heads, seq_len, head_dim]

        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous()  # [batch, seq_len, num_heads, head_dim]
        attn_output = attn_output.view(batch_size, seq_len, self.num_heads * self.head_dim)

        # 输出投影
        output = self.W_o(attn_output)

        return output


# 测试不同配置
batch_size = 2
seq_len = 10
d_model = 512
num_heads = 8

x = torch.randn(batch_size, seq_len, d_model)

# 配置1：标准MHA（num_kv_groups = num_heads）
mha = GroupedQueryAttention(d_model, num_heads, num_kv_groups=8)
out_mha = mha(x)
print(f"MHA输出形状: {out_mha.shape}")

# 配置2：GQA（num_kv_groups = 2）
gqa = GroupedQueryAttention(d_model, num_heads, num_kv_groups=2)
out_gqa = gqa(x)
print(f"GQA输出形状: {out_gqa.shape}")

# 配置3：MQA（num_kv_groups = 1）
mqa = GroupedQueryAttention(d_model, num_heads, num_kv_groups=1)
out_mqa = mqa(x)
print(f"MQA输出形状: {out_mqa.shape}")

# 参数量对比
def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

print(f"\n参数量对比:")
print(f"MHA: {count_parameters(mha):,}")
print(f"GQA: {count_parameters(gqa):,}")
print(f"MQA: {count_parameters(mqa):,}")