OjaKV: Online Low-Rank KV Cache Compression

• OjaKV is a framework for compressing key-value caches in autoregressive models by preserving first and recent tokens at full rank and compressing intermediate tokens.
• It implements a hybrid storage policy that maintains anchor tokens for stable attention while applying low-rank approximations to reduce memory usage.
• The method employs online PCA via Oja’s rule to continuously update the compression subspace, ensuring alignment with evolving context during inference.

Searching arXiv for the OjaKV paper and closely related methods mentioned in the provided data so the article can include accurate arXiv citations. Searching arXiv for "OjaKV Context-Aware Online Low-Rank KV Cache Compression with Oja's Rule". OjaKV is a framework for key-value (KV) cache compression in autoregressive LLMs that combines a hybrid storage policy with online subspace adaptation based on Oja’s rule. It is designed for long-context inference, where the KV cache becomes the dominant memory cost as context grows, and it aims to reduce memory usage without requiring model fine-tuning. The method preserves the first and most recent tokens at full rank, compresses intermediate tokens in a low-rank form, and updates the compression subspace during inference so that it remains aligned with the evolving context (Zhu et al., 25 Sep 2025).

1. Problem setting and motivation

The paper positions the KV cache as the central systems bottleneck for long-context autoregressive inference. Per layer and per head, a model must retain keys and values for every processed token so that subsequent tokens can attend to them. This memory cost scales linearly with sequence length, batch size, layers, and heads, and therefore constrains long-context inference, throughput, and deployment on commodity GPUs (Zhu et al., 25 Sep 2025).

A concrete example given in the paper is that a Llama-3.1-8B model processing a 32K-token prompt at a batch size of 4 requires approximately 16 GB of KV cache memory, exceeding the model’s weights (Zhu et al., 25 Sep 2025). The stated motivation for KV compression is threefold: to fit long prompts without offloading or sharding, to preserve batch throughput by supporting more concurrent sequences, and to reduce cost in serving and edge scenarios while keeping accuracy and latency reasonable.

Within this setting, the paper treats low-rank projection as a promising direction but argues that static, offline-learned subspaces perform poorly under data distribution shifts. The failure mode is described in terms of misalignment between an offline calibration subspace and the actual activation covariance induced by topic changes, style changes, or task transitions within the same sequence. This motivates online subspace tracking during inference rather than a fixed compression basis.

2. Hybrid storage policy and anchor tokens

OjaKV’s first central design choice is a hybrid storage policy. The method assumes that not all tokens are equally compressible and therefore stores two classes of tokens at full rank: the first tokens, described as context roots, and the most recent tokens, described as near-future attention targets. These full-rank tokens act as anchor tokens that stabilize attention and preserve exact fidelity where models most rely on them (Zhu et al., 25 Sep 2025).

Intermediate tokens, which are typically the majority, are stored in low-rank form. The default anchor counts are static: $n_{\text{start}} = 32$ and $n_{\text{recent}} = 32$. Anchors are exempt from compression; their keys and values are stored in full rank, and read/write attention to them is unchanged. Non-anchor tokens store only compressed features (Zhu et al., 25 Sep 2025).

This anchor policy is significant because it separates attention preservation from compression efficiency. The paper explicitly attributes two roles to anchors: they guarantee exact attention bindings to the earliest and latest tokens, and they slightly reduce the compression ratio while stabilizing attention where it matters most. A plausible implication is that OjaKV is not merely a low-rank approximation scheme; it is also a selective fidelity-allocation mechanism over the token axis.

3. Mathematical formulation of compressed attention

Per attention head, the paper uses head dimension $d_h$, cached sequence length $n$ or total context $T$, current queries $m$, keys and values $K, V \in \mathbb{R}^{n \times d_h}$, queries $Q \in \mathbb{R}^{m \times d_h}$, and $H_{kv}$ KV heads in a layer. Each head maintains orthonormal bases $U_k \in \mathbb{R}^{d_h \times r_k}$ and $n_{\text{recent}} = 32$0 satisfying

$n_{\text{recent}} = 32$1

Compressed features are defined as

$n_{\text{recent}} = 32$2

The paper presents two equivalent attention computation paths. In the low-rank kernel, attention is computed in the reduced space and then expanded: $n_{\text{recent}} = 32$3

In the FlashAttention-compatible path, compressed tokens are reconstructed to full rank on the fly: $n_{\text{recent}} = 32$4 and attention is then computed as

$n_{\text{recent}} = 32$5

The equivalence relation stated in the paper is

$n_{\text{recent}} = 32$6

