"""
=============================================================================
第 12 讲 · 代码 5：叠加（superposition）——特征比维度多时，模型怎样把它们挤进去；稀疏自编码器怎样把它们分开
=============================================================================

对应讲义：lectures/lecture12_mechanistic_interpretability.md 第 6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
可解释性研究的一个大麻烦：大模型里的单个神经元常常是"多义的"（polysemantic）——同一个神经元对好几种毫不相干的东西都有反应。
Elhage et al. (2022) 提出一个解释：模型要表示的"特征"比它的维度多，于是把多个特征**叠加**在同一组维度里；
只要特征是稀疏的（大多数时候不出现），互相之间的干扰就可以容忍。

    演示 1：Elhage et al. 的玩具模型：20 个特征（重要性递减）、隐藏层只有 5 维，h = W x，x̂ = ReLU(Wᵀ h + b)。
            特征越稀疏，模型表示的特征越多（> 5 个就是叠加），特征之间的干扰也越大。
    演示 2：叠加时的几何：被表示的特征方向之间是什么关系？5 个隐藏维度里的每一维对应几个特征（"多义"）？偏置起什么作用？
    演示 3：稀疏自编码器（SAE）：只看隐藏层的 5 维激活，学一本"字典"（40 个方向），要求每次只用很少几个——
            它能找回原来的特征方向吗？（这是 2023 年以后可解释性研究的主要工具，第 12 讲 §7）

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_superposition_and_sae.py     # 约 15 秒；只需要 PyTorch
"""

import os
import sys

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
import tiny_gpt  # noqa: E402,F401  （导入时设置 torch.set_num_threads(1)）
from console import pad, print_header  # noqa: E402

N_FEAT, N_HID = 20, 5
IMPORTANCE = torch.tensor([0.7 ** i for i in range(N_FEAT)])     # 第 i 个特征的重要性（损失里的权重）
SPARSITIES = [0.0, 0.7, 0.9, 0.97, 0.99]


def sample(gen, S, batch=1024):
    """特征向量 x：每个特征以概率 S 为 0，否则在 [0, 1) 均匀取值。S 越大越稀疏。"""
    x = torch.rand(batch, N_FEAT, generator=gen)
    return x * (torch.rand(batch, N_FEAT, generator=gen) > S).float()


def train_toy(S, steps=8000, seed=0):
    """
    训练玩具模型：把 20 维的 x 压到 5 维 h = W x，再还原 x̂ = ReLU(Wᵀ h + b)，损失是按重要性加权的平方误差。

    【思路】如果只能表示 5 个特征，最好的选择是最重要的 5 个、互相正交——没有干扰。
            但如果特征很稀疏，两个特征很少同时出现，就可以让它们共用方向：偶尔同时出现时的干扰，
            由 ReLU 和偏置 b 压掉（偏置是负的，小的干扰被截成 0）。
    【返回】(W (5 × 20), b, 最后的损失)
    """
    torch.manual_seed(seed)
    W = torch.nn.Parameter(torch.randn(N_HID, N_FEAT) * 0.1)
    b = torch.nn.Parameter(torch.zeros(N_FEAT))
    opt = torch.optim.Adam([W, b], lr=1e-3)
    gen = torch.Generator().manual_seed(seed + 1)
    for _ in range(steps):
        x = sample(gen, S)
        x_hat = torch.relu(x @ W.T @ W + b)
        loss = (IMPORTANCE * (x - x_hat) ** 2).sum(1).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return W.detach(), b.detach(), loss.item()


def interference(W):
    """每个特征 i 受到的干扰 Σ_{j≠i} (Ŵ_i · W_j)²：其他特征在它的方向上的投影的平方和（Ŵ_i 是单位向量）。"""
    norms = W.norm(dim=0).clamp(min=1e-8)
    Wn = W / norms
    proj = (Wn.T @ W) ** 2                                         # proj[i, j] = (Ŵ_i · W_j)²
    return proj.sum(1) - proj.diagonal()


