Biological Annotation-Guided KV Cache Eviction for Protein Language Models

BACE (Biological Annotation-Guided Cache Eviction), reduces the key-value (KV) cache used during autoregressive protein generation. It trains a lightweight 24,066-parameter PPO-inspired Actor-Critic ranking policy to decide which ProtGPT2 cache tokens can be evicted while preserving biologically important residue positions derived from RefSeq records and AlphaFold confidence annotations.

🔬 Overview

Autoregressive protein language models keep a KV cache for previously generated tokens, and that cache grows with sequence length, batch size, layers, and attention heads. Generic cache-reduction methods use attention or redundancy signals, but they do not know whether a cached token corresponds to an active site, interaction interface, disease-associated variant, domain boundary, or other biologically meaningful position. BACE adds biological context to cache eviction by representing each candidate token with live attention/cache features plus reference-mapped biological annotations, then learning a PPO-inspired ranking policy that balances cache reduction, biological preservation, and sequence quality.

Abstract

Autoregressive protein language models are promising for computational protein generation and design, but their inference is bottlenecked by the key-value (KV) cache. Existing cache-compression methods, including H2O, ScissorHands, and SnapKV, use attention or redundancy heuristics and do not model the biochemical importance of residues. We propose Biological Annotation-Guided Cache Eviction (BACE), a lightweight Actor- Critic ranking policy trained with a proximal policy optimization (PPO)-inspired objective. BACE makes per-token eviction decisions from a 25-dimensional state vector combining 14 live attention/cache statistics with 11 protein residue-level biological features derived from National Center for Biotechnology Infor- mation (NCBI) GenBank annotations and AlphaFold predicted Local Distance Difference Test (pLDDT) scores. We train on 1,442 feature-complete proteins with NCBI RefSeq records and matched AlphaFold confidence scores, split into 1,009 training, 216 validation, and 217 held-out test proteins. We report two operating points. The 25% operating point is safety-first, achieving positive reward, 100% Active Site Preservation Rate (ASPR), and zero critical evictions, with lower cache reduction. The 40% operating point is high-compression, removing 37.19% of KV-cache tokens while maintaining 99.95% ASPR. In the 40% benchmark, BACE evicts only one critical residue, while random eviction, H2O, ScissorHands, and SnapKV evict 12 to 347 critical residues. At the same operating point, BACE achieves the highest relative reward, reduces reward damage by 73.9–83.7%, and has the lowest perplexity among compressed-cache methods. Feature-sensitivity studies show separable contributions: removing biological features increases critical evictions from 1 to 13, while removing AlphaFold pLDDT increases perplexity (PPL) from 224.74 to 1405.65. These results demonstrate the potential of annotation-guided KV cache eviction for compressed protein language model inference and highlight the need to evaluate cache reduction jointly with sequence fluency and biological preservation.

✨ Key Features

• Biological annotation-aware cache eviction: Combines attention/cache state with RefSeq and AlphaFold-derived biological annotations.
• Lightweight PPO-inspired policy training: Learns eviction behavior from rollout rewards rather than fixed hand-written protection rules.
• Multiple cache budgets: Trains and evaluates separate policies at 15%, 20%, 25%, 30%, 35%, and 40% target eviction.
• Held-out benchmarking: Compares BACE against full cache, random eviction, H2O, ScissorHands, and SnapKV on held-out test proteins.
• Feature-sensitivity analysis: Measures the role of biological features, AlphaFold pLDDT, sequence context, domain features, interaction features, and attention/cache features.
• Runtime profiling: Includes throughput, latency, cache-reduction, and GPU-memory profiling artifacts.

🧭 Architecture

flowchart TD
    A[RefSeq NP/NM records] --> B[Protein feature construction]
    C[AlphaFold pLDDT] --> B
    B --> D[Feature-complete protein splits]
    D --> E[ProtGPT2 frozen generator]
    E --> F[Live KV-cache and attention state]
    B --> G[Biological token features]
    F --> H[BACE token-state matrix]
    G --> H
    H --> I[PPO-inspired Actor-Critic policy]
    I --> J[Top-k eviction actions]
    J --> K[Autoregressive continuation]
    K --> L[Reward: biology + PPL + cache efficiency]
    L --> I
    I --> M[Validation-selected checkpoints]
    M --> N[Held-out benchmark vs cache baselines]
    N --> O[Tables, ablations, runtime profile]

Loading

📝 Citation

If you use BACE in your research, please cite:

BACE System:

@article{gelalbace_2026,
  title={Biological Annotation-Guided Key-Value Cache Eviction for Protein Language Models},
  author={Gelal, Rupesh and Aronow, Bruce J. and Prasath, V. B. Surya },
  year={2026},
  note={submitted August 2026}
}