Attention Residuals

出处：Attention Residuals arXiv:2603.15031、官方实现 MoonshotAI/Attention-Residuals；K3 用 Block 版本。

一、伪代码：图中每个框对应一行

# 常驻状态（每 token 一份，与序列长度无关）
#   emb     : [d]         embedding，bank[0]，永远在栈底
#   bank    : [n][d]      已完成的块表示（K3：12 层一块，共 8 块）
#   partial : [d]         当前块的运行和，块边界处入栈

for layer in layers:
    src = [emb] + bank + [partial]         # ① stack：图中底部那一摞，n+1 份 [d]
    k   = rmsnorm(src)                     # ② RMSNorm：出 K

    logits = k @ w
    alpha  = softmax(logits)               # ④ softmax：在来源维上分配权重
    h      = src @ alpha                     

    if layer.idx % 12 == 0:                # 块边界：攒满 12 层的运行和入栈，开新块
        bank.append(partial)
        partial = zeros()

    partial += layer.attn(rmsnorm(h))      # 子层输出照旧等权累进当前块
    partial += layer.moe(...)              # MLP 前会再走一次上面的取回

二、为什么需要这些模块

模块	去掉/换掉会怎样	所以它解决什么
stack（留住历史）	标准残差只留一个 partial，历史一旦被加总就取不回来	让「更早的层」在后续层里仍然可寻址，而不是被稀释进一个和
RMSNorm（只给打分用）	不做归一化，模长大的层靠幅值就能占满 softmax	让权重由方向决定而非幅值；且混合仍用原值，保住各层真实尺度
W／伪查询 w_l	换成固定 1 → 退回标准残差；换成静态标量 → 就是 DenseFormer，几乎无收益	权重逐层、逐 token 自适应；w_l 不带 token 维，所以同块内可批量预算
softmax（在来源维）	用线性权重就不成其为「检索」，只算加权平均	制造竞争：想多读某一层，必须少读别层，权重总量守恒
einsum 取回	没有它，层只能看「上一层压缩过的状态」	这一层的输入 = 按需重组的历史，而不是唯一的累加和
块边界入栈（Block 版）	Full 版要留住全部 $L$ 层输出，每 token $O(Ld)$	状态从 $O(Ld)$ 降到 $O(Nd)$：块内仍等权，只在块间检索

一句话：朴素残差的问题是「历史只累加、不可检索」，这五步分别负责「留住历史、让权重可比、让权重可学、让权重竞争、把历史取回来」。

与 mHC 那类方法正交：mHC 加宽「当前状态」（多条残差流），AttnRes 检索「过去」（深度维 attention）。代价是每 token 多存 $N+1$ 份 $d$ 维向量、每层多一次 $n+1$ 路的 softmax，K3 上约 2% 推理开销。