def demo1(models):
    print_header("演示 1：20 个特征挤进 5 维——特征越稀疏，被表示的特征越多")
    print(f"{pad('稀疏度 S', 10)}{pad('损失', 10, '>')}{pad('被表示的特征（‖W_i‖ > 0.5）', 30, '>')}{pad('平均干扰', 10, '>')}  各特征的 ‖W_i‖（按重要性从高到低）")
    for S in SPARSITIES:
        W, b, loss = models[S]
        norms = W.norm(dim=0)
        rep = norms > 0.5
        inter = interference(W)[rep].mean().item()
        print(f"{pad(S, 10)}{pad(f'{loss:.4f}', 10, '>')}{pad(int(rep.sum()), 30, '>')}{pad(f'{inter:.2f}', 10, '>')}  "
              + " ".join(f"{v:.2f}" for v in norms.tolist()))
    print("\n观察 Observations：")
    print("  1. 特征不稀疏时（S = 0，每个特征都总是出现），模型只表示最重要的 5 个，互相正交，干扰是 0——")
    print("     和主成分分析的结论一样：维度只有 5，就只留 5 个方向。")
    print("  2. 特征越稀疏，被表示的特征越多：S = 0.7 时 8 个，0.9 时 9 个，0.97 时 10 个，0.99 时 11 个（另有 2 个表示了一半）。")
    print("     多出来的特征和别的特征共用方向，干扰从 0 升到 1 左右，S = 0.99 时更大——模型用'偶尔的干扰'换'多表示几个特征'。")
    print("  3. 被表示的总是最重要的那些（‖W_i‖ 从前往后排），不重要的特征被完全放弃（‖W_i‖ ≈ 0）。")
    print("     Elhage et al. 在更极端的稀疏度下看到几乎所有特征都被挤进去；我们 8000 步的训练没有走那么远。")


def demo2(models):
    print_header("演示 2：叠加时的几何（S = 0.97）——特征方向之间的夹角，以及每个隐藏维度对应几个特征")
    W, b, _ = models[0.97]
    norms = W.norm(dim=0)
    feats = [i for i in range(N_FEAT) if norms[i] > 0.5]
    Wn = W[:, feats] / norms[feats]
    cos = Wn.T @ Wn
    print("被表示的特征两两之间的余弦（只列绝对值 > 0.3 的）：")
    pairs = [(feats[a], feats[c], cos[a, c].item()) for a in range(len(feats)) for c in range(a + 1, len(feats)) if abs(cos[a, c]) > 0.3]
    print("  " + "，".join(f"特征 {i} 与 {j}：{v:+.2f}" for i, j, v in pairs))
    print(f"\n每个隐藏维度上，权重绝对值 > 0.3 的特征（'这个神经元对哪些特征有反应'）：")
    for d in range(N_HID):
        hits = [i for i in feats if abs(W[d, i]) > 0.3]
        print(f"  隐藏维度 {d}：特征 {hits}")
    print(f"偏置 b（被表示的特征）：{[round(b[i].item(), 2) for i in feats]}")
    print("\n观察 Observations：")
    print("  1. 被表示的 10 个特征正好组成 5 对，每一对指向**相反的方向**（余弦 −1.00）：两个特征共用一条直线，")
    print("     一个占正方向、一个占负方向，10 个特征正好占满 5 维空间里的 5 条直线。Elhage et al. 把这种结构叫做'二角形'（digon），")
    print("     更稀疏时还会出现三角形、五边形等更复杂的排列。")
    print("  2. 每个隐藏维度都对 8–10 个特征有反应：如果把隐藏维度当成'神经元'去解释，它就是多义的——")
    print("     不是因为它'在做好几件事'，而是因为特征在叠加，没有和坐标轴对齐。")
    print("  3. 偏置几乎都是 0：在二角形里，伙伴特征出现时，在我的方向上的投影是**负的**（余弦 −1），")
    print("     ReLU 自己就把它截掉了，不需要负的偏置。这就是为什么只要两者很少同时出现，就几乎没有代价。")


