Kimi K3のAttention Residualsを解説する

こんにちは、DeNAでデータサイエンティストをやっているまつけんです。

今回は、最近公開されましたKimi K3 Technical Reportで解説されているモデルの重要なコンポーネントのうち、

• Attention Residuals
• Kimi Delta Attention(KDA)

について、興味深い仕組みだと思ったので、これらの２つの仕組みについて直感的に理解できるように解説していきたいと思います。この記事はその第一弾としてFull Attention Residualsを扱います。
Technical Reportでは、まず基本形であるFull Attention Residualsが説明されています。実際のKimi K3では、メモリと通信量を削減したBlock Attention Residualsが採用されていますが、本記事では、仕組みの理解を主眼に置き、Full Attention Residualsを解説します。

まず、通常のResidualsをおさらいして、そこからFull Atttention Residualsがどのような特徴を持っているかを見ていくことにしましょう。

通常のResiduals

まず、Residual connectionを持つ各サブレイヤーを次のように表します。

\begin{aligned}
h_{l+1} &= h_{l} + f_{l}(h_{l}) 
\end{aligned}

よくあるResidual Layerです。

1層目では (l=1) なので、

h_2 = h_1 + f_1(h_1) \cdots (1)

となります。Token Embeddingである$h_1$に、ある関数を通した$h_1$の値$f_1(h_1)$を足し合わせて次に送っていますね。

2層目では、

h_3 = h_2 + f_2(h_2)

となります。ここに、1層目で得られた $(1)$式を代入します。

\begin{aligned} 
h_3 &= h_2 + f_2(h_2) \\ 
&= \left(h_1 + f_1(h_1)\right) + f_2(h_2) \\
&= h_1 + f_1(h_1) + f_2(h_2) 
\end{aligned}

同様に3層目では、

\begin{aligned}
h_4
&= h_3 + f_3(h_3) \\
&= \left(
h_1 + f_1(h_1) + f_2(h_2)
\right)
+ f_3(h_3) \\
&= h_1
+ f_1(h_1)
+ f_2(h_2)
+ f_3(h_3)
\end{aligned}

同じ操作を $(l-1)$層目まで繰り返すと、

\begin{aligned}
h_l
&= h_{l-1} + f_{l-1}(h_{l-1}) \\
&= h_{l-2}
+ f_{l-2}(h_{l-2})
+ f_{l-1}(h_{l-1}) \\
&= h_{l-3}
+ f_{l-3}(h_{l-3})
+ f_{l-2}(h_{l-2})
+ f_{l-1}(h_{l-1}) \\
&\quad \vdots \\
&= h_1
+ f_1(h_1)
+ f_2(h_2)
+ \cdots
+ f_{l-1}(h_{l-1})
\end{aligned}

したがって、総和記号を使うと、

\boxed{
h_l
=
h_1
+
\sum_{i=1}^{l-1} f_i(h_i)
}

と書けます。

図で表すとこういった関係です。¹

Full Attention Residuals

通常のResidualを踏まえてFull Attention Residualsです。Technical Reportでは

\begin{aligned}
\boldsymbol{q}_l &= \boldsymbol{w}_l\in\mathbb{R}^{d},\\
\boldsymbol{k}_i=\boldsymbol{v}_i&=
\begin{cases}
\boldsymbol{h}_1, & i=0,\\
f_i(\boldsymbol{h}_i), & 1\le i\le l-1.
\end{cases}
\end{aligned}\tag{8}

\alpha_{i\to l}
=\frac{\phi(\boldsymbol{q}_l,\boldsymbol{k}_i)}{\sum_{j=0}^{l-1}\phi(\boldsymbol{q}_l,\boldsymbol{k}_j)},
\qquad
h_l=\sum_{i=0}^{l-1}\alpha_{i\to l}\boldsymbol{v}_i.
\tag{9}

のように数式で表されています。これを読み解いていきます。

$w_l$ は学習可能なパラメータです。インプットデータに基づくものではないのでpseudo-query（擬似クエリ）と表現されています。

$v_i$ は先ほどの通常のResidualsの図で下側に並んでいるものと同じです。これをKeyとみなして$k_i$と表現します。

次に$\alpha_{i\to l}$です。

\alpha_{i\to l} = \operatorname{LayerWeight}_{i\to l​}(\boldsymbol{q}_l,\boldsymbol{k}_0,\boldsymbol{k}_1,\cdots,\boldsymbol{k}_{l-1}) 
=\frac{\phi(\boldsymbol{q}_l,\boldsymbol{k}_i)}{\sum_{j=0}^{l-1}\phi(q_l,\boldsymbol{k}_j)}

これは、Self Attentionと同じsoftmax型の重み付けを、トークン方向ではなく深さ方向に適用した仕組みです。今回で言うと$v_i \ (i: 0 〜 l-1)$ を加重平均するウェイトをq(Query)とk(Key)から作ろうと言うことです。この関数を便宜上 LayerWeightと名付けます。（この記事の中だけでの独自関数名称）

この式の中にある$\phi$と言う関数は
$\phi(q,k)=\exp\left(q^\top\operatorname{RMSNorm}(k)\right)$
です。
RMSNormは Root Mean Square Normalization を略した名称で、ベクトル全体の大きさを「各要素の二乗平均平方根」で割ってそろえる正規化関数です。式は以下です。

\operatorname{RMSNorm}(\boldsymbol{x})
=
\boldsymbol{g}
\odot
\frac{
    \boldsymbol{x}
}{
    \sqrt{
        \frac{1}{d}
        \sum_{j=1}^{d} x_j^2
        +
        \varepsilon
    }
}

$\boldsymbol{g}$ は通常、学習されるモデルパラメータです。1で固定されるケースもありますが、実際の実装
https://huggingface.co/moonshotai/Kimi-K3/blob/main/modeling_kimi_linear.py#L226-L236
を見てみると

class KimiRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        dtype = hidden_states.dtype
        x = hidden_states.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
        return self.weight * x.to(dtype)

のような実装であり、ここの self.weight $\boldsymbol{g}$のことです。

ここまでで作ったattention weight(のようなもの)に$\boldsymbol{v}_i$をかけて、完成です。

h_l=\sum_{i=0}^{l-1}\alpha_{i\to l}\boldsymbol{v}_i.

こうみると、TransformerのSelf Attentionの仕組みによく似ています。

参考：TransformerのSelf AttentionのQKVを直感的に解説する

以上から、この図のような関係性となります。ここで$h_1$はToken Embeddingを表します。

（ほんとうは$h_i, f_1(h_1), \cdots, f_{l-1}(h_{l-1})$からそれぞれのLayerWeightに線が伸びているのですが、ゴチャついたので省略しています）

わかりやすいのかわかりにくいのか、どっちかわからなくなりましたが、要はすごいところは、通常のResidualが以下のように、重みをつけずに各レイヤーの出力を足し合わせて最終出力としていたところを、

Self Attentionのような仕組みで各レイヤーに重みをつけて加重平均して最終出力に繋げられているところです。
これは、レイヤー方向・深さ方向のAttentionと見なすことができ、これは非常に面白い仕組みだと私は考えています！

Reference

Kimi K3: Open Frontier Intelligence : https://arxiv.org/abs/2607.24653

1. 例えば$f_{l}(h_{l})$はその前の$f_{l-1}(h_{l-1})$と$h_{l-1}$の影響を受けていることは図では省略。 ↩