Browse · Writeups
Writeup

Post-Transformer Architectures in 2024: Cheaper Sequences, Finite Memory

Linear attention, state-space models and RWKV reduce the cost of processing sequences. Their harder challenge is preserving the information a task needs—and making that computation fast on real hardware.

Source and teaching reviews passed · Human editorial review pending.

A whole book, a quadratic bill

Feed an entire book into Llama, and dense self-attention compares its tokens against one another. Dan Fu uses this example to explain the problem motivating post-transformer architectures: pairwise work grows quadratically with sequence length, even though the output still contains one representation per input token.

Longer sequences can mean larger documents, more visual tokens or more generated output. In this 2024 discussion, Fu connects that pressure to OpenAI o1 and test-time compute: spending additional computation while producing an answer, including generating more tokens. Could a different sequence operator preserve quality while its work grows as N log N or N^(3/2), rather than N², for a sequence of length N?

The efficiency goal already had a deployment story. In the recording’s introduction, Charlie reports that RWKV v5, codenamed Eagle, had shipped to 1.5 billion Windows 10 and 11 machines for energy-sensitive, on-device Copilot use cases. The same introduction announces QRWKV6, a Qwen 32B model modified with RWKV linear attention layers.

Linear attention: cheaper algebra does not guarantee faster execution

Dense self-attention compares token queries with token keys to form a pairwise score matrix, applies a row-wise softmax, and multiplies the resulting weights by the values. Fu’s explanation of linear attention starts by removing that intervening softmax. The matrix multiplications can then be regrouped: compute the key–value product first, avoiding the quadratic intermediate. Feature maps and attention approximations offer variants of this approach.

Around 2020, these approaches encountered two obstacles: - Quality did not reliably match dense attention. Reducing computation could compromise the model’s ability to handle the task. - Better scaling did not guarantee lower latency. Naive feature-map implementations could cost enough to make the supposedly efficient operator slower in wall-clock time.

That experience established a recurring requirement: architecture and implementation must improve together. An attractive growth rate alone will not secure adoption if the implementation runs more slowly.

State-space models improve memory updates and parallel execution

A recurrent model carries a hidden state, a running summary of earlier inputs that it updates as new tokens arrive. State-space models brought ideas from dynamical systems and signal processing to the design of those updates. The aim was to use the available state more effectively. Fu reports meaningful improvements on long-sequence evaluations such as Long Range Arena and on time-series tasks.

