🤖 machine learning

Online Vector Quantized Attention

This paper introduces Online Vector-Quantized (OVQ) attention, a sequence mixing layer that achieves linear compute and constant memory costs while significantly improving long-context performance through sparse memory updates, offering a competitive alternative to self-attention with a fraction of the memory usage.

Original authors: Nick Alonso, Tomas Figliolia, Beren Millidge

Published 2026-05-18

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Nick Alonso, Tomas Figliolia, Beren Millidge

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

The Problem: The "Too Big" and the "Too Small" Brain

Imagine you are trying to read a massive, 100,000-page novel to answer a specific question about a character mentioned on page 5.

The Old Way (Self-Attention): This is like having a super-intelligent assistant who reads the entire book every time you ask a question. They remember every single word perfectly. However, to do this, they need a library the size of a city to store all those notes. As the book gets longer, the library gets bigger and the assistant gets slower. It's accurate, but it's expensive and slow.
The Efficient Way (Linear Attention/SSMs): This is like an assistant who only keeps a tiny notepad. They summarize the book as they read it, keeping only the most important stats. They are incredibly fast and only need a pocket-sized notepad. But, because the notepad is so small, they often forget the specific details needed for long, complex stories. They struggle to find that character mentioned on page 5 when the book is 100,000 pages long.

The Goal: The authors wanted to build an assistant that is as fast and memory-efficient as the "pocket notepad" guy but as smart and detailed as the "city library" guy.

The Solution: The "Smart Filing Cabinet" (OVQ-Attention)

The authors created a new method called Online Vector Quantized (OVQ) Attention. Think of it as a Smart Filing Cabinet that updates itself in real-time.

Here is how it works, step-by-step:

1. The Dictionary (The Folders)

Instead of remembering every single word (like the city library) or just a summary (like the pocket notepad), this system uses a Dictionary of Folders.

Imagine you have a cabinet with, say, 4,000 empty folders.
As the assistant reads the book, they don't write down every word. Instead, they look at the current sentence and ask: "Which of these 4,000 folders does this sentence fit into best?"
They drop the sentence into that folder.

2. The "Online" Magic (Updating the Folders)

In previous versions of this idea, the folders were pre-written before the assistant started reading. If the book used words the folders didn't expect, the assistant got confused.

In OVQ-Attention, the folders are empty at the start. As the assistant reads the book:

They create new folders on the fly if they see something unique.
They update the existing folders to better match what they are seeing right now.
Crucially: They only update the specific folder that the current sentence belongs to. They don't rewrite the whole cabinet. This keeps the work fast (linear) and the memory usage low (constant).

3. The "Sparse" Update (The Efficient Trick)

Usually, if you want to remember more details, you need a bigger cabinet, which takes up more space.

The Paper's Trick: The authors realized that even if you have a cabinet with 100,000 folders, you only ever touch a few of them at once.
Because the updates are sparse (you only touch the specific folder relevant to the current sentence), the effort to update the cabinet doesn't grow just because the cabinet is huge.
Result: You can have a massive cabinet (high memory capacity) without paying the "effort tax" of a massive cabinet. You get the best of both worlds: huge storage, low cost.

The Results: How Well Did It Work?

The authors tested this "Smart Filing Cabinet" on several challenges:

The "Needle in a Haystack" Test (In-Context Recall):
- The Task: Hide a specific key-value pair (like a phone number) in a huge block of text and ask the model to find it later.
- The Result: The old "pocket notepad" models failed miserably after a certain length. The "city library" models worked perfectly but were slow. The OVQ-attention model worked almost as perfectly as the city library, even with a much smaller memory footprint, and it could handle texts up to 64,000 words long.
The "Learning While Reading" Test (In-Context Learning):
- The Task: Give the model a bunch of examples of a new rule (e.g., "If you see X, do Y") and ask it to apply that rule to new sentences later in the text.
- The Result: The OVQ model learned the rules just as well as the heavy, slow models, while the efficient-but-dumb models failed to learn the complex patterns.
Reading Long Stories (Language Modeling):
- When asked to predict the next word in a long story (using the PG19 dataset), the OVQ model performed competitively with the best standard models, despite using a fraction of the memory.

