注意残差（Attention Residuals）

Paper | arXiv | Overview | Results | Citation

(a) 標準残差（Standard residuals）＋均一な加算集計
(b) Full AttnRes：各層がすべての前出力を注目（attend）する
(c) Block AttnRes：レイヤーをブロックに分割し、メモリ使用量を O(L d) から O(N d) に削減

このリポジトリについて

Attention Residuals (AttnRes) は、Transformer の標準残差接続の代替として設計された「ドロップイン」モジュールです。
各層が学習済みで入力依存的な注意機構（attention over depth）を通じて以前の表現を選択的に集約できるようになります。

概要

標準残差接続は、すべてのレイヤー出力を固定係数（単位重み）で加算します。
層が深くなるにつれてこの均一な集計は各レイヤーの寄与を希薄化し、隠れ状態の大きさが無制限に増大するという PreNorm でよく知られる問題を引き起こします。

AttnRes はこの固定集計を、前層出力に対する softmax 注意機構へ置き換えます：

[ \mathbf{h}l ;=;\sum{i=0}^{l-1}\alpha_{i\to l},\mathbf{v}_i ]

ここで重み (\alpha_{i\to l}) は各層ごとに学習された擬似クエリ (\mathbf{w}_l\in\mathbb{R}^d) から計算されます。
これにより、すべての以前の表現に対して内容依存型（content‑aware）で選択的にアクセスできるようになります。

Block AttnRes

Full AttnRes は直感的ですが、大規模化すると O(L d) のメモリを必要とします。
Block AttnRes ではレイヤーを (N) 個のブロックに分割し、各ブロック内は標準残差で集約、注意機構はブロックレベル表現（block‑level representations）のみで行います。
約8つのブロックであれば Full AttnRes のほぼ全ての利点を回復しつつ、実用的なドロップイン置換として余分なオーバーヘッドを最小限に抑えられます。

PyTorch スタイルの擬似コード

def block_attn_res(blocks: list[Tensor], partial_block: Tensor,
                   proj: Linear, norm: RMSNorm) -> Tensor:
    """
    ブロック間注意：ブロック表現 + 部分和を注目する。
    blocks : N 個の [B, T, D] テンソル – 以前のブロックの完成済み表現
    partial_block : [B, T, D] – 現在のブロック内で途中まで集計したもの (b_n^i)
    """
    V = torch.stack(blocks + [partial_block])          # [N+1, B, T, D]
    K = norm(V)
    logits = torch.einsum('d, n b t d -> n b t', 
                          proj.weight.squeeze(), K)
    h = torch.einsum('n b t, n b t d -> b t d',
                     logits.softmax(0), V)
    return h


def forward(self, blocks: list[Tensor], hidden_states: Tensor) \
        -> tuple[list[Tensor], Tensor]:
    partial_block = hidden_states

    # 自注意機構前のブロック残差
    h = block_attn_res(blocks, partial_block,
                       self.attn_res_proj, self.attn_res_norm)

    # ブロック境界で新しいブロックを開始
    if self.layer_number % (self.block_size // 2) == 0:
        blocks.append(partial_block)
        partial_block = None

    # 自注意層
    attn_out = self.attn(self.attn_norm(h))
    partial_block = partial_block + attn_out if partial_block is not None \
                                            else attn_out

    # MLP 前のブロック残差
    h = block_attn_res(blocks, partial_block,
                       self.mlp_res_proj, self.mlp_res_norm)

    # MLP 層
    mlp_out = self.mlp(self.mlp_norm(h))
    partial_block = partial_block + mlp_out

    return blocks, partial_block

実験結果

スケーリング法則

AttnRes はすべての計算予算においてベースラインよりも一貫して優れた性能を示します。
Block AttnRes は 1.25 倍 の計算量で訓練されたベースラインと同等の損失（loss）を達成します。

下流タスクでの性能

(48B / 3B モデル、1.4T トークン)

AttnRes は全タスクで改善を示し、特に多段階推論（GPQA‑Diamond）で +7.5、コード生成（HumanEval）で +3.1 の最大ゲインを達成しています。

訓練ダイナミクス

AttnRes は PreNorm による希薄化問題を軽減し、層深さにわたって出力の大きさが安定し、勾配ノルムも層間でより均一になります。

引用

お役立ていただけましたら、以下の BibTeX をご利用ください。

@misc{chen2026attnres,
  title         = {Attention Residuals},
  author        = {Kimi Team and Chen, Guangyu and Zhang, Yu and Su, Jianlin and Xu, Weixin and Pan, Siyuan and Wang, Yaoyu and Wang, Yucheng and Chen, Guanduo and Yin, Bohong and Chen, Yutian and Yan, Junjie and Wei, Ming and Zhang, Y. and Meng, Fanqing and Hong, Chao and Xie, Xiaotong and Liu, Shaowei and Lu, Enzhe and Tai, Yunpeng and Chen, Yanru and Men, Xin and Guo, Haiqing and Charles, Y. and Lu, Haoyu and Sui, Lin and Zhu, Jinguo and Zhou, Zaida and He, Weiran and Huang, Weixiao and Xu, Xinran and Wang, Yuzhi and Lai, Guokun and Du, Yulun and Wu, Yuxin and Yang, Zhilin and Zhou, Xinyu},
  year          = {2026},
  archiveprefix = {arXiv},
  eprint        = {2603.15031},
  primaryclass  = {cs.CL}
}