Search papers, labs, and topics across Lattice.
This paper introduces AnchorKV, a novel method for key-value (KV) cache compression in large language models that enhances safety alignment while maintaining performance. By implementing a soft penalty token selection rule based on an offline safety anchor, AnchorKV effectively biases token retention away from harmful prompts, addressing vulnerabilities in existing compression methods that fail under adversarial conditions. The results demonstrate that AnchorKV significantly improves safety alignment without compromising the utility of the model, offering a robust solution to the KV cache bottleneck in LLMs.
A small trade-off in utility can lead to significantly enhanced safety alignment in LLMs, protecting against jailbreak attacks.
Large language models (LLMs) outperform earlier architectures on generative inference and long-context tasks, but their large size introduces significant challenges in memory usage, energy cost, and on-device deployment. Since scaling pre-trained language models improves downstream capability \cite{zhao2023survey}, the key-value (KV) cache becomes a dominant inference bottleneck. Recent KV cache compression methods \cite{jo2025fastkv,li2024snapkv,zhou2024dynamickv} reduce this cost by retaining only a subset of attention-relevant tokens. However, while these approaches preserve accuracy on benign workloads, their compression policies either fail to defend against jailbreak attacks \cite{jiang2024robustkv} or degrade safety alignment under aggressive eviction. We propose AnchorKV, a drop-in modification to KV cache compression that biases token retention scores away from directions in key space associated with harmful prompts. AnchorKV constructs an offline safety anchor by adapting a difference-of-means representation engineering approach \cite{arditi2024refusal,zou2023representation} to the layer-specific key projection space used in KV caching. Based on this anchor, a soft penalty token selection rule trades a small amount of utility for substantially improved safety alignment, while reducing to the original compressor when the penalty is zero.