with output equivalence between the FlashAttention-compatible path and $n_{\text{recent}} = 32$7 from the low-rank kernel path (Zhu et al., 25 Sep 2025). The paper further emphasizes that using $n_{\text{recent}} = 32$8, rather than $n_{\text{recent}} = 32$9, preserves the effective temperature of the original head; changing the scaling to $d_h$0 would alter the temperature and typically requires calibration.

These equations define the core representational claim of OjaKV: exact attention behavior is preserved on anchors, while non-anchor tokens are stored compactly and either attended in reduced space or reconstructed for standard kernels.

4. Online PCA via Oja’s rule

The second central design choice is online adaptation of the low-rank subspace using Oja’s algorithm for online principal component analysis. The method updates $d_h$1 and $d_h$2 from token activations during both prompt prefilling and decoding (Zhu et al., 25 Sep 2025).

For prefilling, with learning rate $d_h$3, the paper uses a batch Oja update over selected key vectors: $d_h$4 and analogously for values,

$d_h$5

The bases are orthonormalized after the update.

During decoding, with learning rate $d_h$6, new $d_h$7 pairs are buffered and the same form of update is applied periodically every $d_h$8 steps: $d_h$9 followed by orthonormalization and buffer clearing (Zhu et al., 25 Sep 2025).

The paper connects this matrix-batch update to classical multi-component Oja updates. In its formulation, the batch variant

$n$0

is presented as the natural extension over a set of samples $n$1, with the same fixed points as the principal subspace under appropriate step sizes and periodic re-orthonormalization. Theoretical discussion in the paper states that Oja’s rule is a streaming PCA method that provably tracks leading eigenvectors of the data covariance under mild conditions, and that in KV compression the relevant data are key/value activations whose covariance shifts as topics, styles, and tasks change within a long context (Zhu et al., 25 Sep 2025).

The significance of this design is explicit: static subspaces learned offline become misaligned when principal angles to the current activation covariance grow, whereas online Oja updates reduce these angles over time. This suggests that OjaKV’s principal advantage is not low rank alone, but low rank with context tracking.

5. Initialization, inference workflow, and systems integration

The initialization procedure is based on per-head calibration. For each head $n$2, the method collects $n$3, $n$4, and $n$5 activations from $n$6 sequences of length $n$7, builds matrices $n$8, $n$9, $T$0, and concatenates $T$1. A compact SVD is then computed,

$T$2

and the minimal rank $T$3 is chosen to satisfy the energy criterion

$T$4

$T$5 is set to the top-$T$6 columns of $T$7, while $T$8 is obtained from an SVD of $T$9. The method maintains per-head bases but enforces a consistent effective rank across heads in a layer, set to the maximum $m$0 observed (Zhu et al., 25 Sep 2025).

The prefill-stage update uses importance scoring over the prompt, inspired by SnapKV, based on the last $m$1 queries. The method selects a salient set $m$2, with top-$m$3 fraction $m$4, applies a single relatively high-$m$5 batch Oja update over the corresponding key and value matrices, orthonormalizes the bases, and then marks anchors (Zhu et al., 25 Sep 2025).

During decoding, the default periodic update interval is $m$6, with $m$7, $m$8, $m$9, $K, V \in \mathbb{R}^{n \times d_h}$0, and $K, V \in \mathbb{R}^{n \times d_h}$1. Non-anchor writes store

$K, V \in \mathbb{R}^{n \times d_h}$2

whereas anchor writes store full $K, V \in \mathbb{R}^{n \times d_h}$3 (Zhu et al., 25 Sep 2025).

For orthonormalization, the paper permits QR factorization or Gram–Schmidt on $K, V \in \mathbb{R}^{n \times d_h}$4 matrices, with per-head cost $K, V \in \mathbb{R}^{n \times d_h}$5. Per update and per head, for buffer size $K, V \in \mathbb{R}^{n \times d_h}$6, the paper gives the following complexity terms: computing $K, V \in \mathbb{R}^{n \times d_h}$7 costs $K, V \in \mathbb{R}^{n \times d_h}$8, computing $K, V \in \mathbb{R}^{n \times d_h}$9 costs $Q \in \mathbb{R}^{m \times d_h}$0, the rank-$Q \in \mathbb{R}^{m \times d_h}$1 update $Q \in \mathbb{R}^{m \times d_h}$2 costs $Q \in \mathbb{R}^{m \times d_h}$3, and orthonormalization costs $Q \in \mathbb{R}^{m \times d_h}$4 (Zhu et al., 25 Sep 2025).

The systems claim is that OjaKV is fully compatible with FlashAttention. The practical note in the paper is that the FlashAttention-compatible path preserves existing kernel efficiency and masking logic. Reconstruction GEMMs are linear in $Q \in \mathbb{R}^{m \times d_h}$5, and memory savings remain because only compressed tensors are stored for non-anchors. The method is described as plug-and-play and does not require model fine-tuning (Zhu et al., 25 Sep 2025).

