Chunkwise-Parallel Execution of KronLSTM

Why People Wanted It

A recurrent memory update is naturally sequential: each token changes the memory used by the next token. The problem is the price of that dependency when training on long sequences. So the real question was never just how to batch tokens; it was whether their causal interactions could be reorganized without explicitly constructing the large combined key-value features.

KronLSTM has a useful property: the coefficient of each rank-one write depends on memory through a scalar score. Solve for those coefficients within a chunk, and you can assemble its writes and readouts using matrix operations. The memory still passes from one chunk to the next.

“All token interactions inside a chunk are expressed as dense matrix products and one small triangular solve, and the memory is propagated only between chunks.” — A8, p. 49

The Core Move

Replace the within-chunk recurrence with a unit lower-triangular system whose interactions factor into key similarities and value similarities.

For a chunk of $C$ tokens, the write coefficients satisfy

$$(I+N)u=c-\beta\odot p.$$

Here $p$ contains the scores against the incoming memory, $c$ contains the write gates, and $\beta$ contains the erase gates. The strictly lower-triangular matrix $N$ accounts for how earlier writes change later scores. The diagonal of $I+N$ is all ones: this is a triangular solve, not an assumption that the tokens are independent.

“Hence $u_r+\sum_{s<r}N_{rs}u_s=c_r-\beta_r p_r$, which is row $r$ of $(I+N)u=c-\beta\odot p$.” — Proposition A5 proof, p. 50

Quiz

What replaces token-by-token evaluation of interactions inside a KronLSTM chunk?

A unit lower-triangular solve using elementwise products of key and value Gram interactions

A full dense inverse of the memory state

Independent per-token writes that ignore earlier writes

Explicit construction of every combined key-value feature

$ check

Why the Gram Product Replaces Kronecker Features

Here is the part that gets technical: the combined-feature interaction factors into a key inner product times a value inner product.

Forgetting is absorbed into asymmetric keys: $\tilde{k}_r$ is the decay-scaled scoring key, while $\bar{k}_s$ is the writing key in the decay-normalized frame. Their corresponding combined features are

$$\tilde{a}_r=\tilde{k}_r\otimes v_r, \qquad \bar{a}_s=\bar{k}_s\otimes v_s.$$

The mixed-product identity gives

$$\tilde{a}_r^\top\bar{a}_s =(\tilde{k}_r^\top\bar{k}_s)(v_r^\top v_s).$$

So the interaction can be computed without forming either combined feature. Let $G_{kk}$ be the strictly lower-triangular decay-weighted key Gram matrix and $G_{vv}$ the value Gram matrix. The system matrix uses

$$N=\operatorname{diag}(\beta)(G_{kk}\odot G_{vv}).$$

The elementwise product combines the two similarities; multiplication by $\operatorname{diag}(\beta)$ weights each row by that token's erase gate. Strict lower-triangular masking preserves causality: an earlier write can change a later write coefficient, but not the reverse.

The asymmetric keys matter. With channel-wise forgetting, this is not simply the ordinary symmetric key Gram matrix. The scoring and writing sides carry different decay scalings.

“The $d_kd_v$-dimensional features are never formed: two $C\times C$ Gram matrices of $d_k$- and $d_v$-dimensional vectors and one Hadamard product suffice.” — A8.3, p. 50

This removes the large feature construction. But how does the incoming memory enter the solve?

Why the Incoming Memory Enters as Scalar Scores

Now the version for readers who want more of the machinery: unrolling the decay-normalized recurrence separates the incoming-memory score from the scores contributed by earlier writes.

The token-level recurrence is

$$S_r=S_{r-1}F_r+u_rv_rk_r^\top, \qquad u_r=c_r-\beta_r y_r, \qquad y_r=v_r^\top S_{r-1}F_rk_r,$$

where $F_r$ is the diagonal key-side forget gate. The readout uses the memory after the current write:

$$h_r=S_rq_r.$$

Because the forget matrices are diagonal and commute, cumulative decay can be absorbed into the asymmetric keys. In that frame, the incoming state contributes

$$p_r=v_r^\top S_0\tilde{k}_r.$$

Substituting the unrolled memory into each score produces the coefficient equation from the core move. Its terms have distinct jobs: $c$ sets the candidate write; $\beta\odot p$ corrects for the incoming memory; and $Nu$ corrects for earlier writes within the chunk.

The write gate therefore enters only the right-hand side. The erase gate enters both the system matrix and the incoming-memory term. Since $N$ is strictly lower triangular, $I+N$ is invertible for every erase-gate value and every input. This is an algebraic reformulation of the recurrence, not an approximation to it.

