IndexMem-Qwen3-4B

A retention gate for a frozen Qwen3-4B, from The Gating Mechanism of KV Cache Eviction. It scores every KV entry once, when the entry is created, so the lowest-ranked entries can be evicted from the cache and never recomputed.

This repository holds the gate only (7.5M parameters). The backbone is an unmodified Qwen/Qwen3-4B and is not redistributed here.

At 75% KV-cache compression: RULER 8K 83.89 (Full-KV 92.41).

Inference

Needs the code at https://github.com/Chooovy/kvpress (branch main).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress.indexmem import IndexMemConfig, IndexMemTextGenerationPipeline

model_id = "Qwen/Qwen3-4B"
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2"
).cuda().eval()

pipeline = IndexMemTextGenerationPipeline(
    model=model,
    tokenizer=AutoTokenizer.from_pretrained(model_id),
    scorer_checkpoint="indexmem.pt",   # this repo
    config=IndexMemConfig(cache_budget=2048),
)
print(pipeline("Your document", question="Your question", max_new_tokens=256))

cache_budget is the mean capacity per KV head and includes 64 centroid slots. Four attention sinks and a 128-token local window are always attended; the rest of the budget is allocated across heads by retained attention mass, with an evictable floor of 512.

scorer.json records the scorer geometry and is only needed if you re-convert the original training checkpoint yourself.

Training recipe

Only the gate is trained. The backbone, embeddings and LM head stay frozen.

Scorer	2-layer MLP on the hidden state, width 256, one score per KV head
Loss	reverse KL against the same model running at full KV
Retention rate	learned per entry, fixed token-age scale 16384
Always attended	4 sinks + 128-token window, exempt from the gate's normalizer
Gate mass	256, shared by the evictable entries
Curriculum	300 steps at 8K, then 300 at 16K; global batch 8
Data	LongMino, subsets 2e16 and 2e17