The Secret Sauce: Gaussian Mixture Regression

The paper explains the math behind this using a concept called Gaussian Mixture Regression.

Simple Analogy: Imagine you are trying to guess the weather based on a cloud.
- Standard models try to match the cloud to every past cloud they've ever seen.
- OVQ-Attention groups clouds into "types" (e.g., "Storm Cloud," "Fluffy Cloud").
- As new clouds appear, the system doesn't just memorize them; it adjusts the definition of "Storm Cloud" to fit the new data better. It learns the shape of the cloud types while it reads, making its guesses much more accurate over long periods.

Summary

The paper introduces OVQ-Attention, a new way for AI to process long texts.

Old Efficient Models: Fast and small, but forgetful.
Old Powerful Models: Smart and detailed, but slow and memory-hungry.
OVQ-Attention: Uses a dynamic, self-updating filing system. It keeps a small, constant amount of active memory but can store a massive amount of information by organizing it into clusters. It learns these clusters while it reads, allowing it to be both fast and incredibly smart over very long contexts (up to 64k tokens).

The authors conclude that this is a major step forward in making AI models that are both efficient and capable of handling long, complex tasks without needing super-computers.

Technical Summary: Online Vector Quantized Attention

1. Problem Statement

Standard sequence mixing layers in Large Language Models (LLMs) face a fundamental trade-off between computational efficiency and long-context performance.

Self-Attention: While highly effective for long-context tasks (in-context learning, reasoning, retrieval), it incurs quadratic computational costs ( $O(T^2)$ ) and linearly growing memory costs ( $O(T)$ ), making it expensive for very long sequences.
Linear Attention and State Space Models (SSMs): These offer linear computational costs ( $O(T)$ ) and constant memory ( $O(1)$ ). However, they suffer from significant performance deficits on tasks requiring precise long-range information processing due to their limited memory capacity.
Vector Quantized (VQ) Attention: A prior approach (Lingle, 2023) that achieves linear compute and constant memory by quantizing keys into a fixed dictionary of centroids. However, preliminary testing revealed that the original VQ-attention struggles with long-context tasks. The primary bottleneck was identified as the static, pre-trained key dictionary ( $D_k$ ), which fails to model keys effectively during the forward pass, leading to quantization errors that degrade performance. Increasing the dictionary size provided diminishing returns.

2. Methodology: Online Vector Quantized (OVQ) Attention

The authors propose Online Vector Quantized (OVQ) Attention, a novel sequence mixing layer that learns both the key dictionary ( $D_k$ ) and value dictionary ( $D_v$ ) online during the forward pass.

Theoretical Basis: Gaussian Mixture Regression (GMR)

The paper establishes a theoretical link between OVQ-attention and Gaussian Mixture Regression (GMR).

Self-Attention as GKR: Standard self-attention is equivalent to Gaussian Kernel Regression (GKR) under specific assumptions.
VQ-Attention as GMR: The authors demonstrate that VQ-attention (and by extension, OVQ) is more closely related to GMR. In GMR, a Gaussian Mixture Model (GMM) is fitted to a dataset of input-output pairs (keys and values).
The Connection: The prediction formula for GMR, under assumptions of unit-norm queries/keys and isotropic covariance, simplifies to the linear-time VQ-attention equation:
$O_{ovq} = \text{softmax}(\beta Q D_k^\top + \log(c)) D_v$
where $c$ represents the count of keys assigned to each centroid.
Implication: Unlike standard VQ-attention which learns $D_k$ and $D_v$ separately (pre-training vs. online), GMR learns them simultaneously to fit the same dataset. This motivates the OVQ approach of learning both dictionaries online to accurately model the conditional distribution $P(V|K)$ .

Online Learning Formulation

OVQ-attention operates within a test-time training framework using a chunk-parallel formulation:

