---
title: 'OjaKV: Online Low-Rank KV Cache Compression'
url: https://www.emergentmind.com/topics/ojakv
type: topic
---

# OjaKV: Online Low-Rank KV Cache Compression

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 large language models 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 [2509.21623].

## 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 [2509.21623].

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 [2509.21623]. 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 [2509.21623].

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 [2509.21623].

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 \(U_v \in \mathbb{R}^{d_h \times r_v}\) satisfying
\[
U_k^T U_k = I_{r_k}, \qquad U_v^T U_v = I_{r_v}.
\]

Compressed features are defined as
\[
\tilde{Q} = Q U_k \in \mathbb{R}^{m \times r_k}, \qquad
\tilde{K} = K U_k \in \mathbb{R}^{n \times r_k}, \qquad
\tilde{V} = V U_v \in \mathbb{R}^{n \times r_v}.
\]

The paper presents two equivalent attention computation paths. In the low-rank kernel, attention is computed in the reduced space and then expanded:
\[
\tilde{A} = \operatorname{softmax}\!\left(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d_h}}\right), \qquad
\tilde{O} = \tilde{A}\tilde{V}, \qquad
\hat{O} = \tilde{O} U_v^T.
\]

In the FlashAttention-compatible path, compressed tokens are reconstructed to full rank on the fly:
\[
\hat{K} = \tilde{K} U_k^T, \qquad
\hat{V} = \tilde{V} U_v^T,
\]
and attention is then computed as
\[
\hat{O} = \operatorname{softmax}\!\left(\frac{Q\hat{K}^T}{\sqrt{d_h}}\right)\hat{V}.
\]

The equivalence relation stated in the paper is
\[
Q\hat{K}^T = \tilde{Q}\tilde{K}^T,
\]
with output equivalence between the FlashAttention-compatible path and \(\tilde{O}U_v^T\) from the low-rank kernel path [2509.21623]. The paper further emphasizes that using \(\sqrt{d_h}\), rather than \(\sqrt{r_k}\), preserves the effective temperature of the original head; changing the scaling to \(\sqrt{r_k}\) 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 \(U_k\) and \(U_v\) from token activations during both prompt prefilling and decoding [2509.21623].

For prefilling, with learning rate \(\eta_{\text{pre}}\), the paper uses a batch Oja update over selected key vectors:
\[
\tilde{K} = U_k^T K, \qquad
U_k \leftarrow U_k + \eta_{\text{pre}} (K - U_k \tilde{K}) \tilde{K}^T,
\]
and analogously for values,
\[
\tilde{V} = U_v^T V, \qquad
U_v \leftarrow U_v + \eta_{\text{pre}} (V - U_v \tilde{V}) \tilde{V}^T.
\]
The bases are orthonormalized after the update.

During decoding, with learning rate \(\eta_{\text{dec}}\), new \((k_t, v_t)\) pairs are buffered and the same form of update is applied periodically every \(T\) steps:
\[
U_k \leftarrow U_k + \eta_{\text{dec}} (K - U_k \tilde{K}) \tilde{K}^T, \qquad
U_v \leftarrow U_v + \eta_{\text{dec}} (V - U_v \tilde{V}) \tilde{V}^T,
\]
followed by orthonormalization and buffer clearing [2509.21623].

The paper connects this matrix-batch update to classical multi-component Oja updates. In its formulation, the batch variant
\[
U \leftarrow U + \eta (K - U U^T K)(U^T K)^T
\]
is presented as the natural extension over a set of samples \(K\), 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 [2509.21623].

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 \(i\), the method collects \(Q\), \(K\), and \(V\) activations from \(n_s\) sequences of length \(n\), builds matrices \(R_i^Q\), \(R_i^K\), \(R_i^V \in \mathbb{R}^{(n_s \cdot n)\times d_h}\), and concatenates \(R_i^{KQ} = [R_i^Q, R_i^K] \in \mathbb{R}^{(n_s \cdot n)\times 2d_h}\). A compact SVD is then computed,
\[
R_i^{KQ} = U \Sigma V^T,
\]
and the minimal rank \(r\) is chosen to satisfy the energy criterion
\[
\frac{\|(R_i^{KQ})_r\|_F^2}{\|R_i^{KQ}\|_F^2} \ge \epsilon_{\text{th}}.
\]
\(U_k\) is set to the top-\(r\) columns of \(U\), while \(U_v\) is obtained from an SVD of \(R_i^V\). The method maintains per-head bases but enforces a consistent effective rank across heads in a layer, set to the maximum \(r\) observed [2509.21623].

