Problems
Loading...
Step 1 of 11
KDA Recurrence
Hard

Kimi Delta Attention updates one matrix state per batch item and head after every token. Queries and keys have shape (B,S,H,Dk)(B,S,H,D_k), values and output-gate logits have shape (B,S,H,Dv)(B,S,H,D_v), and the state has shape (B,H,Dk,Dv)(B,H,D_k,D_v). The write strength supplies one scalar per token and head, while the decay logits have the same shape as the keys.

For retention vector αt\alpha_t, write strength βt\beta_t, key ktk_t, value vtv_t, and previous state St−1S_{t-1}, update the state in sequence order with

St=(I−βtktktT)Diag⁡(αt)St−1+βtktvtT.S_t=\left(I-\beta_tk_tk_t^{\mathsf T}\right)\operatorname{Diag}(\alpha_t)S_{t-1}+\beta_tk_tv_t^{\mathsf T}.

For lower log-decay bound gmin⁡<0g_{\min}<0 and decay logit ztz_t, compute

αt=exp⁡(gmin⁡sigmoid⁡(zt)).\alpha_t=\exp\left(g_{\min}\operatorname{sigmoid}(z_t)\right).

Read from the updated state, not the previous state:

o~t=StTqt.\widetilde{o}_t=S_t^{\mathsf T}q_t.

Apply RMS normalization independently over each head's DvD_v values, multiply by the sigmoid output gate, concatenate the HH heads, and apply the supplied projection from width HDvHD_v to model width. Return a dictionary with outputs and final_state. Their shapes are (B,S,Dmodel)(B,S,D_{model}) and (B,H,Dk,Dv)(B,H,D_k,D_v), and both preserve the input floating-point dtype and device.

Loading visualization...

Examples

Input: query = [[[[1]]]], key = [[[[1]]]], value = [[[[2]]]], decay_logits = [[[[0]]]], write_strength = [[[[0.5]]]], output_gate_logits = [[[[0]]]], output_projection = [[1]], initial_state = [[[[0]]]], g_min = -5, eps = 1e-06

Output: {"outputs": tensor([[[0.4999997500001875]]]), "final_state": tensor([[[[1.0]]]])}

Explanation: The write creates a state value of 1. RMS normalization keeps the scalar readout near 1, and the sigmoid gate scales it by 0.5.

Input: query/key/decay logits shape (1, 2, 1, 2), value shape (1, 2, 1, 1), initial state shape (1, 1, 2, 1), output projection shape (2, 1)

Output: {"outputs": tensor of shape (1, 2, 2), "final_state": tensor of shape (1, 1, 2, 1)}

Input: query/key/decay logits shape (1, 2, 2, 1), value shape (1, 2, 2, 1), initial state shape (1, 2, 1, 1), output projection shape (2, 2)

Output: {"outputs": tensor of shape (1, 2, 2), "final_state": tensor of shape (1, 2, 1, 1)}

Hint 1

Broadcast each retention vector across the value dimension before applying it to the state.

Hint 2

Use outer products for both the erase term and the key-value write term.

Hint 3

Read from the updated state, normalize each head over its value width, then concatenate the heads.

Requirements

  • Return a dictionary with exactly outputs and final_state. The output tensor has shape (batch, sequence, model width).
  • Preserve the floating-point dtype and device of the tensor inputs.
  • Do not mutate any input tensor.
  • Match the stated recurrence within absolute tolerance 1e-6 and relative tolerance 1e-5.

Constraints

  • Query and key tensors have shape (batch, sequence, heads, key width).
  • Value and gate tensors have shape (batch, sequence, heads, value width).
  • Write strengths lie in [0, 1], and the lower log-decay bound is negative.
  • The output projection accepts the concatenated head values.
Prev1 / 11Next