NagaYu/kvine-resident-policy

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.