RL-learned KV retention
A small DPO-trained controller that decides, at every step, which entries to keep in the KV cache during long-context streaming inference. It replaces the fixed sliding window with a learned policy, and it improves perplexity by up to 1.37 over sliding-window attention on PG-19 at small cache sizes, with under 0.1% latency overhead.
- Role
- Course research with three teammates, EECS 598 Foundations of LLMs. I developed the theory and the RL controller.
- Timeline
- Fall 2025
- Stack
- PyTorch, Qwen-3 8B, DPO, PG-19
- Status
- Paper and poster below
The problem
Streaming inference has to throw tokens away at some point. StreamingLLM keeps the attention sinks plus a sliding window, which is a fixed rule that discards remote context regardless of what's actually in it. The question we asked was whether a small policy could learn what's worth keeping.
What I built
A lightweight scoring policy that decides which remote tokens stay in the cache at each step. We first trained it by behavioral cloning from an oracle's best remote sets, then refined it with DPO on token-level preference pairs (oracle-preferred versus rejected tokens), treating cache compression as the ranking problem it really is. The policy runs alongside the base model and adds under 0.1% latency. We compared it against behavioral cloning alone and against the StreamingLLM sliding-window baseline at cache budgets from 32 to 96 tokens.
Results
The DPO policy beat both behavioral cloning and the sliding-window baseline, with the biggest gap at the smallest budgets: up to 1.37 perplexity at cache size 32. The attention visualizations were my favourite part. You can see the policy hanging on to semantically important tokens from far back in the context rather than just the most recent ones.
What I'd do differently
Perplexity is only a proxy. The result I'd really want next is on a task where retention actually matters, like long-document QA, at a scale where the controller's cost is measured against real serving. I'd also love to try learning the policy online from the serving traffic itself, which is the same instinct behind Transmute.