"""
=============================================================================
第 4 讲 · 代码 2：把状态看成"联想记忆" —— 线性注意力、Delta 规则、门控
=============================================================================

对应讲义：lectures/lecture04_attention_alternatives_moe.md 第 3.3、3.4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
线性注意力的状态 S（一个 d_k × d_v 矩阵）可以看成一个**联想记忆**（associative memory）：
    写入：S ← S + k v^T          （把"键 k → 值 v"这条关联存进去）
    读取：v̂ = S^T q               （用查询 q 去取；q 等于某个存过的键时，希望取回对应的值）
这个视角能直观地解释线性注意力的弱点，以及后来的改进为什么有效：
    1. 容量：d_k 维的键最多只能"互不干扰"地存 d_k 条关联。存多了，读出来就是一团混合物。
    2. 更新：同一个键再写一次新值，纯累加会把新旧值混在一起，而不是覆盖。
       → **Delta 规则**（DeltaNet）：先读出这个键当前存的值，只写入"差值"：
             S ← S + β k (v - S^T k)^T   ⟺   S ← (I - β k k^T) S + β k v^T
    3. 遗忘：信息太旧就该淡出。→ **门控**（γ_t）：S ← γ_t S + ...
    4. Gated DeltaNet（Qwen3-Next / Qwen3.5 使用）把 2 和 3 结合起来：
             S ← γ_t (I - β_t k k^T) S + β_t k v^T
对比参照组：softmax 注意力（KV 缓存）把每一条都原样存下来，读取很准，但内存随 token 数线性增长。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_associative_memory.py
"""

import torch
import torch.nn.functional as F


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def random_unit_vectors(n, d, generator):
    """生成 n 个 d 维的随机单位向量（作为"键"）。"""
    x = torch.randn(n, d, generator=generator)
    return x / x.norm(dim=-1, keepdim=True)


def cosine(a, b):
    """逐行余弦相似度的平均值：1 表示取回的值完全正确，0 表示毫不相关。"""
    return F.cosine_similarity(a, b, dim=-1).mean().item()


# =============================================================================
# 三种"写入"规则
# =============================================================================
def write_additive(S, k, v):
    """线性注意力：S ← S + k v^T。"""
    return S + torch.outer(k, v)


def write_delta(S, k, v, beta=1.0):
    """
    Delta 规则：S ← S + β k (v - S^T k)^T。

    【直觉】先用 k 读出当前存的值 S^T k，与想存的 v 比较，只把"差值"写进去。
    当 k 是单位向量、β = 1 时，写完之后 S^T k 恰好等于 v：旧值被**覆盖**，而不是混合。
    （这正是经典的"最小均方误差"在线学习规则，也可以看成对 ||S^T k - v||² 做一步梯度下降。）
    【参数】beta ∈ [0, 1]：写入强度；β = 0 表示这一步什么也不写（Gated DeltaNet 的"不写入"门）。
    """
    return S + beta * torch.outer(k, v - S.T @ k)


def write_gated_delta(S, k, v, gamma, beta=1.0):
    """Gated DeltaNet：S ← γ (I - β k k^T) S + β k v^T，等价于先整体衰减 γ，再做 Delta 写入。"""
    return write_delta(gamma * S, k, v, beta)


def read(S, q):
    """读取：v̂ = S^T q。"""
    return S.T @ q


# =============================================================================
# 实验 1：容量
# =============================================================================
def demo_capacity():
    """演示 1：存的关联越多，读出来越不准（d_k = 64）。"""
    print_header("演示 1：容量 —— 固定大小的记忆能存多少条关联？")
    g = torch.Generator().manual_seed(0)
    d = 64
    print(f"  d_k = d_v = {d}；存入 N 条（随机键 → 随机值），再用每个键把值读出来，看余弦相似度：")
    print(f"  {'N':>6}{'线性注意力（累加）':>20}{'Delta 规则':>14}{'softmax 注意力（KV 缓存）':>26}")
    for N in [8, 32, 64, 128, 512]:
        keys = random_unit_vectors(N, d, g)
        vals = torch.randn(N, d, generator=g)
        S_add = torch.zeros(d, d)
        S_del = torch.zeros(d, d)
        for i in range(N):
            S_add = write_additive(S_add, keys[i], vals[i])
            S_del = write_delta(S_del, keys[i], vals[i])
        r_add = torch.stack([read(S_add, keys[i]) for i in range(N)])
        r_del = torch.stack([read(S_del, keys[i]) for i in range(N)])
        # softmax 注意力：保存全部键和值，用"尖锐"的 softmax 读取（相当于找最相似的键）
        attn = torch.softmax(keys @ keys.T * 20.0, dim=-1)
        r_soft = attn @ vals
        print(f"  {N:>6}{cosine(r_add, vals):>20.3f}{cosine(r_del, vals):>14.3f}{cosine(r_soft, vals):>26.3f}")
    print("""
  - 线性注意力：读出的是"目标值 + 所有其他值的干扰"。键越多、干扰越大；N 远大于 d 时几乎读不出来。
  - Delta 规则：**平均**准确度并不比累加更高（N 很大时甚至更低）：N 超过 d 之后，
    新写入会覆盖（挤掉）旧的关联，所以早期的关联几乎读不出来。它的优势不在总容量，
    而在于"更新"（演示 2）和"让最近的记忆保持清晰"（演示 3）。
  - softmax 注意力：原样保存每一条，读取几乎完美，代价是内存随 N 线性增长（KV 缓存）。
  这就是"固定状态"模型与"注意力"模型的根本权衡：**记忆容量 vs 内存/计算成本**。""")


