Kimi (AI)

Kimi Delta Attention (KDA)

Kimi Delta Attention (KDA) is an attention mechanism that runs in linear time and constant memory.^[44]

In the standard causally masked dot-product attention mechanism, the input sequence of vectors ${\displaystyle x_{1},x_{2},\dots }$ are converted to a triple sequence of query-key-value vectors ${\displaystyle q_{i}=W^{Q}x_{i},k_{i}=W^{K}x_{i},v_{i}=W^{V}k_{i}}$ via matrices ${\displaystyle W^{Q},W^{K},W^{V}}$, and the t-th output vector is ${\displaystyle o_{t}=\sum _{i\leq t}{\frac {e^{k_{i}^{T}q_{t}}}{\sum _{j\leq t}e^{k_{j}^{T}q_{t}}}}v_{i}}$.

A naive linear attention mechanism is to take off the softmax: ${\displaystyle o_{t}=\left(\sum _{i\leq t}v_{i}k_{i}^{T}\right)q_{t}}$. The thing in the brackets is a memory matrix ${\displaystyle S_{t}}$, such that ${\displaystyle o_{t}=S_{t}q_{t}}$, and the memory updates via ${\displaystyle S_{t}=S_{t-1}+v_{t}k_{t}^{T}}$.

This update resembles a linear recurrent network, so one considers modifying it in a way similar to LSTM:$${\displaystyle {\begin{aligned}{\tilde {S}}_{t}&=\operatorname {diag} (\alpha _{t})S_{t-1}\quad &{\text{ (forget gating)}}\\{\tilde {v}}_{t}&={\tilde {S}}_{t}k_{t}&{\text{ (predict)}}\\S_{t}&={\tilde {S}}_{t}+\beta _{t}(v_{t}-{\tilde {v}}_{t})k_{t}^{T}&{\text{ (memory updating)}}\\\end{aligned}}}$$where the vector ${\displaystyle \alpha _{t}}$ is obtained by applying a learned function to ${\displaystyle x_{t}}$.

Attention Residuals

Attention Residuals is an attention mechanism over the residual stream.^[44] Consider a neural network with layers 1, 2, 3, .... Let the input of layer i be ${\displaystyle h_{i}}$, and let its output be ${\displaystyle h_{i}'}$. We also write ${\displaystyle h_{0}':=h_{1}}$ to make the notation cleaner.

In non-residual neural networks, ${\displaystyle h_{i}=h_{i-1}'}$. In residual neural networks, ${\displaystyle h_{i}=h_{0}'+h_{1}'+\dots +h_{i-1}'}$. In Attention Residuals, we use the standard dot-product attention:$${\displaystyle h_{i}=\operatorname {SoftmaxAttention} (q_{i},K_{i},V_{i}),\quad K_{i}=V_{i}=[h_{0}',h_{1}',\dots ,h_{i-1}']}$$where the vectors ${\displaystyle q_{1},q_{2},\dots }$ are learnable parameters, one for each layer. This is similar to the idea of learnable positional encoding vectors.

In practice, the Kimi K3 paper used a variant called Blocked Attention Residuals, that they claimed preserved essentially all the benefits of Attention Residuals while saving memory. The layers are partitioned into consecutive blocks, each containing ${\displaystyle n_{1},n_{2},\dots }$ layers. Let ${\displaystyle (n,i)}$ denote the i-th layer in the n-th block, then $${\displaystyle {\begin{aligned}h_{(n,i)}&=\operatorname {SoftmaxAttention} (q_{(n,i)},K_{(n,i)},V_{(n,i)}),\\K_{(n,i)}=V_{(n,i)}&={\begin{cases}\left[\sum _{j}h_{(1,j)}',\sum _{j}h_{(2,j)}',\dots ,\sum _{j}h_{(n-1,j)}'\right]&{\text{if }}i=1,\\\left[\sum _{j}h_{(1,j)}',\sum _{j}h_{(2,j)}',\dots ,\sum _{j}h_{(n-1,j)}',\sum _{j=1}^{i-1}h_{(n,j)}'\right]&{\text{if }}i\geq 2.\end{cases}}\end{aligned}}}$$

Others

Stable LatentMoE is a variant of LatentMoE adapted to Kimi K3, which uses 896 routed experts with 16 active experts per token, corresponding to a sparsity of 56. This makes normal LatentMoE unstable, and their variant makes it stable again.^[44]