def demo3_sae(models):
    print_header("演示 3：稀疏自编码器——只看 5 维的隐藏激活，能找回 10 个特征方向吗？（S = 0.97 的模型）")
    W, _, _ = models[0.97]
    gen = torch.Generator().manual_seed(9)
    n_lat, l1 = 40, 3e-3
    torch.manual_seed(1)
    W_dec = torch.nn.Parameter(torch.randn(N_HID, n_lat))
    W_enc = torch.nn.Parameter(torch.randn(n_lat, N_HID) * 0.1)
    b_enc = torch.nn.Parameter(torch.zeros(n_lat))
    b_dec = torch.nn.Parameter(torch.zeros(N_HID))
    opt = torch.optim.Adam([W_dec, W_enc, b_enc, b_dec], lr=3e-3)
    for _ in range(6000):
        with torch.no_grad():
            W_dec.div_(W_dec.norm(dim=0, keepdim=True))            # 字典里每个方向保持单位长度（否则 L1 惩罚可以靠缩放绕过）
        h = sample(gen, 0.97, 4096) @ W.T                           # 模型的隐藏激活（SAE 只看得到这个）
        z = torch.relu((h - b_dec) @ W_enc.T + b_enc)               # 稀疏的"特征激活"
        h_hat = z @ W_dec.T + b_dec
        loss = ((h - h_hat) ** 2).sum(1).mean() + l1 * z.abs().sum(1).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        W_dec.div_(W_dec.norm(dim=0, keepdim=True))
        norms = W.norm(dim=0)
        feats = [i for i in range(N_FEAT) if norms[i] > 0.5]
        best = [(W[:, i] / norms[i] @ W_dec).max().item() for i in feats]
        h = sample(gen, 0.97, 4096) @ W.T
        z = torch.relu((h - b_dec) @ W_enc.T + b_enc)
        unexplained = (((h - (z @ W_dec.T + b_dec)) ** 2).sum() / ((h - h.mean(0)) ** 2).sum()).item()
        active = (z > 1e-3).float().sum(1).mean().item()
        alive = int(((z > 1e-3).float().sum(0) > 0).sum())
    print(f"字典大小 {n_lat}，L1 系数 {l1}，训练 6000 步")
    print(f"每个真实特征方向与字典里最接近的方向的余弦：{[round(c, 3) for c in best]}")
    print(f"余弦 > 0.9 的（找回来了）：{sum(c > 0.9 for c in best)} / {len(feats)}")
    print(f"平均每个样本激活的字典方向：{active:.2f} 个（真实特征平均每个样本出现 {N_FEAT * 0.03:.2f} 个）；"
          f"曾经被激活过的字典方向 {alive} / {n_lat}；没被解释的方差比例 {unexplained:.1e}")
    print("\n观察 Observations：")
    print("  1. SAE 从来没见过真实的特征 x，只看到 5 维的 h；但它学到的字典里，每一个被模型表示的特征都有一个几乎完全对齐的方向")
    print("     （余弦 0.96–0.998）。叠加在一起的特征被重新'拆开'了。")
    print("  2. 关键是稀疏性约束：每个样本只用不到 1 个字典方向——和真实特征的稀疏程度相当。")
    print("     如果不要求稀疏，5 维的空间用任意 5 个方向就能完美重建，找不回任何有意义的东西。")
    print("  3. 这就是 Anthropic 2023 年 'Towards Monosemanticity' 的核心想法：在真实模型的激活上训练 SAE，")
    print("     得到比单个神经元更'单义'的特征。真实模型里我们不知道真实的特征是什么，只能靠人去看每个方向在什么输入上激活——")
    print("     玩具模型的好处是可以直接验证'找对了'。")


if __name__ == "__main__":
    models = {S: train_toy(S) for S in SPARSITIES}
    demo1(models)
    demo2(models)
    demo3_sae(models)