KVine resident-set policy — learned revisit predictor 🌿
A tiny MLP that scores a branch of an agent's trajectory tree by how likely it is to be revisited soon. It is the learned policy for KVine's resident-set manager: high-scoring branches stay in RAM, low-scoring ones are offloaded to disk and restored on demand. A better predictor ⇒ fewer disk restores on revisit.
- Code / write-up: https://github.com/NagaYu/kvine
- Training data: NagaYu/kvine-agent-trajectory-bench
- Interactive demo: NagaYu/kvine-demo
What it is
6 → 16 → 16 → 1 MLP (ReLU), with input standardization (mu/sigma) baked into
the checkpoint. Output is a logit; sigmoid(logit) = P(revisited within 5 steps),
used directly as the resident-set keep-score.
Input features (structural — they depend on the trajectory tree and access pattern, not on the LLM's weights):
| feature | meaning |
|---|---|
depth |
depth from root (shallow branches — e.g. subagents forked at root — are shared/revisited more) |
log_tokens |
log1p(branch token count) |
staleness |
steps since the branch was last accessed |
log_access |
log1p(access count) |
siblings |
number of sibling nodes |
children |
number of child nodes |
Training
Self-supervised from KVine execution logs: for each logged branch snapshot the label is "was this branch accessed again within the next 5 steps?". Data generated by replaying the synthetic agent workload (see the dataset) across 12 seeds.
- Train examples: 10,544 · positive rate 39.4%
- Loss: BCE 0.70 → 0.10 (400 epochs, Adam 1e-2)
- Source model for log collection:
Qwen/Qwen2.5-0.5B-Instruct(features are model-independent)
Why it matters (benchmark)
On a revisit-heavy branch regime, branch-revisit hit rate (KV still resident when a closed branch is revisited), averaged over 5 seeds:
| policy | revisit hit rate | disk restores |
|---|---|---|
| random evictor (baseline) | 60% | 443 |
| heuristic (LRU + depth) | 81% | 198 |
| this learned MLP | 87% | 138 |
The learned policy keeps the right branches resident and needs the fewest restores.
Usage
import torch
from huggingface_hub import hf_hub_download
from modeling_kvine_policy import KVineResidentPolicyMLP # bundled in this repo
ckpt = hf_hub_download("NagaYu/kvine-resident-policy", "model.safetensors")
model = KVineResidentPolicyMLP.from_safetensors(ckpt).eval()
import math
# a fresh, shallow subagent branch vs a very stale deep branch
x = torch.tensor([
[1.0, math.log1p(400), 0.0, math.log1p(1), 0.0, 0.0],
[9.0, math.log1p(120), 15.0, math.log1p(1), 2.0, 0.0],
])
print(model.revisit_prob(x)) # tensor([1.000, 0.000]) -> keep, offload
Limitations
Research prototype. The scorer's leverage is largest when independent branch KV is significant relative to the shared trunk; with a very large shared prompt, COW prefix-sharing dominates and any eviction policy performs similarly (this is a strength of KVine, documented in the repo). fp32 / small models.