Study Plan

Gated DeltaNet State Update

Language ModelsLinear AlgebraNeural Networks

Hard

Implement the Gated DeltaNet recurrent update. At position $t$ , $q_t$ , $k_t$ , and $v_t$ are query, normalized key, and value vectors; $\gamma_t$ and $\beta_t$ are scalar decay and update gates. $S_t$ is the state, $y_t$ is the output, and $I$ is the Dk-by-Dk identity matrix. Superscript T denotes transpose.

S_t=\gamma_t\left(I-\beta_tk_tk_t^{\mathsf T}\right)S_{t-1}+\beta_tk_tv_t^{\mathsf T}.

y_t=q_t^{\mathsf T}S_t.

Each batch state starts at zero. Return a dictionary containing outputs with shape (B, S, Dv) and final_state with shape (B, Dk, Dv), both PyTorch tensors with the same dtype and device as q.

Loading visualization...

Examples

Input: q = [[[1,0],[1,0]]], k = [[[1,0],[1,0]]], v = [[[2,0],[5,0]]], gamma = [[1,1]], beta = [[1,1]]

Output: {"outputs":[[[2,0],[5,0]]],"final_state":[[[5,0],[0,0]]]}

Explanation: The repeated unit key erases its previous state direction before the second value is written, so 5 replaces 2.

Input: q = [[[1,0],[1,0]]], k = [[[1,0],[1,0]]], v = [[[4,0],[99,0]]], gamma = [[1,0.5]], beta = [[1,0]]

Output: {"outputs":[[[4,0],[2,0]]],"final_state":[[[2,0],[0,0]]]}

Input: q = [[[1,0]],[[0,1]]], k = [[[1,0]],[[0,1]]], v = [[[7]],[[9]]], gamma = [[1],[1]], beta = [[1],[1]]

Output: {"outputs":[[[7]],[[9]]],"final_state":[[[7],[0]],[[0],[9]]]}

Hint 1

torch.eye creates the identity matrix on a chosen device.

Hint 2

torch.bmm applies a different matrix to each batch item.

Hint 3

torch.unsqueeze supports batched outer products through broadcasting.

Requirements

Use PyTorch, with float32 intermediate calculations and results cast to the query dtype.
Each batch item has its own independent recurrent state.
Do not modify any input tensor.

Constraints

q and k have shape (B, S, Dk), v has shape (B, S, Dv), and gamma and beta have shape (B, S).
B is batch size, S is sequence length, and Dk and Dv are key and value widths; all are positive.
All inputs are PyTorch tensors on the same device; q, k, and v share dtype float32 or bfloat16, while gamma and beta are float32.
Keys are unit-normalized before conversion to the input dtype; bfloat16 rounding may change their norms slightly. Use the supplied keys without renormalizing them.
All gate values lie in [0, 1].
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

Parallel and Recurrent Linear AttentionMedium Mamba 2 Gated State ScanMedium Top-k MoE Router with Load StatisticsHard Orthogonalize a Muon Momentum UpdateEasy RMSNorm Forward PassEasy Decode Attention with a KV CacheHard Recover an IsoFLOP Compute-Optimal FrontierHard Fit a Data Power Law with a Noise FloorHard Jacobian of a Neural Network LayerHard JVP and VJP for a 2-Layer NetworkHard Scaled Dot-Product Attention in PythonMedium RBF Kernel MatrixHard

Gated DeltaNet State Update

Language ModelsLinear AlgebraNeural Networks

Hard

S_t=\gamma_t\left(I-\beta_tk_tk_t^{\mathsf T}\right)S_{t-1}+\beta_tk_tv_t^{\mathsf T}.

y_t=q_t^{\mathsf T}S_t.

Each batch state starts at zero. Return a dictionary containing outputs with shape (B, S, Dv) and final_state with shape (B, Dk, Dv), both PyTorch tensors with the same dtype and device as q.

Loading visualization...

Examples

Input: q = [[[1,0],[1,0]]], k = [[[1,0],[1,0]]], v = [[[2,0],[5,0]]], gamma = [[1,1]], beta = [[1,1]]

Output: {"outputs":[[[2,0],[5,0]]],"final_state":[[[5,0],[0,0]]]}

Explanation: The repeated unit key erases its previous state direction before the second value is written, so 5 replaces 2.

Input: q = [[[1,0],[1,0]]], k = [[[1,0],[1,0]]], v = [[[4,0],[99,0]]], gamma = [[1,0.5]], beta = [[1,0]]

Output: {"outputs":[[[4,0],[2,0]]],"final_state":[[[2,0],[0,0]]]}

Input: q = [[[1,0]],[[0,1]]], k = [[[1,0]],[[0,1]]], v = [[[7]],[[9]]], gamma = [[1],[1]], beta = [[1],[1]]

Output: {"outputs":[[[7]],[[9]]],"final_state":[[[7],[0]],[[0],[9]]]}

Hint 1

torch.eye creates the identity matrix on a chosen device.

Hint 2

torch.bmm applies a different matrix to each batch item.

Hint 3

torch.unsqueeze supports batched outer products through broadcasting.

Requirements

Use PyTorch, with float32 intermediate calculations and results cast to the query dtype.
Each batch item has its own independent recurrent state.
Do not modify any input tensor.

Constraints

q and k have shape (B, S, Dk), v has shape (B, S, Dv), and gamma and beta have shape (B, S).
B is batch size, S is sequence length, and Dk and Dv are key and value widths; all are positive.
All inputs are PyTorch tensors on the same device; q, k, and v share dtype float32 or bfloat16, while gamma and beta are float32.
Keys are unit-normalized before conversion to the input dtype; bfloat16 rounding may change their norms slightly. Use the supplied keys without renormalizing them.
All gate values lie in [0, 1].
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

q =

Accepts: array

k =

Accepts: array

v =

Accepts: array

gamma =

Accepts: array

beta =

Accepts: array