The prefill-stage update uses importance scoring over the prompt, inspired by SnapKV, based on the last \(w\) queries. The method selects a salient set \(S_{\text{imp}} = \mathrm{TopK}(s)\), with top-\(k\) fraction \(k_{\text{pre}} = 0.05n\), applies a single relatively high-\(\eta_{\text{pre}}\) batch Oja update over the corresponding key and value matrices, orthonormalizes the bases, and then marks anchors [2509.21623].

During decoding, the default periodic update interval is \(T = 32\), with \(\eta_{\text{pre}} = 0.10\), \(\eta_{\text{dec}} = 0.05\), \(w = 32\), \(n_{\text{start}} = 32\), and \(n_{\text{recent}} = 32\). Non-anchor writes store
\[
\tilde{K}_t = K_t U_k, \qquad \tilde{V}_t = V_t U_v,
\]
whereas anchor writes store full \(K_t, V_t\) [2509.21623].

For orthonormalization, the paper permits QR factorization or Gram–Schmidt on \(d_h \times r\) matrices, with per-head cost \(O(d_h r^2)\). Per update and per head, for buffer size \(s\), the paper gives the following complexity terms: computing \(\tilde{K} = U^T K\) costs \(O(d_h r s)\), computing \((K - U\tilde{K})\) costs \(O(d_h r s)\), the rank-\(r\) update \((K - U\tilde{K})\tilde{K}^T\) costs \(O(d_h r s)\), and orthonormalization costs \(O(d_h r^2)\) [2509.21623].

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 \(n\), 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 [2509.21623].

## 6. Memory model, empirical findings, and limitations

The paper gives explicit KV memory formulas. Per head and per token, full rank stores \(2d_h\) scalars for key and value, whereas low rank stores \(r_k + r_v\) scalars. For batch size \(B\), sequence length \(T\), layers \(L\), KV heads \(H_{kv}\), and bytes per scalar \(b\), full-rank KV memory is
\[
M_{\text{full}} = BTLH_{kv}(2d_h)b.
\]
Without anchors, low-rank memory is
\[
M_{\text{low}} = BTLH_{kv}(r_k + r_v)b,
\]
with saving
\[
1 - \frac{r_k + r_v}{2d_h} = 1 - \frac{r}{d_h}
\quad \text{when } r_k = r_v = r.
\]
With anchors \(A = n_{\text{start}} + n_{\text{recent}}\), the paper gives
\[
M_{\text{OjaKV}} = BLH_{kv}\big[2Ad_h + (T-A)(r_k + r_v)\big]b.
\]
Anchors therefore modestly reduce the compression ratio while preserving exact representation where the model is most sensitive [2509.21623].

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 [2509.21623].

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 [2509.21623]. 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 [2509.21623].

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 [2509.21623]. 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 [2509.21623].

Several limitations and edge cases are identified. If context shifts are extremely rapid and nonstationary, online updates may lag, especially with large \(T\) or conservative \(\eta_{\text{dec}}\). Excessive compression with very small \(r\) 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 \(\eta_{\text{pre}}\), \(\eta_{\text{dec}}\), \(T\), \(r\), and \(A\) require tuning for different models and tasks, and poor settings may underperform static methods [2509.21623].

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 [2509.21623].

Source: https://www.emergentmind.com/topics/ojakv