“The write gate $c_r$ enters only the right-hand side of the solve, the erase gate $\beta_r$ enters the system matrix and the coupling to $S_0$.” — A8.3, p. 50

“$N$ is strictly lower triangular, so $I+N$ is unit lower triangular and invertible for every $\beta$ and every input.” — Proposition A5, p. 50

Once those coefficients are known, do we still need every intermediate memory matrix?

Why Memory Propagates Only at Chunk Boundaries

Here is the deeper chunk-level consequence: solved write coefficients let you assemble the outgoing memory without materializing every intermediate revision.

Each token's pseudo-value is its original value scaled by $u_r$. Using $\Gamma_C$ for cumulative diagonal decay through the chunk, the end state is

$$S_C=S_0\Gamma_C+\sum_{s\le C}u_s v_s\bar{k}_s^\top\Gamma_C.$$

The first term carries the incoming memory through the chunk's forgetting. Each term in the sum carries a solved write through the decay remaining after its position. This is the end-state expression used in the proof of Theorem A11.

Readouts follow the same additive representation:

$$h_r=S_0\tilde{q}_r+\sum_{s\le r}u_s(\bar{k}_s^\top\tilde{q}_r)v_s.$$

Here $\tilde{q}_r$ is the query scaled by cumulative decay. Stacking the readouts turns the sum into multiplication by the causal readout Gram matrix. Its diagonal is included because each query sees its own token's write.

The chunk-local Gram matrices can be prepared across chunks. With $\Lambda=(I+N)^{-1}$, the state-independent terms $\Lambda c$ and $\Lambda\operatorname{diag}(\beta)$ can also be prepared in parallel. This inverse notation describes the algebra; the computation is organized around a triangular solve.

The remaining dependency is incoming memory, then scores $p$, then coefficients $u$, then outgoing memory. The next chunk must receive that outgoing state before its state-dependent part can be evaluated.

“Everything in (A152)–(A159) except $p$, $u$ and [the pseudo-value matrix] is local to a chunk and is computed for all chunks at once.” — A8.4, p. 51

“The dependence $S^{[n]}\to p\to u\to S^{[n+1]}$ is the only sequential part.” — A8.4, p. 51

How much work should you group into a chunk when this state handoff remains?

Minimal Causal Code

This plain PyTorch loop implements the recurrence in A149. It is a reference for the semantics, not a fast chunkwise implementation: the triangular solve replaces the loop's within-chunk coefficient dependencies.

import torch

def kronlstm_chunk_reference(S, q, k, v, beta, c, forget):
    # S: [dv, dk]; q, k: [C, dk]; v: [C, dv]
    # beta, c: [C]; forget: [C, dk], with entries in (0, 1)
    reads = []
    for r in range(k.shape[0]):
        # Apply key-side forgetting before scoring the current pair.
        decayed_S = S * forget[r].unsqueeze(0)
        score = v[r] @ (decayed_S @ k[r])
        u = c[r] - beta[r] * score

        # Visible state transition: decayed memory plus rank-one write.
        S = decayed_S + u * torch.outer(v[r], k[r])
        # The query sees the current token's write.
        reads.append(S @ q[r])
    return S, torch.stack(reads)

What You Gain, and What You Lose

You gain dense chunk-local computation without constructing combined features of dimension $d_kd_v$. The write coefficients come from a structured triangular system, and readouts can be assembled without explicitly stepping through every intermediate memory.

You do not gain independent chunks. The incoming state still determines the scalar scores and therefore the outgoing state. You also introduce chunk-local Gram matrices and a triangular solve rather than simply making the recurrent loop disappear. Chunk length controls how much work is grouped into that machinery; this formulation alone does not specify an optimal choice.

Finally, the rewrite depends on the update structure: diagonal key-side decay and a rank-one write whose coefficient is affine in a scalar score. It is not a drop-in transformation for arbitrary recurrent networks.

“The derivation uses only the structure ‘diagonal decay on the key side, rank-one write along $v_rk_r^\top$ with a coefficient affine in the scalar score’.” — A8.6, p. 51

Quiz

Which dependency remains sequential in chunkwise-parallel KronLSTM?

Every pairwise interaction inside a chunk must still be evaluated in a token loop

The value Gram matrix depends on the previous chunk's memory

No sequential dependency remains after splitting the sequence into chunks

The state handoff between chunks, through incoming-memory scores and solved writes

$ check