📊 statistics

Express Language Modeling

The paper introduces Express, a tool that converts non-causal attention approximations into causal ones with improved theoretical guarantees and an efficient Triton implementation, enabling significant speedups over FlashAttention 2 and overcoming key resource bottlenecks in long-context language modeling.

Original authors: Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a very long book to answer a question. As you read, you need to remember every important detail you've encountered so far to make sense of the current sentence. In the world of AI language models, this "remembering" process is called Attention.

The problem is that as the book gets longer, the effort required to remember everything grows explosively. It's like trying to hold a conversation with a room full of people where, for every new word you speak, you have to re-introduce yourself to everyone in the room and re-read your entire conversation history. This gets so slow and memory-heavy that it becomes impossible to handle very long stories or complex reasoning tasks.

This paper introduces a new tool called Express (and a specific version called Thinformer Express) that acts like a super-efficient librarian for these AI models. Here is how it works, using simple analogies:

1. The Problem: The "Quadratic" Bottleneck

Normally, to understand a new sentence, an AI looks at every single previous word. If you have 1,000 words, it does 1,000 checks. If you have 1 million words, it has to do 1 trillion checks. This is the "quadratic" cost mentioned in the paper. It's like trying to find a specific needle in a haystack by checking every single piece of hay one by one, over and over again.

2. The Solution: The "Smart Summarizer" (Thinning)

The authors realized that you don't actually need to remember every word perfectly. You just need a small, high-quality "summary" or "core set" of the most important words that represents the whole story.

They built a tool called Express that takes a "non-causal" summarizer (a tool that can look at a whole book and pick the best 100 pages) and turns it into a "causal" summarizer.

Non-causal is like reading the whole book first, then picking the highlights.
Causal is like reading the book page-by-page in real-time, deciding instantly which highlights to keep in your pocket as you go, without ever looking ahead.

Express is the magic trick that allows the AI to do this real-time summarization without losing accuracy. It ensures that even though the AI is only looking at a tiny fraction of the past (a "thinned" version), it still gets the right answer.

3. The "Inflation" Trick

The paper describes a clever mechanism where the AI's memory "inflates" and "deflates" like a balloon.

The Exact Phase: At the very beginning, the AI remembers everything perfectly.
The Thin Phase: As more words come in, the AI groups them into batches. Instead of keeping all of them, it uses a special algorithm to "compress" the batch down to a smaller, weighted summary.
The HALVE Phase: When the summary gets too big, the AI performs a "halving" operation, cutting the size in half while carefully preserving the most important information.

This process ensures the memory size stays small and constant, no matter how long the story gets. It's like having a backpack that automatically shrinks its contents to fit, keeping only the most essential items, so you never run out of space.

4. Real-World Results: Speeding Up the AI

The authors didn't just do the math; they built a fast version of this tool using a programming language called Triton (which is like a high-speed engine for computer chips). They tested it in four specific scenarios:

Long-Context Prefill (Reading the Book): When the AI is given a massive document to read before answering, Express made it 82 times faster than the current best method (FlashAttention 2) for very long texts (512,000 words).
Compressing the Memory (The KV Cache): AI models store past words in a "cache." Express helped compress this cache using other popular methods, making the whole process faster without losing accuracy.
Memory-Efficient Decoding (Thinking Step-by-Step): When solving hard math problems that require long chains of reasoning, Express allowed the AI to use only 61% of the memory required by the standard method while getting the same correct answers.
Compute-Efficient Decoding (Speeding Up Thinking): For the same math problems, Express allowed the AI to finish in only 56% of the time it usually takes, while still getting the right answer.

Summary

In short, Express is a new way for AI to handle long conversations and complex tasks. It acts like a smart filter that constantly summarizes the past, keeping the AI's memory small and its thinking speed fast, without making the AI "forget" important details. It solves the problem of AI getting too slow or running out of memory when dealing with long texts or complex reasoning.

Technical Summary: Express Language Modeling

Problem Statement

Attention mechanisms in large language models (LLMs) scale quadratically with sequence length ( $n$ ), creating significant computational and memory bottlenecks, particularly for long-context prefilling, sequential decoding, and Key-Value (KV) cache management. While state-of-the-art sub-quadratic attention approximations exist (e.g., Thinformer [4]), they primarily address the unmasked attention problem. The more critical challenge for language modeling is causal masking, where each query must attend only to an evolving context of prior key-value pairs. Existing causal approximations either lack strong theoretical guarantees, suffer from high memory overhead, or fail to provide efficient streaming updates suitable for the incremental nature of token generation.

Methodology

The EXPRESS Meta-Procedure

The authors introduce EXPRESS, a meta-procedure designed to convert any high-quality, non-causal (unmasked) attention approximation based on thinning (subselecting input points) into a high-quality, causal (masked) approximation with streaming guarantees.

