Understanding Attention Residuals in Kimi K3

Hi, I'm Matsuken, a data scientist at a Japanese technology company.

In this series, I will explain two key components described in the recently released Kimi K3 Technical Report:

• Attention Residuals
• Kimi Delta Attention (KDA)

I found both mechanisms fascinating, so I would like to explain them in a way that makes them intuitive to understand. As the first installment in this series, this article focuses on Full Attention Residuals.

This is an English version of my original Japanese article, which I originally wrote and published on Qiita.

---

The Technical Report first introduces Full Attention Residuals as the basic form of the mechanism. Kimi K3 itself uses Block Attention Residuals, which reduce memory consumption and communication volume. This article, however, focuses on Full Attention Residuals so that we can concentrate on understanding the underlying idea.

The core idea: Standard residual connections add preceding layer outputs without weighting them. Attention Residuals learn how much each previous layer should contribute.

Let's begin with a quick review of standard residual connections, then see what makes Full Attention Residuals different.

Standard Residual Connections

We can express a sublayer with a residual connection as follows:

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

This is a familiar residual layer.

For the first layer, where $l=1$ , we have:

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

Here, the token embedding $h_1$ is added to its transformed value $f_1(h_1)$ , and the result is passed to the next layer.

For the second layer:

$$h_3 = h_2 + f_2(h_2)$$

Substituting equation $(1)$ , obtained from the first layer, gives:

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

Similarly, for the third layer:

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

Repeating the same expansion through layer $(l-1)$ , we obtain:

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

Using summation notation, this becomes:

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

The relationship can be visualized as follows.¹

Full Attention Residuals

With standard residual connections reviewed, let's move on to Full Attention Residuals. The Technical Report defines them as follows:

Let's unpack these equations.

$\boldsymbol{w}_l$ is a learnable parameter. Because it does not depend on the input data, the report refers to it as a pseudo-query.

The $\boldsymbol{v}_i$ terms are the same values shown along the bottom of the preceding diagram of standard residual connections. We also treat each of these values as a Key and denote it by $\boldsymbol{k}_i$ .

Next, consider $\alpha_{i\to l}$ :

This mechanism applies softmax-style weighting, similar to self-attention, across the depth dimension rather than across the token dimension. In this case, the idea is to use the Query $\boldsymbol{q}_l$ and Keys $\boldsymbol{k}_i$ to construct the weights for a weighted average of $\boldsymbol{v}_i$ , where $i=0,\ldots,l-1$ . For convenience, I will call this function LayerWeight—a name used only in this article.

The function $\phi$ in the equation is defined as:

$$\phi(\boldsymbol{q},\boldsymbol{k}) =\exp\left(\boldsymbol{q}^\top\operatorname{RMSNorm}(\boldsymbol{k})\right)$$

RMSNorm stands for Root Mean Square Normalization. It normalizes the overall magnitude of a vector by dividing it by the root mean square of its elements. The formula is:

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

The vector $\boldsymbol{g}$ is usually a learnable model parameter. It can also be fixed to one, but the actual Kimi K3 implementation uses the following code:

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)

Here, self.weight corresponds to $\boldsymbol{g}$ .

Finally, we multiply each $\boldsymbol{v}_i$ by the attention-like weight derived above and sum the results:

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

Seen this way, the mechanism closely resembles self-attention in a Transformer.

Related article: An Intuitive Explanation of Q, K, and V in Transformer Self-Attention

Putting everything together gives the relationship shown below. Here, $h_1$ represents the token embedding.

(Strictly speaking, each LayerWeight receives connections from all of $h_1, f_1(h_1), \ldots, f_{l-1}(h_{l-1})$ . I omitted those lines because the diagram became too cluttered.)

At this point, I am no longer sure whether the diagrams make this easier or harder to follow! The key idea, though, is simple. A standard residual connection produces its final output by adding the outputs of the layers together without weighting them:

Full Attention Residuals instead use a self-attention-like mechanism to assign weights to the individual layers, take their weighted average, and pass that result to the final output.

In other words, the mechanism can be viewed as attention across layers—or across the depth dimension. I think this is a particularly fascinating idea!

References

• Kimi K3: Open Frontier Intelligence

---

1. For example, the diagram omits the fact that $f_l(h_l)$ is influenced by both $f_{l-1}(h_{l-1})$ and $h_{l-1}$ . ↩