6. Memory model, empirical findings, and limitations

The paper gives explicit KV memory formulas. Per head and per token, full rank stores $Q \in \mathbb{R}^{m \times d_h}$6 scalars for key and value, whereas low rank stores $Q \in \mathbb{R}^{m \times d_h}$7 scalars. For batch size $Q \in \mathbb{R}^{m \times d_h}$8, sequence length $Q \in \mathbb{R}^{m \times d_h}$9, layers $H_{kv}$0, KV heads $H_{kv}$1, and bytes per scalar $H_{kv}$2, full-rank KV memory is

$H_{kv}$3

Without anchors, low-rank memory is

$H_{kv}$4

with saving

$H_{kv}$5

With anchors $H_{kv}$6, the paper gives

$H_{kv}$7

Anchors therefore modestly reduce the compression ratio while preserving exact representation where the model is most sensitive (Zhu et al., 25 Sep 2025).

The empirical evaluation is reported on a single NVIDIA H100 NVL with PyTorch 2.6.0, Transformers 4.44.0, FlashAttention 2.7.4.post1, and float16 precision. The models are Llama-2-7B, Llama-3.1-8B-Instruct, and LongChat-7B for RULER. Initial bases are calibrated on WikiText-2. Benchmarks include lm-eval-harness tasks such as PiQA, WinoGrande, and HellaSwag, LongBench, and RULER. The stated metrics are accuracy per benchmark protocol, KV memory in GB, and TTFT in ms (Zhu et al., 25 Sep 2025).

The paper’s high-level empirical conclusion is that OjaKV maintains or even improves zero-shot accuracy at high compression ratios on very long-context tasks, especially those that require complex reasoning, and that these gains come from online subspace tracking that follows context shifts (Zhu et al., 25 Sep 2025). A quantitative example on LongBench with Llama-3.1-8B reports the following averages relative to full KV memory: Full KV cache uses 100% memory with accuracy 53.0; SnapKV at 50% keep uses 50% memory with accuracy 52.66; OjaKV at 0.6× memory, approximately 1.67× compression, uses 60% memory with accuracy 43.13; and OjaKV + SnapKV uses 30% memory with accuracy 43.33. The paper interprets this as validation of orthogonality between feature-dimension compression and sequence-length compression (Zhu et al., 25 Sep 2025).

A qualitative case study on LongBench multi-news summarization contrasts OjaKV with a static PCA baseline at 0.6× compression. The static baseline is described as fixating on a secondary topic, Philadelphia, and omitting the primary Los Angeles events, whereas OjaKV’s online subspace adaptation preserves both Los Angeles and Philadelphia content by tracking new dominant information (Zhu et al., 25 Sep 2025). This is presented as a direct instance of subspace misalignment under context shift.

The paper also situates OjaKV against static or offline low-rank projections, token selection such as SnapKV, heavy hitter methods such as H2O, and various low-rank methods including Palu, Eigen Attention, MatryoshkaKV, and ReCalKV. Its stated distinctive contribution is online, plug-and-play subspace tracking at inference time without fine-tuning, together with the anchor-aware storage policy (Zhu et al., 25 Sep 2025).

Several limitations and edge cases are identified. If context shifts are extremely rapid and nonstationary, online updates may lag, especially with large $H_{kv}$8 or conservative $H_{kv}$9. Excessive compression with very small $U_k \in \mathbb{R}^{d_h \times r_k}$0 can degrade accuracy and exacerbate temperature mismatch unless scaling and calibration are handled carefully. Oja updates and orthonormalization add compute overhead, even if the paper characterizes this overhead as small. Hyperparameters $U_k \in \mathbb{R}^{d_h \times r_k}$1, $U_k \in \mathbb{R}^{d_h \times r_k}$2, $U_k \in \mathbb{R}^{d_h \times r_k}$3, $U_k \in \mathbb{R}^{d_h \times r_k}$4, and $U_k \in \mathbb{R}^{d_h \times r_k}$5 require tuning for different models and tasks, and poor settings may underperform static methods (Zhu et al., 25 Sep 2025).

Taken together, these results define OjaKV as a long-context KV-cache compression method whose novelty lies in combining exact preservation of first and recent tokens with online low-rank subspace adaptation. In the paper’s formulation, this combination explains why the method is particularly effective on long-context reasoning tasks: it avoids catastrophic forgetting of early context while continuing to align the compressed representation with newly dominant content (Zhu et al., 25 Sep 2025).