Prediction: For each chunk of queries, the model computes attention using the current state of the dictionaries ( $D_k, D_v$ ) and count vectors ( $c$ ). Unlike the original VQ-attention, the current chunk's keys and values are not quantized before being used in the attention calculation; they are concatenated directly to the dictionaries.
State Update: After generating predictions, the model updates the dictionaries using the current chunk's keys and values.
- Dictionary Growth: The number of centroids ( $N_t$ ) grows asymptotically toward a maximum $N$ using a plateauing growth function inspired by the Dirichlet process. This allows the model to adapt its memory capacity dynamically.
- Centroid Initialization: New centroids are initialized using a "spread-maximizing" scheme (similar to k-means++), selecting data points with the smallest dot products to existing centroids to ensure diversity.
- Update Rule: Existing centroids are updated via an online version of the batch K-means update. The authors prove that under their assumptions, the first batch EM update on a GMM is equivalent to a Newton update on the Negative Log Likelihood (NLL), providing a second-order approximation for minimizing the loss.

Computational Properties

Memory: OVQ maintains constant memory ( $O(N)$ ) regardless of sequence length, as the state is defined by the dictionaries and counts, not the full sequence.
Compute: The layer incurs linear compute costs ( $O(T)$ ). The update operation is sparse; each token in a chunk updates only a single row of the dictionaries. This allows the state size ( $N$ ) to be increased significantly without increasing the memory footprint of the update tensor, unlike linear attention where state updates scale with the state dimension.

3. Key Contributions

Proposal of OVQ-Attention: A new sequence mixing layer that learns both key and value dictionaries online, overcoming the static dictionary limitation of prior VQ-attention.
Theoretical Framework: A novel derivation linking OVQ-attention to Gaussian Mixture Regression, providing a principled basis for the online learning algorithm, including dictionary growth and sparse update rules.
Empirical Validation: Comprehensive testing on synthetic tasks (In-Context Recall, In-Context Learning) and long-context language modeling (PG19).

4. Experimental Results

The authors evaluated OVQ-attention against strong baselines, including full self-attention, linear attention, and SSMs (e.g., Mamba, Gated Delta Net).

Synthetic In-Context Recall (ICR):
- OVQ-attention significantly outperformed original VQ-attention and linear attention baselines on "Basic ICR" and "Positional ICR" tasks.
- It matched or slightly deviated from full self-attention baselines up to 64k sequence lengths, despite using a state size only 10-25% of that used by full self-attention.
- Performance improved with larger dictionary sizes at test time, demonstrating effective length extrapolation.
Synthetic In-Context Learning (ICL):
- On a long-context linear regression task with multiple functions, OVQ-attention successfully learned all functions in scenarios with up to 128 distinct functions, matching the performance of full self-attention.
- Original VQ-attention and linear attention baselines failed to learn even a single function in these difficult long-context scenarios.
Long-Context Language Modeling (PG19):
- Hybrid architectures interleaving OVQ with sliding window layers drastically improved performance over pure sliding window or Gated Delta Net models.
- OVQ-attention matched or slightly surpassed standard self-attention baselines on the PG19 dataset at 16k context lengths.
Short-Context Benchmarks:
- On standard short-context benchmarks (PIQA, HellaSwag, etc.), OVQ-attention performed comparably to standard self-attention, confirming it does not degrade performance on tasks where compression is not required.

5. Significance and Claims

The paper claims that OVQ-attention represents a significant step forward in balancing efficiency and performance for long-context processing.

Efficiency vs. Capacity: Unlike linear attention and SSMs, OVQ-attention uses sparse memory updates that are independent of the total state size. This allows the model to utilize very large memory states (high capacity) while maintaining constant memory footprint and linear compute costs.
Continual Learning Potential: The authors suggest that the sparse, online clustering nature of OVQ-attention makes it inherently more resilient to "catastrophic forgetting" compared to dense gradient updates used in linear attention/SSMs. This positions OVQ as a promising candidate for continual learning scenarios where data distributions shift over time.
Scalability: While current experiments are at a relatively small scale (<500M parameters), the architecture's theoretical properties suggest it is a viable path toward efficient, long-context LLMs without the quadratic costs of self-attention.

The paper concludes that OVQ-attention offers a compelling alternative to self-attention, achieving competitive performance on long-context tasks with a fraction of the memory requirements, marking an important research direction for efficient sequence modeling.