Yang et al., MIT and MIT-IBM Watson AI Lab (2023) — ARXIV-2312.06635
Key takeaways
- The practical problem it names. Linear attention has better asymptotics and can be written as an RNN with a matrix-valued hidden state, so inference is linear-time. It was still slower in wall-clock than a good softmax kernel, because the existing implementations were not I/O-aware. That is a systems failure being mistaken for an architectural one.
- FlashLinearAttention is the fix: a chunkwise algorithm that trades memory movement against parallelism, with secondary chunking inside each chunk so the inter-sub-chunk part runs in half precision on tensor cores while the intra-sub-chunk part stays in full precision in log space. Faster than FlashAttention-2 as a standalone layer even at 1K sequence length, which is where the asymptotic argument has not begun to pay.
- The gate. Plain linear attention has no decay term, so it cannot forget; the paper cites that as a suspected cause of its instability on long contexts. RetNet-style fixes use a global, non-data-dependent decay. GLA makes the forget gate data-dependent and per-channel, and the chunkwise algorithm is generalised to carry the cumulative decay across chunks.
- Results. Competitive with a LLaMA-architecture Transformer and with RetNet and Mamba at moderate scale, with higher training throughput than a similarly-sized Mamba. The standout is length generalization: trained at 2K, no significant perplexity degradation past 20K.
Standing in the anthology
A correction to what this note first said. It claimed GLA was one of the two parents of Gated DeltaNet (LIT-137) — the gate half, to LIT-195’s delta rule. That is wrong, and declaring the lineage edges is what caught it. Gated DeltaNet’s own abstract names the two mechanisms it combines as Mamba2’s gate and DeltaNet’s update rule, and its title is Improving Mamba2 with Delta Rule. The gating lineage runs through LIT-162, not through here.
What GLA does hold is a parallel claim to the same idea — a data-dependent, per-channel forget gate for linear attention, a year before Mamba-2 — reached independently and by a different construction. Two groups arriving at data-dependent gating is worth having in the record; it is not the same as being the step Gated DeltaNet built on.
It is also the origin of the FlashLinearAttention library that both parents ship in and that Kimi K3 (LIT-131) names in its kernel work — so the line this record carries as an architecture story is partly a kernel story, and this is where that starts.
Filed via the reference pass in #40, from Kimi K3 §2.1.1.
No practice of its own: parallel invention, not a step the record’s line took. Data-dependent per-channel gating arrived here a year before the paper that this record’s recommended module actually builds on, and by a different construction — which is worth holding as a second data point and is not a recommendation the record can source to anything it recommends. Its other contribution, the kernel library both parents ship in, is infrastructure the practices assume rather than advice they give.