# =============================================================================
# 实验 2：更新同一个键
# =============================================================================
def demo_update():
    """演示 2：同一个键写入新值，累加会混合，Delta 规则会覆盖。"""
    print_header("演示 2：更新 —— 同一个键写入新值")
    g = torch.Generator().manual_seed(1)
    d, N = 64, 16
    keys = random_unit_vectors(N, d, g)
    old = torch.randn(N, d, generator=g)
    new = torch.randn(N, d, generator=g)
    S_add, S_del = torch.zeros(d, d), torch.zeros(d, d)
    for i in range(N):                              # 先存旧值
        S_add = write_additive(S_add, keys[i], old[i])
        S_del = write_delta(S_del, keys[i], old[i])
    for i in range(N):                              # 再用同样的键存新值
        S_add = write_additive(S_add, keys[i], new[i])
        S_del = write_delta(S_del, keys[i], new[i])
    r_add = torch.stack([read(S_add, keys[i]) for i in range(N)])
    r_del = torch.stack([read(S_del, keys[i]) for i in range(N)])
    print(f"  {'':<22}{'与新值的相似度':>14}{'与旧值的相似度':>14}")
    print(f"  {'线性注意力（累加）':<18}{cosine(r_add, new):>16.3f}{cosine(r_add, old):>14.3f}")
    print(f"  {'Delta 规则':<20}{cosine(r_del, new):>16.3f}{cosine(r_del, old):>14.3f}")
    print("""
  累加规则读出的是"新值 + 旧值"的混合（与两者的相似度都约为 0.65）；Delta 规则基本用新值覆盖了旧值
  （与新值的相似度不是 1，是因为随机的键并不完全正交，写入别的键时会轻微干扰这一条）。
  对语言模型来说，这就像"变量被重新赋值"、"话题发生转移"：旧信息应该被替换，而不是叠加。""")


# =============================================================================
# 实验 3：门控 —— 让旧记忆淡出
# =============================================================================
def demo_gating():
    """演示 3：写入远超容量时，谁被保留下来？比较累加、门控累加、Delta、门控 Delta。"""
    print_header("演示 3：记忆满了之后 —— 门控与 Delta 规则")
    g = torch.Generator().manual_seed(2)
    d, N = 64, 256                                   # 写入 256 条，远超容量 64
    keys = random_unit_vectors(N, d, g)
    vals = torch.randn(N, d, generator=g)
    rules = [
        ("累加（线性注意力）", lambda S, k, v: write_additive(S, k, v)),
        ("门控累加，γ=0.95", lambda S, k, v: write_additive(0.95 * S, k, v)),
        ("Delta 规则", lambda S, k, v: write_delta(S, k, v)),
        ("门控 Delta，γ=0.9", lambda S, k, v: write_gated_delta(S, k, v, 0.9)),
    ]
    print(f"  依次写入 {N} 条关联后，按写入时间分组读取（余弦相似度）：")
    print(f"  {'':<20}{'最近 8 条':>10}{'倒数 33~64 条':>14}{'最早 64 条':>12}")
    for name, write in rules:
        S = torch.zeros(d, d)
        for i in range(N):
            S = write(S, keys[i], vals[i])
        sims = F.cosine_similarity(torch.stack([read(S, keys[i]) for i in range(N)]), vals, dim=-1)
        print(f"  {name:<20}{sims[-8:].mean().item():>10.3f}{sims[-64:-32].mean().item():>14.3f}{sims[:64].mean().item():>12.3f}")
    print("""
  - 累加：所有记忆都一样模糊（约 0.44）—— 256 条关联挤在 64 维里，互相干扰。
  - 门控累加：整体衰减让旧记忆淡出，干扰减少，最近的记忆清晰得多；代价是较早的记忆几乎全部丢失。
  - Delta 规则：只修改与当前键相关的部分，最近的记忆很清晰，中等距离的记忆也保留得比门控累加更多。
  - 门控 Delta：两种机制叠加。这里用的是固定的 γ，看不出它比 Delta 规则更好；
    真实模型中 γ_t、β_t 都由当前 token 计算出来，模型可以自己决定：遇到新文档时大幅遗忘，
    平时尽量保留，无关的 token 干脆不写入（β = 0）。这种"按内容决定"的灵活性才是门控的价值所在。""")


def main():
    demo_capacity()
    demo_update()
    demo_gating()
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()