Recurrence also creates a computational bottleneck. In a naively evaluated recurrent neural network, or RNN, each token’s update waits for the previous one, restricting parallel execution. [S4](https://arxiv.org/abs/2111.00396) reformulates its recurrent computation as a convolution, which can be evaluated using a fast Fourier transform, or FFT. Fu describes this as N log N work in sequence length, using a relatively well-optimized hardware primitive.

Principled state updates addressed quality; FFT-based convolution addressed execution. Specialized kernels—optimized implementations of the underlying operations—such as FlashFFTConv continued the hardware work. By 2022–2023, alternative models were showing meaningful wall-clock advantages over transformers in some regimes, while still sometimes lagging in language-modeling quality.

Selection makes finite state more useful

A compact state helps only if the model can extract the information it needs. Selection mechanisms let that extraction respond to the input. The discussion traces three approaches: - H3 and Hyena use element-wise gates to improve selection from the recurrent state. - [Mamba](https://arxiv.org/abs/2312.00752) brings selection inside the state update. Data-dependent state-space parameters respond to the current input, extending selection beyond gates surrounding the layer. - BASED combines linear and sliding-window attention. Credited to Simran Arora and collaborators, it uses a Taylor approximation of softmax attention alongside sliding-window attention, reportedly improving the tradeoff between sequence recall and recurrent-state size.

Mamba also carries forward the implementation lesson: its treatment of the GPU memory hierarchy makes hardware efficiency part of the architecture’s design. A more selective state update still needs an efficient kernel.

Selection exposes a weakness in a familiar prompt layout. When a long document precedes a question, a fixed-state model must compress the document before knowing which detail will matter. A niche fact may disappear before the model learns that it was essential.

Just Read Twice, credited to Simran Arora, changes the prompt to document → question → document → question. On the second reading, the model already knows what to look for. The talk reports better quality, especially on recall-intensive tasks. This makes inference strategy part of architecture research: cheaper token processing can make an additional pass worthwhile.

The overview closes with examples of where these ideas were being applied in 2024: - Jamba combined architectures. The talk presents this hybrid mixture-of-experts model as a leading alternative-architecture result before that morning’s releases, without specifying evaluation conditions for the ranking. - SANA scaled image generation with linear attention. Attributed to MIT, the diffusion model replaces attention layers with linear attention to process larger images and longer sequences more efficiently. - DNA modeling extended the application range. A gated state-space model used for DNA modeling appeared on the cover of Science.

RWKV: keep the GPU busy, then improve quality

Eugene describes [RWKV](https://arxiv.org/abs/2305.13048) as a separate branch rooted in recurrent networks and linear attention, with ideas exchanged between its community and state-space researchers. Contributors developed it through Discord and GitHub; EleutherAI helped them write their first paper.

The project’s stated goal was accessible AI across languages, nations and compute budgets. Eugene reports training across more than 100 languages, with an ambition to reach 200, while lowering compute costs enough to support devices such as Raspberry Pis.

The starting constraint was recurrence’s sequential dependency. An early experiment removed a dependency and improved GPU utilization, but initially had poor loss. The enduring priority became cascading computation across tokens and layers to keep the GPU occupied. Other architectural choices could change; some successful changes began as working code and received a theoretical rationale afterward.

RWKV’s high-level design has two main blocks, with calculations that change between versions: - `TimeMix` handles long-term memory state, carrying information through the recurrent computation. - `ChannelMix` incorporates nearby-token information, using token shifting to bring in a neighboring token.

Eagle and Finch, versions 5 and 6, use matrix-valued state. [GoldFinch](https://arxiv.org/abs/2407.12077) is a hybrid model. At the time of the talk, a paper for version 7, codenamed Goose, was being written.

Reuse pretrained weights to test a new attention mechanism

Limited compute prompted a different experimental strategy. The RWKV team reports doing most research on a handful of H100 GPUs. Converting a pretrained transformer let the team reuse expensive pretraining and focus new training on a replacement attention mechanism.

QRWKV6, announced that morning as a 32B instruction-model preview, started from Qwen 32B Instruct. The conversion followed three steps: 1. Freeze the feed-forward layers and replace query–key–value attention with RWKV linear layers. The conversion keeps TimeMix and omits RWKV’s ChannelMix. 2. Train the replacement layers with the feed-forward layers frozen. Freezing is presented as necessary for the new attention mechanism to learn initially. 3. Unfreeze and train all layers together, using a custom learning-rate schedule so the components adapt to one another.

The team reports roughly matching the original model’s performance after a few hours on two nodes, with no more than three days of training across the preview effort. This was a preliminary result: some benchmark performance declined, and the discussion supplies neither a full benchmark comparison nor the node configuration.

A 70B conversion was underway, with version 7 and hybrid conversions planned. The enterprise opportunity remained conditional: if a 70B conversion reached 128K context and matched the relevant benchmarks, it might cover substantial existing workloads. That projection relied on the speaker’s assertion that most enterprise workloads used 70B models below 32K context; successful 128K scaling had not been reported. Applying the conversion approach to Mamba was also an expectation.

Hybrids sometimes outperformed both constituent baselines, according to the RWKV discussion and conversations with the Black Mamba and Jamba teams. Why remained unresolved. The speaker described evidence from only roughly four experiments or teams, making the result a reason for further investigation rather than a general rule.

Design the state around matrix hardware

[ThunderKittens](https://github.com/HazyResearch/ThunderKittens) reduces the burden of kernel development. The CUDA library grew from the experience of spending about two months implementing efficient computation for a new architecture. Even changing one line of PyTorch could entail another week of CUDA work.

The library organizes computation around matrices rather than individual floats. Its hardware example is H100 warp-group matrix multiplication: small operations, such as multiplying two 64-by-64 matrices, can inform state sizes and update functions. The team uses ThunderKittens both to reimplement existing architectures and to design new ones around these tensor-core operations. Hardware becomes an input to model design.

The proposed applications ask what that efficiency could enable: - Persistent game state: extend interactive video generators beyond the few seconds of memory they were reported to retain. - Faster video generation: preserve quality while reducing generation time. - Long-lived assistant memory: retain useful information from a week of observations.

Finite memory leaves a role for external search

The conversation turns to retrieval-augmented generation, or RAG, and whether state-based models could make it unnecessary. One speaker offers a limited personal observation: large differences in embedding-model quality repeatedly made little difference to the final application. Without specified datasets or evaluations, the observation leaves the larger question open—as the speaker acknowledges.

The central distinction is processing indefinitely versus remembering everything. Fixed-state architectures—including RWKV, state-space models and xLSTM—use bounded memory rather than an ever-expanding state. The research problem is how efficiently they use that capacity: what they retain, what they forget and whether the retained information supports the task.

The speaker mentions an RWKV state of about 40 MB and a possible future 400 MB state, but state size alone does not establish recall capacity. The associated token-capacity estimates were speculative. External search could complement this finite memory: a model that cannot remember an article could query a database, much as a person searches Notion for a forgotten company document.

A related proposal is an editable external factual store. One speaker takes small models’ conversational fluency but weaker factuality as a clue that larger models devote substantial capacity to facts. That is an intuition about model weights, motivating a store—potentially involving gradient descent—where knowledge could be updated or deleted without relying solely on what the model encodes in its weights.

Evaluate useful memory, not just a context limit

The Q&A challenges the context-length race itself. One speaker questions demand for two-million-token prompts and suggests specialized models for concrete applications, recalling that the original motivation was to challenge quadratic computation. A participant twice responds that they use Gemini’s long context frequently; the discussion acknowledges those uses. The exchange establishes a disagreement about demand, not evidence that very long context has no users.

Fixed-state inference does not eliminate long-sequence training memory costs. Backpropagation needs intermediate states across the token sequence. Training RWKV on a million-token sequence therefore requires rolling out that sequence, leaving a bottleneck in GPU memory.

Reasoning adds another tradeoff. In a hypothetical comparison, a 70B model uses twice as many reasoning tokens as a 400B model to reach the same result. The speaker argues that the smaller model could still consume fewer resources, motivating fast, efficient architectures below 200B parameters. The proposal is to compare resource cost at equivalent task success, allowing both model size and token budget to vary; no measured comparison is offered.

An audience question about models that continuously watch an input stream brings the issue back to evaluation. Recurrent models can in principle keep running beyond their training context, but continued execution does not establish useful long-range memory. A speaker argues that architecture research outpaced the benchmarks needed to assess it, and would have introduced task-specific evaluations alongside the push for longer contexts in 2022.

The closing discussion reports that alternative mechanisms can remain stable beyond their training lengths while increasingly forgetting information. It also cites encouraging third-party time-series and weather-modeling results at controlled parameter sizes, without naming studies or supplying measurements.

Weather forecasting makes the application distinction concrete. A continuously updated state may help predict tomorrow’s weather without preserving exact recall of events years earlier. Predicting what comes next and recalling a specific past event place different demands on memory. The useful architecture depends on which question the system must answer.

Latent Space

Join Latent Space.

Sign in with Google, or get the newsletter by email.

Loading Google sign-in…

or subscribe by email

Free. Unsubscribe any time.