EXPRESS operates as an updatable weighted cache object initialized with a target cache size $n_{out}$ , an inflation factor $m$ , and a base halving algorithm (e.g., Kernel Halving). It manages the cache through three distinct phases:

Exact Phase: The first $n_{out}$ points are stored directly without compression.
Thin Phase: Incoming points are batched and incrementally thinned down to size $n_{out}$ using a combination of uniform subsampling and a generalized compression algorithm (COMPRESS2) before entering the main cache.
HALVE Phase: When the cache size reaches a threshold ( $4n_{out}$ ), the entire cache is compressed back to $n_{out}$ by applying the base halving algorithm twice, and the inflation factor $m$ is increased.

This structure ensures that the cache size remains bounded (at most $6n_{out}$ ) regardless of the sequence length $n$ , while maintaining sub-Gaussian approximation guarantees.

Thinformer Express

The authors instantiate EXPRESS with Kernel Halving (KH), the state-of-the-art non-causal thinning algorithm, to create Thinformer Express. This system utilizes the key-value attention kernel $\kappa_{att}$ to ensure the approximation preserves the specific averages required for attention computation.

Key technical components include:

Streaming Guarantees: Theoretical proofs demonstrate that Thinformer Express achieves an approximation error of $O(\log^{3/2}(n)/s)$ with only $O(s)$ memory and $O(s^2 \log^2(n))$ compression overhead for a sequence of length $n$ and cache size $s$ .
I/O-Aware Implementation: Inspired by FlashAttention 2, the authors implemented Thinformer Express using Triton kernels. This implementation utilizes tiling to minimize high-bandwidth memory (HBM) reads/writes, parallelizes across batch and head dimensions, and avoids materializing intermediate kernel matrices.

Key Contributions

Causal Conversion Framework: The introduction of EXPRESS, which bridges the gap between unmasked and masked attention approximations, transforming offline thinning algorithms into efficient, streaming causal approximations.
Improved Theoretical Bounds: Thinformer Express provides tighter approximation guarantees than existing causal methods (e.g., HyperAttention [15], BalanceKV [19]). Specifically, it offers a faster error decay rate ( $n^{-a}$ vs. $n^{-a/6}$ ) and improved dependence on the value matrix norm.
Efficient Implementation: A Triton-based implementation that maximizes GPU utilization, demonstrating substantial speedups over both exact attention and existing approximate methods.
Comprehensive Evaluation: Empirical validation across four critical resource-constrained scenarios in the language modeling pipeline.

Experimental Results

The authors evaluated Thinformer Express against FlashAttention 2 and other leading KV cache compression methods (SnapKV, StreamingLLM, PyramidKV, etc.) on standard benchmarks:

Long-Context Prefill:
- At a sequence length of 512k tokens, Thinformer Express achieved an 82× speedup over FlashAttention 2.
- It uniformly improved both perplexity and speedup compared to HyperAttention.
- Unlike HyperAttention, which ran out of memory at 512k tokens on an A6000 GPU, Thinformer Express completed the task.
Accelerated KV Cache Compression:
- When paired with SnapKV, StreamingLLM, and PyramidKV, Thinformer Express reduced total attention runtime on LongBench-E tasks without sacrificing accuracy.
- It effectively decouples the quadratic cost of generating the initial keys/values from the compression step.
Memory-Efficient Long-Form Decoding:
- On MATH-500 reasoning problems, Thinformer Express dominated the accuracy-vs.-cache-size trade-off.
- It matched the accuracy of exact attention while using only 61% of the cache size.
Accelerated Long-Form Decoding:
- On the same MATH-500 tasks, it dominated the accuracy-vs.-runtime trade-off.
- It matched exact attention accuracy with only 56% of the computation time.

Significance and Claims

The paper claims that Thinformer Express represents a significant advancement in making long-context language modeling feasible under resource constraints. By providing per-token accuracy guarantees, constant memory usage (independent of sequence length), and sub-quadratic query time, it addresses the primary bottlenecks in LLM inference: long-context prefilling, KV cache management, and memory/compute-constrained decoding.

The authors emphasize that their approach is not merely a heuristic but is grounded in rigorous sub-Gaussian thinning theory, offering stronger theoretical guarantees than prior causal approximations. Furthermore, the efficient Triton implementation demonstrates that these theoretical gains translate into real-world performance improvements, enabling the processing of sequences up to 512k tokens with substantial speedups over the current state-of-the-art (FlashAttention 2).

The paper concludes by noting limitations, including the current focus on English and mathematical reasoning, and the potential for further optimization using newer GPU features (e.g., Tensor Memory Accelerator, FP8 support) not yet exploited in their implementation.