OK-DMD

OK-DMD: Online Koopman Dynamic Mode Decomposition for KV-Cache Eviction

OK-DMD offers a robust and mathematically grounded framework for query-agnostic KV-cache compression and eviction specifically designed for Long-Context Large Language Models (LLMs). This innovative approach treats the sequence of attention Key vectors as a physical trajectory within a semantic space, effectively tracking and preserving stable mathematical "attractors" of the text to mitigate catastrophic context forgetting during extended text generations.

The Core Concept: Language as a River

Traditional KV-cache eviction strategies, such as H2O or local windowing, treat attention keys as isolated, static vectors, relying heavily on rolling attention scores that are often volatile and susceptible to attention drift. This volatility means a token that may be overlooked now could become crucial later on as the context evolves.

OK-DMD redefines this paradigm by conceptualizing the flow of text as a river, where every new word is analogous to a drop of water moving with the current. Instead of focusing on isolated drops, OK-DMD estimates the underlying current (the semantic attractor) employing Koopman Operator Theory, which assumes that the transition from one Key state ($k_{t-1}$) to the next ($k_t$) can be represented locally by a linear transition operator $A_t$:

k_t \approx A_t k_{t-1}

By calculating the persistent eigenvalues of $A_t$ in real-time, OK-DMD discerns which historical tokens robustly represent the underlying structure of the document (stable modes) and which are merely transient noise (local syntax or conversational tangents) that can be safely evicted.

Domain Shift Resilience: The "Lasagna Trap"

Conventional eviction heuristics often experience catastrophic forgetting during domain shifts. For instance, if an LLM is engaged in a lengthy discussion about Astronomy, any sudden injection of a transient Lasagna Recipe causes a shift in attention entirely towards cooking terminology, leading to the purge of older Astronomy keys to stay within memory limits. When the conversation reverts to Astronomy, the model is left with gaps in its recall, potentially resulting in hallucinations.

OK-DMD counteracts this issue by tracking the spectral footprint (the persistent eigenvalues) of the global trajectory, ensuring that the "Astronomy" semantic subspace remains stored while transient noise is permitted to be evicted.

Quick Start & Benchmarking

A benchmark demonstrating the efficacy of OK-DMD under a multi-domain stress test (Astronomy $ ightarrow$ Lasagna $ ightarrow$ Astronomy Recall) can be executed to evaluate how OK-DMD preserves the global attention manifold compared to H2O:

python benchmarks/run_stress_test.py

The anticipated output, even under an aggressive 92.2% compression budget (restricting the KV-cache of Qwen-2.5-0.5B from 817 to just 64 active tokens), consistently shows OK-DMD outpacing traditional attention-accumulating methods:

------------------------------------------------------------
 FINAL PERFORMANCE REPORT: DOMAIN SHIFT STRESS TEST
 Total Sequence Length:    817 tokens
 KV-Cache Budget Limit:    64 tokens (Aggressive Compression)
------------------------------------------------------------
 -> Mean Attention Coherence (H2O):      0.997211
 -> Mean Attention Coherence (OK-DMD):   0.997972
------------------------------------------------------------
[SUCCESS] OK-DMD outperformed H2O! The dynamic attractor preserved the old Astronomy context.

Mathematical Architecture

OK-DMD is based on three integrated mathematical mechanisms:

1. Causal Recursive Update (O(d²))
   This method updates the Koopman matrix $A_t$ at each generation step without requiring costly matrix inversions:

   u_t = P_{t-1} k_{t-1}
   g_t = \frac{u_t}{\lambda + k_{t-1}^T u_t + \eta}
   P_t = \frac{P_{\text{next}} + P_{\text{next}}^T}{2} + \sigma I
   A_t = A_{t-1} + e_t g_t^T
   

2. Warm-Started QR Subspace Iteration
   This technique extracts the orthonormal basis $Q_t$ representing the most persistent Schur modes using QR iteration initialized by the previous step's subspace:

   Y = A_t Q_{t-1}
   Q_t, R = \text{QR}(Y)
   

3. Topological Eviction Scoring
   The survival score $S_i$ for any token in the cache represents the projection energy of its key onto the persistent subspace:

   S_i = \| Q_t^T k_i \|_2^2
   

Memory Break-Even Analysis

Maintaining state matrices like $P$ and $A$ adds a modest memory overhead of 41 MB per active sequence (for Llama-3-70B with 80 layers and 8 KV heads). The analysis determines the exact sequence length $N$ where OK-DMD starts to save VRAM under specific compression rates:

N > \frac{40,960 \text{ KB}}{(1 - C) \cdot 320 \text{ KB}}

• 50% Eviction: Efficiency is achieved for sequences longer than 256 tokens.
• 80% Eviction: Efficiency is achieved for sequences longer than 160 tokens.

Limitations

As a mathematically rigorous project, the main limitations of OK-DMD include:

• The Needle in a Haystack (NIAH) Constraint: OK-DMD is optimized to extract global semantic trends, and transient anomalies may be misclassified as noise and evicted.
• The Production Solution: In enterprise architectures, OK-DMD should not compress cache during prompt prefilling. Instead, prompts and final instructions should be parsed in parallel, using a Prompt-Aware filter to lock essential tokens. OK-DMD can then be employed during the long answer generation phase to prevent semantic drift.