---
title: 'Thinformer Express: Streaming Causal Attention'
url: https://www.emergentmind.com/topics/thinformer-express
type: topic
---

# Thinformer Express: Streaming Causal Attention

Thinformer Express is a causal attention approximation algorithm that combines the state-of-the-art Thinformer thinning procedure with the Express meta-procedure for converting non-causal attention approximations into streaming, causal ones, while maintaining strong theoretical error guarantees. The method achieves uniform $\varepsilon = O(\log^{3/2} n / s)$ approximation error, $O(s)$ memory, and $O(s^2 \log^2 n)$ compression overhead for a sequence of length $n$. Its highly optimized implementation in Triton delivers significant speedups compared to prior state-of-the-art methods such as FlashAttention 2 and HyperAttention, supporting efficient deployment for long-context prefill, KV-cache compression, and memory-/compute-constrained decoding workloads [2606.10944].

## 1. Problem Statement and Theoretical Foundations

Causal attention, foundational to sequence modeling, requires restricting each position $j$ to attend only to prior or current positions ($i\leq j$). For query–key–value triplets $(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$, the exact masked (causal) attention output for $q_j$ is
\[
o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.
\]
Exact computation requires $\Theta(d n^2)$ time and $\Theta(d n)$ memory, making scaling to long sequences intractable for standard hardware constraints.

Coreset-based attention approximation addresses this by maintaining a compressed, weighted subset of key–value pairs, enabling an approximate attention output
\[
\hat o_j = \frac{\sum_{(k, w)} w\,\exp(q_j^\top k/\sqrt{d})\,v}{\sum_{(k, w)} w\,\exp(q_j^\top k/\sqrt{d})}
\]
with $\| \hat o_j - o_j \|$ provably small for all $j$.

The key theoretical device is *sub-Gaussian thinning*: a randomized algorithm produces a weighted coreset $\{(x, w)\}$ with guarantees
\[
\exp\left(\frac{1}{n}\sum_{i=1}^n f(x_i) - \sum_{(x, w)} w f(x)\right) \leq \exp\left(\frac{1}{2}\sigma^2 \|f\|_H^2\right)
\]
for all test functions $f$ in a suitable RKHS, with high probability. Thinformer provides non-causal thinning with $\sigma = O(\sqrt{\log n} / s)$. The Express meta-procedure converts such thinning to the causal (streaming) domain with only an $O(\log n)$ inflation in $\sigma$, yielding Thinformer Express.

## 2. Algorithmic Structure and Execution Phases

Thinformer Express maintains a weighted cache $C$ of target size $s$, supporting the following workflow for each token:
- **Coreset retrieval**: Obtain $\{(k_i, v_i, w_i)\} = C.\mathrm{wtd()}$.
- **Approximate attention formation**: Compute
  \[
  \hat o_j = \frac{\sum_i w_i\,\exp(q_j^\top k_i/\sqrt{d})\,v_i}{\sum_i w_i\,\exp(q_j^\top k_i/\sqrt{d})}.
  \]
- **Update**: Call $C.\mathrm{update}(k_j, v_j)$.

Express operates internally in three phases:
1. **Exact**: Accumulates the first $s$ pairs exactly.
2. **Thin**: Processes subsequent input in blocks ($2^m$ in size), performing stratified subsampling followed by ThinformerHalve to reduce block size to $s$; summaries are appended to cache.
3. **Halve**: When cache exceeds $4s$, recursively apply halving (down to $s$ points), incrementing $m$.

This scheme guarantees at most $6s$ points in cache and streaming update complexity $O(\log^2 n)$ halving calls per token. Thinformer’s kernel, quadratic in block size $O(s)$, leads to total compression cost $O(d s^2 \log n \log (n/s))$.

## 3. Approximation Guarantees

The approximation error for Thinformer Express is rigorously characterized. Fix $\delta>0$ and run Thinformer Express with cache size $s$. With probability at least $1-\delta$, for every $1\leq j\leq n$,
\[
\|\hat o_j - o_j\| \leq C\,\exp\left(\frac{2R^2}{\sqrt{d}}\right)\sqrt{\ln\left(\frac{(d+1)n}{\delta'}\right)}\,\frac{\log^{3/2} n}{s}\|v\|_{\max} 
\]
where $R = \max_i \|q_i\|, \|k_i\|$, $\|v\|_{\max} = \max_i \|v_i\|$, $\delta' \approx 1/(4\ln n)$, and $C$ is a small absolute constant. In big-O notation,
\[
\varepsilon = \max_j \|\hat o_j - o_j\| = O\left(\frac{\log^{3/2} n}{s}\right)
\]
uniformly over all tokens. This guarantee arises by chaining sub-Gaussian error bounds across $O(\log n)$ thinning stages and subsequently applying an attention-output stability lemma [2606.10944].

## 4. Resource Complexity and Asymptotics

Thinformer Express achieves asymptotically optimal scaling for memory and runtime overheads in attention computation.

| Resource             | Scaling                        | Notes                          |
|----------------------|-------------------------------|--------------------------------|
| Memory (KV storage)  | $O(s)$                        | At most $6s$ weighted points   |
| Per-token query      | $O(d\,s)$                     | Full attention over cache      |
| Compression overhead | $O(d\,s^2\,\log^2 n)$         | Dominated by $O(\log n)$ small halving calls per token |

Total streaming compression time over $n$ tokens satisfies $O(d\,n\,s^2\,\log(n/s))$. Memory is independent of sequence length $n$; only computational overhead grows slowly as $\log^2 n$.

## 5. Implementation Optimizations

The combination of Triton and algorithmic design underpins Thinformer Express’s efficiency. Major optimizations include:
- **Tiling**: Key–query exponentiation and halving inner loops are partitioned into $B \times T$ tiles to maximize on-chip cache utilization.
- **Fused operations**: \(\exp(q^\top k)\) multiplications are fused with summation, preventing formation of $s \times s$ or $s \times B$ matrices in high-bandwidth memory (HBM).
- **Parallelism**: Same-size halving tasks are executed in parallel during offline prefill phases.
- **Double indirection**: Keys and values remain contiguous in HBM, while coresets store references in shared memory, eliminating scatter/gather overhead.

Performance gains include:
- Unmasked prefill ($n=32$K, $s=256$): Torch-compiled Thinformer $15\times$, Triton Thinformer $27\times$ faster than FlashAttention 2.
- Masked prefill (ChatGLM2-6B-32K, $n=512$K, $s=512$ or 1024): Thinformer Express up to $82\times$ faster than FlashAttention 2. HyperAttention is out-of-memory above $128$K.

## 6. Empirical Evaluation and Benchmarks

Rigorous empirical evaluations span diverse large-language-model workloads:

(a) **Long-context prefill, masked** (ChatGLM2-6B-32K, $s=512$):

| Length $n$ | FlashAttn2 | HyperAttn | Thinf-Expr |
|------------|------------|-----------|------------|
| 64K        | 1×         | 5×        | 12×        |
| 128K       | 1×         | 10×       | 25×        |
| 256K       | 1×         | --        | 45×        |
| 384K       | 1×         | --        | 60×        |
| 512K       | 1×         | --        | 82×        |

(b) **KV-cache compression** (Llama 3.1 8B, LongBench-E): Wrapping compressors (SnapKV, StreamingLLM, PyramidKV) with Express yields $2\times$–$4\times$ reduction in attention time at preserved end-task accuracy.

(c) **Memory-constrained decoding** (MATH-500, DeepSeek-R1-Distill-LLama-8B):

| Cache size | Exact Acc. | Thinf-Expr Acc. |
|------------|------------|-----------------|
| 1000       | 28.4%      | 28.3%           |
| 2000       | 31.2%      | 31.1%           |
| 4000       | 34.7%      | 34.6%           |

| Cache size | Exact Mem. | Thinf-Expr Mem. |
|------------|------------|-----------------|
| 4000       | 100%       | 61%             |

(d) **Compute-constrained decoding** (same settings):

| Time  | Exact Acc. | Thinf-Expr Acc. |
|-------|------------|-----------------|
| 1.0×  | 34.7%      | 34.7%           |
| 0.75× | 28.4%      | 28.4%           |

| Time     | Exact Cost | Thinf-Expr Cost |
|----------|------------|-----------------|
| Decoding | 100%       | 56%             |

Thinf-Expr matches exact accuracy with $61\%$ of KV memory and $56\%$ of computational cost.

## 7. Deployment Considerations and Parameter Selection

- **Cache size $s$**: Set $s \approx (\log^{3/2} n) / \varepsilon$ to target desired maximum approximation error $\varepsilon$. In practice, $s=256$–2048 suffices for $n$ up to several hundred thousand.
- **Sequence length $n$**: Memory usage is $O(s)$, independent of $n$. The only $n$-dependence is in the compression overhead ($\log^2 n$ factor), negligible when $s \ll n$.
- **Hardware**: The Triton reference implementation utilizes on-chip shared memory and contiguous HBM allocation. The algorithm is AISA-friendly. Warp specialization and FP8 are not used but can be integrated.
- **Pipeline integration**:
  - **Prefill**: Thinformer Express is applied in all layers to compress long input contexts prior to decoding.
  - **KV-cache**: Generation-time storage only requires the $s$-sized Express cache, not the full $n$.
  - **Decoding**: Express.update is called only after computing $\hat o_j$; halving costs per-token ($O(\log n)$) are negligible compared to attention.

Thinformer Express thus provides an end-to-end solution for causal attention with proven error bounds, optimal memory scaling, and practical speedups for long-context and resource-constrained neural language modeling [2606.10944].

Source: https://www.emergentmind.com/topics/thinformer-express