Gated KronLSTM Sequence Mixer

Why People Wanted It

A recurrent memory update has to decide what to retain, what old association to erase, and what new association to write. If one gate is made to do several of those jobs, the controls are coupled: changing how much new information is written can also change how much old information gets removed. The practical question is whether a linear-attention block can expose these decisions separately while keeping its recurrent state simple.

KronLSTM takes the Kronecker Delta memory rule and embeds it in a gated sequence-mixer block. The block carries its head memories, rather than extra normalizer or stabilizer memories.

The Core Move

The core move is to give forgetting, erasing, and writing separate controls, then query the resulting matrix memory.

For one head, let $S_t\in\mathbb{R}^{d_v\times d_k}$ be memory, with key $k_t$, value $v_t$, and query $q_t$. Its update is

$$S_t = S_{t-1}F_t + \beta_t\bigl(v_t-S_{t-1}F_t k_t\bigr)k_t^\top,$$

where $F_t$ forgets along key directions, $\beta_t$ controls erasure, and the target $c_t$ controls writing through the parameterization described below. The read is $S_tq_t$.

Quiz

In KronLSTM, what is the central role of separating the erase and write gates?

It replaces the query read with a normalizer state.

It makes the forget gate operate over value directions instead of key directions.

It removes the need for a recurrent memory matrix.

It lets correction strength and the written target be controlled independently.

$ check

Why Erase and Write Need Separate Gates

Now the version for readers who want more of the machinery: the key distinction is that the update can correct the decayed memory toward a gated target rather than forcing the target to be one. Let the desired target be $\tilde c_t=c_t/\beta_t$. Substituting this into the correction form gives

$$S_t=S_{t-1}F_t+\beta_t(\tilde c_t v_t-S_{t-1}F_tk_t)k_t^\top =S_{t-1}F_t+(c_tv_t-S_{t-1}F_tk_t)k_t^\top.$$

In words, the erase amount remains set by $\beta_t$, while the value contribution is scaled by $c_t$. Both gates are per-head scalars computed from the input with sigmoid projections, so each is bounded; they are not two names for one shared knob. This is the architectural separation: forgetting first acts on memory, then the key-directed correction erases the existing response and writes the gated target.

Why Forgetting Acts Per Key Direction

Here is the part that gets technical: the forget operator is diagonal on the key dimension, so it can retain different memory directions by different amounts. Write $F_t=\operatorname{diag}(f_t)$, with each component strictly between zero and one. The update multiplies memory on the right by $F_t$, meaning each key direction is independently attenuated before the correction.

The forget vector is parameterized as $f_t=\sigma(W_fx_t+b_f)^{\exp(a_{\log})}$, with the exponent applied componentwise and one learned scalar per head. In words, the input sets direction-specific retention, while the learned positive power adjusts the head's overall retention horizon. This is not a value-channel forget gate: its diagonal spans key directions, and the matrix memory is contracted from the right.

Why the Readout Is Gated Without Extra Memory

Now the block-level detail: queries, keys, and values come from input projections followed by short causal depthwise convolutions and SiLU activations; queries and keys are L2-normalized. The memory read is $\tilde h_t=S_tq_t$, which means the query selects a value-space output from the stored key-to-value associations.

A sigmoid output gate $o_t=\sigma(W_ox_t+b_o)$ modulates that read, followed by per-head RMSNorm and an output projection over concatenated heads. The gate changes what the current layer emits; it is not another recurrent memory. The block's recurrent state consists of the head memories $S_t$, with no separate normalizer or stabilizer state. So what does this simpler state bookkeeping buy, and what does it leave the gates responsible for?

Minimal Causal Code

This loop makes the memory update visible. It omits the projections, convolutions, normalization, and output gate that wrap the rule in the full block.

import torch

def kronlstm_step(S, k, v, q, F, beta, c):
    # S: [value_dim, key_dim]; F: [key_dim]
    decayed = S * F.unsqueeze(0)
    predicted = decayed @ k
    error = c * v - predicted
    S_next = decayed + error.unsqueeze(-1) * k.unsqueeze(0)
    read = S_next @ q
    return S_next, read

The vector F scales each key column of S; beta is omitted from this compact expression because the write target is represented directly as c in the decoupled update. In the equivalent target form, the correction is scaled by beta and targets c / beta.

What You Gain, and What You Lose

The gate roles are explicit: a per-key-direction forget gate contracts memory, while separate per-head erase and write gates control the correction, and an output gate controls the read. The full block uses only its head memories as recurrent state, with no additional normalizer or stabilizer memories. With L2-normalized keys and values and sigmoid $\beta_t$, the update is non-expansive; the forget gate contracts it further.

The cost is that you now have several input-dependent controls to parameterize and reason about, rather than one coupled update knob. The forget gate only distinguishes key directions, while erase and write gates are scalar per head; the output gate controls emission, not stored memory. KronLSTM is thus not a claim that every gate can express every kind of memory edit. It is a design that assigns distinct jobs to distinct gates.

Quiz

Which recurrent state does the KronLSTM block carry, according to its architecture description?

The per-head memory matrices only.

A memory matrix plus a separate normalizer and stabilizer state.

Only the latest query and key vectors.

A single scalar state shared across all heads.

$ check