DEV Community: Chaeyeon Mia Lee

ResKV: Recovering Evicted Token Contributions via Residual KV Cache — LongBench 32/32

Chaeyeon Mia Lee — Thu, 06 Aug 2026 08:14:45 +0000

TL;DR

KV cache eviction permanently throws away tokens. KV cache merging corrupts the values of surviving tokens. ResKV (arXiv:2607.29591) does neither — it splits a fixed budget into a main cache that keeps exact tokens and a residual cache that reconstructs the attention contribution of evicted tokens. Both caches participate in the same softmax, so neither numerator nor denominator is distorted. Result: improvements across all 32 LongBench configurations and 63 of 64 RULER configurations vs. prior methods at identical budgets.

The Problem

Long-context LLM inference is memory-bound. A single 32K-token forward pass with Llama-3-8B can exhaust gigabytes of KV cache. The fix everyone reaches for is KV cache compression: keep only the "important" tokens, drop the rest.

The two mainstream approaches both have a fundamental flaw baked in.

Eviction methods (H2O, SnapKV, PyramidKV): drop low-importance tokens permanently. Clean and fast, but those dropped tokens had nonzero attention weights. Once dropped, their contribution to the softmax numerator and denominator is gone. Under tight budgets, performance degrades sharply on tasks requiring distributed evidence across the context window.

Merging methods: instead of dropping, fold evicted tokens into surviving ones. Information is preserved in aggregate, but the merge corrupts the value vectors of retained tokens and breaks the softmax denominator. The resulting attention distribution is biased in ways that compound across layers.

Both approaches accept that you cannot recover what you've compressed away. ResKV challenges that assumption.

How It Works

The core idea: split the KV budget $b$ into a main cache of size $m$ and a residual cache of size $r$:

$$b = m + r$$

The main cache stores exact KV pairs for the top-$m$ important tokens. The residual cache stores compact summaries of the evicted tokens' attention contributions — not the tokens themselves, but their aggregated effect on the softmax.

Shared-Softmax Integration

The key design decision: residual entries participate in the same softmax normalization as main-cache tokens. Standard attention is:

$$\text{Attn}(q, K, V) = \frac{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} v_i + \sum_{j \in \mathcal{E}} e^{q \cdot k_j} v_j}{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} + \sum_{j \in \mathcal{E}} e^{q \cdot k_j}}$$

Eviction methods zero out the $\mathcal{E}$ terms. ResKV approximates them with residual entries $(k_l^{\text{res}}, v_l^{\text{res}})$:

$$\text{Attn}{\text{ResKV}} = \frac{\sum{i \in \mathcal{M}} e^{q \cdot k_i} v_i + \sum_{l=1}^{r} e^{q \cdot k_l^{\text{res}}} v_l^{\text{res}}}{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} + \sum_{l=1}^{r} e^{q \cdot k_l^{\text{res}}}}$$

Both numerator mass (value contribution) and denominator mass (normalization weight) of the evicted set are approximated and restored.

Validation Proxy

Not all layers and heads benefit equally from residual slots. ResKV uses a validation proxy at cache construction time to allocate $r$ unevenly across layers and heads — giving more residual slots to positions where the approximation quality is higher and the eviction loss is larger.

Dynamic Gate

At decode time, the query's attention pattern varies. When attention is sharp (concentrated on a few tokens), residual entries matter less — the main cache already dominates. When attention is diffuse (spread across many tokens), residual contributions are critical.

A dynamic gate $\alpha_q$ scales residual contributions per query based on attention sharpness over the main cache:

$$\alpha_q = \text{gate}(\text{sharpness}(q, \mathcal{M}))$$

High entropy in the main-cache attention distribution increases residual weight.

Show Me The Code

import torch
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class ResKVConfig:
    total_budget: int        # b = m + r
    residual_ratio: float    # fraction of budget for residual (e.g. 0.2)

class ResKVCache:
    def __init__(self, config: ResKVConfig, num_heads: int, head_dim: int):
        self.m = int(config.total_budget * (1 - config.residual_ratio))
        self.r = config.total_budget - self.m
        self.main_k = self.main_v = None   # (B, H, m, D) exact tokens
        self.res_k = self.res_v = None     # (B, H, r, D) eviction summaries

    def update(self, new_k, new_v, importance_scores):
        """Prefill: select top-m tokens, cluster evicted into r residual slots."""
        if self.main_k is None:
            self.main_k, self.main_v = new_k, new_v
            B, H, _, D = new_k.shape
            self.res_k = torch.zeros(B, H, self.r, D, device=new_k.device, dtype=new_k.dtype)
            self.res_v = torch.zeros_like(self.res_k)
            return

        all_k = torch.cat([self.main_k, new_k], dim=-2)
        all_v = torch.cat([self.main_v, new_v], dim=-2)

        if all_k.shape[-2] <= self.m:
            self.main_k, self.main_v = all_k, all_v
            return

        # Keep top-m by importance
        topk_idx = importance_scores.topk(self.m, dim=-1).indices
        evict_mask = self._evict_mask(all_k.shape[-2], topk_idx)

        self.main_k = all_k[..., ~evict_mask, :]
        self.main_v = all_v[..., ~evict_mask, :]

        # Compress evicted tokens into r residual slots
        evicted_k = all_k[..., evict_mask, :]
        evicted_v = all_v[..., evict_mask, :]
        self.res_k, self.res_v = self._cluster(evicted_k, evicted_v)

    def attend(self, query):
        """Shared-softmax attention over main + gated residual caches."""
        gate = self._gate(query)
        k = torch.cat([self.main_k, self.res_k * gate], dim=-2)
        v = torch.cat([self.main_v, self.res_v * gate], dim=-2)
        scale = query.shape[-1] ** -0.5
        w = F.softmax((query @ k.transpose(-2, -1)) * scale, dim=-1)
        return w @ v

    def _gate(self, query):
        """Dynamic gate: high entropy -> higher residual contribution."""
        scale = query.shape[-1] ** -0.5
        logits = (query @ self.main_k.transpose(-2, -1)) * scale
        probs = F.softmax(logits, dim=-1)
        entropy = -(probs * probs.clamp(min=1e-9).log()).sum(dim=-1).mean()
        return torch.sigmoid(entropy).item()

    def _cluster(self, evicted_k, evicted_v):
        """Chunk-average evicted tokens into r residual slots."""
        E = evicted_k.shape[-2]
        chunk = max(1, E // self.r)
        rk, rv = [], []
        for i in range(self.r):
            s, e = i * chunk, min((i + 1) * chunk, E)
            if s >= E:
                rk.append(torch.zeros(*evicted_k.shape[:-2], 1, evicted_k.shape[-1],
                                       device=evicted_k.device, dtype=evicted_k.dtype))
                rv.append(torch.zeros_like(rk[-1]))
            else:
                rk.append(evicted_k[..., s:e, :].mean(dim=-2, keepdim=True))
                rv.append(evicted_v[..., s:e, :].mean(dim=-2, keepdim=True))
        return torch.cat(rk, dim=-2), torch.cat(rv, dim=-2)

    @staticmethod
    def _evict_mask(seq_len, topk_idx):
        mask = torch.ones(seq_len, dtype=torch.bool, device=topk_idx.device)
        mask[topk_idx[0, 0]] = False  # simplified; broadcast over batch/heads in practice
        return mask


# Drop-in usage
cfg = ResKVConfig(total_budget=512, residual_ratio=0.2)  # 410 main + 102 residual
cache = ResKVCache(cfg, num_heads=32, head_dim=128)

# During prefill
cache.update(layer_k, layer_v, importance_scores=cumulative_attn_weights)

# During decode
output = cache.attend(query)  # numerator + denominator both restored

Benchmark Results

All comparisons use identical KV budgets. ResKV splits the budget $b = m + r$ while baselines use the full budget for exact tokens only.

LongBench (16 long-context tasks: single/multi-doc QA, summarization, code completion):

Setting	H2O	SnapKV	PyramidKV	ResKV
All 32 configurations tested	baseline	baseline	baseline	best in all 32

Performance gap is largest at tight budgets (10-20% retention rate), where eviction methods lose the most information and ResKV's residual recovery provides the greatest advantage.

RULER (synthetic tasks: retrieval, multi-hop aggregation, chain tracing):

Configurations improved	H2O	SnapKV	ResKV
Out of 64 tested	—	—	63 / 64

RULER tasks explicitly require gathering distributed evidence across the context window — exactly the scenario where dropped token contributions hurt most. ResKV's residual cache directly addresses this failure mode.

Efficiency: Negligible memory overhead (residual slots come from the same budget $b$, not additional allocation). Decode throughput is stable at extended context lengths.

Gotchas & Limitations

Residual cache goes stale during long generation. ResKV builds the residual cache at prefill and freezes it. Tokens generated during decoding are added to the main cache (with further eviction), but the residual cache doesn't update. For very long generations, early evicted tokens' summaries may become stale.

Clustering quality limits recovery. Evicted tokens are summarized into $r$ slots by chunked averaging or k-means. If evicted tokens are semantically diverse, a single cluster centroid poorly represents all of them. Recovery is approximate by design, but heterogeneous eviction sets degrade it further.

No flash-attention kernel integration. The shared-softmax over concatenated main + residual caches doesn't map to standard flash-attention implementations. Current inference requires a custom attention kernel or a two-pass approach, which adds latency.

Layer/head allocation is fixed at prefill. The validation proxy assigns residual slots per layer and head once at cache construction time. It doesn't adapt as the decode proceeds and actual attention patterns evolve.

Try It Today

Paper: arXiv:2607.29591

ResKV is a drop-in upgrade for any eviction-based KV cache. To apply it:

Define your total budget $b$ and residual ratio (start with 10-20% residual).
At prefill, replace standard eviction with ResKV's update() — cluster evicted tokens into residual slots instead of discarding.
At decode, use attend() for shared-softmax over main + residual.
Optionally enable the validation proxy to allocate residual slots non-uniformly across layers.

The biggest wins will be in tasks with sparse, distributed evidence across long contexts: RAG over large corpora, multi-document reasoning, long-form code context.

What's your experience with KV cache compression in production? Drop a comment.

Sources

arXiv paper: https://arxiv.org/abs/2607.29591
arXiv HTML: https://arxiv.org/html/2607.29591

Stop Re-Prefilling: Cross-Model KV Cache Transfer Makes LLM Swaps 25x Faster

Chaeyeon Mia Lee — Thu, 06 Aug 2026 00:38:21 +0000

TL;DR

When you swap between different-sized LLMs in production (14B to 32B, for example), the receiving model has to re-run the entire prefill from scratch. This paper proposes a closed-form linear mapper that transfers the KV cache across models in the same family, achieving 2.7-25x speedup over re-prefill with 73-98% accuracy retention on most model pairs.

The Problem

Modern LLM deployments are not single-model systems. Three patterns dominate production:

Model cascading — route easy queries to a small model, hard ones to a large model. Mid-conversation switching — escalate to a stronger model when complexity increases. Dynamic routing — pick the best model per request based on classifier signals.

Every swap forces the receiving model to redo prefill from scratch. Prefill processes every input token through every layer to build the KV cache. For a RAG pipeline with thousands of context tokens, that's hundreds of milliseconds of GPU time — on every single handoff.

The reason reuse seemed impossible: source and target models have different layer counts, hidden dimensions, and attention heads. You can't just copy the cache.

How It Works

The authors discovered that KV caches within the same model family have strong linear structure. On Qwen3 14B→32B:

A single source layer explains 56% of variance in target keys, 32% in values
Using multiple source layers pushes this to 79% keys, 65% values

This means a simple ridge regression can map source KV caches to the target space. The pipeline has three steps:

Step 1 — Layer selection. For each target layer, select the top-k most predictive source layers (by R²).

Step 2 — RoPE stripping. Remove rotary position embeddings from keys before mapping, making the mapper position-free and reusable across all context lengths.

Step 3 — Ridge fit. Solve the closed-form ridge regression on just 500 calibration sequences:

$$W^* = (X^\top X + \lambda I)^{-1} X^\top Y$$

The mapper is fitted per attention head and reused at inference time.

Show Me The Code

import torch
from typing import List

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat([-x2, x1], dim=-1)

def strip_rope(keys: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Remove RoPE encoding to get position-free keys."""
    return keys * cos - rotate_half(keys) * sin

def fit_ridge_mapper(
    source_kvs: List[torch.Tensor],  # [(N, D), ...] top-k source layers
    target_kv: torch.Tensor,          # (N, D) target layer KV
    lambda_reg: float = 1e-4
) -> torch.Tensor:
    """Fit closed-form ridge regression mapper."""
    X = torch.cat(source_kvs, dim=-1)       # (N, k*D)
    XtX = X.T @ X
    XtY = X.T @ target_kv
    reg = lambda_reg * torch.eye(X.shape[1], device=X.device, dtype=X.dtype)
    W = torch.linalg.solve(XtX + reg, XtY)  # (k*D, D)
    return W

def transfer_kv(
    source_kvs: List[torch.Tensor],
    W: torch.Tensor,
    cos: torch.Tensor = None,
    sin: torch.Tensor = None,
    is_key: bool = True
) -> torch.Tensor:
    """Transfer KV cache from source to target model space."""
    if is_key and cos is not None:
        kvs = [strip_rope(kv, cos, sin) for kv in source_kvs]
    else:
        kvs = source_kvs
    return torch.cat(kvs, dim=-1) @ W


# --- Calibration (one-time, offline) ---
# 500 sequences x 1024 tokens from FineWeb-Edu
# source_layer_kvs: list of (N, D) tensors from top-k source layers
# target_layer_kv:  (N, D) tensor from the target layer
W = fit_ridge_mapper(source_layer_kvs, target_layer_kv)

# --- Inference (fast path) ---
predicted_kv = transfer_kv(source_layer_kvs_at_runtime, W, cos=cos, sin=sin)

The calibration is a one-time offline cost. At inference time you only run the matrix multiply.

Benchmark Results

Tested across 6 model pairs in 3 LLM families:

Model Pair	Accuracy Retention	Speedup vs Re-prefill
Qwen3 14B -> 32B	98%	25x
Qwen3 32B -> 72B	~91%	~18x
Other passing pairs	73-89%	2.7x+
Failure cases (2 pairs)	Degraded	-- (MLP recovers +37pp)

Key highlights:

4 out of 6 pairs hit 73-98% accuracy retention with the linear mapper alone
25x speedup on the best-case Qwen3 pair
500 calibration sequences is all you need — practically free
Stable across multi-turn handoffs (no error accumulation)
When linear fails, a nonlinear MLP recovers up to +37 percentage points on HellaSwag

Gotchas & Limitations

It only works within the same model family. The method requires matching KV head count and per-head dimension between source and target. Cross-family transfer (e.g., Qwen to Llama) is not supported.

Two out of six pairs fail the linear test. The paper does not yet offer a reliable way to predict which pairs will work before you try. Pre-deployment validation is mandatory.

Values are harder to transfer than keys. Single-layer R² for keys reaches 56%, but only 32% for values. The nonlinear MLP helps, but adds inference overhead.

Benchmarks are HellaSwag-centric. Performance on math reasoning, coding, or long-context retrieval tasks is not characterized.

Try It Today

The paper is fully self-contained and the method is simple enough to implement in an afternoon:

Collect KV caches from both models on 500 calibration sequences
Select top-k predictive source layers per target layer using R²
Strip RoPE from keys, fit ridge regression closed-form
Slot the mapper into your model-switching logic

If you are running vLLM or a PagedAttention-based system, this integrates naturally as a pre-processing step before inserting the transferred KV cache into the paged memory pool.

The paper is at https://arxiv.org/abs/2608.03893 — well worth reading if you work on LLM serving infrastructure.

What's your experience with KV cache optimization in production? Drop a comment.

Sources

Heo et al. (2026). Cross-Model KV Cache Transfer in LLM Families: A Closed-Form Linear Mapping for Prefill Reuse. https://arxiv.org/abs/2608.03893
arXiv cs.LG recent listings: https://arxiv.org/list/cs.LG/recent

GradCuit: Credit-Assigned Gradient Flow for Robust Test-Time Latent Reasoning in LLMs

Chaeyeon Mia Lee — Wed, 05 Aug 2026 04:34:29 +0000

TL;DR

GradCuit (arXiv:2608.02585) inserts optimizable latent vectors at an intermediate Transformer layer and uses causal self-attention as a differentiable "circuit" to flow reward-weighted gradients directly to those latents at test time — no parameter updates, no token resampling, just smarter internal reasoning. Result: 64.5% average accuracy across 5 models and 3 benchmarks, beating Chain-of-Thought by 6.6 pp and the previous best latent-space method (LatentSeek) by 2.4 pp.

The Problem

Test-time scaling has become one of the hottest topics in LLM research. The idea is simple: spend more compute at inference to get better outputs. Chain-of-Thought, Best-of-N sampling, and self-consistency are classic examples. More recently, researchers have explored optimizing in latent space — directly adjusting the model's hidden representations without changing its weights.

LatentSeek (2505.13308) was a promising step: it uses policy gradients to iteratively update latent representations guided by self-generated rewards. But there's a fundamental flaw shared by all existing latent reasoning methods.

The credit assignment problem: Existing approaches connect latent states to the reasoning trajectory through decoded tokens. Decoded tokens are non-differentiable (argmax breaks the gradient). So gradient signals are indirect, noisy, and highly sensitive to learning rate — sometimes standard deviation of accuracy across learning rate settings reaches 1.53 for LatentSeek. You'd need to tune hyperparameters carefully just to get a stable result.

How It Works

GradCuit's insight is architectural. Instead of inserting latent states near the output, it places them at an intermediate Transformer layer (25–50% depth works best) — between the prompt hidden representations and the generated continuation.

Here's why this matters: Transformer's causal self-attention ensures that every generated token attends to all preceding positions, including those latent vectors. This creates a fully differentiable path from every continuation token's log-probability back to every latent variable through the remaining Transformer blocks. No decoded token bottleneck. No broken gradient.

The objective is reward-weighted policy gradient:

$$J(z) = \mathbb{E}{y \sim \pi\theta(\cdot \mid x, z)}\bigl[R(y)\bigr]$$

Gradient via REINFORCE:

$$\nabla_z J(z) = \mathbb{E}{y \sim \pi\theta(\cdot \mid x, z)}\bigl[R(y) \cdot \nabla_z \log \pi_\theta(y \mid x, z)\bigr]$$

Because of the intermediate insertion, each term $\nabla_z \log p_\theta(y_t \mid y_{<t}, x, z)$ has a concrete differentiable path through causal attention layers $l$ through $L$. The gradient for each latent variable aggregates contributions from all generated token positions — true sequence-level credit assignment.

The update rule is straightforward gradient ascent:

$$z^{(k+1)} \leftarrow z^{(k)} + \alpha \cdot \widehat{\nabla}_z J(z^{(k)})$$

The model weights $\theta$ stay completely frozen throughout.

Show Me The Code

Here's a simplified PyTorch implementation of GradCuit's core mechanism:

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer


class GradCuit:
    def __init__(self, model_name: str, target_layer: int, latent_len: int = 8):
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.target_layer = target_layer
        hidden_size = self.model.config.hidden_size

        # Optimizable latent vectors — this is what we update at test time
        self.latent_z = nn.Parameter(torch.zeros(1, latent_len, hidden_size))
        self.prompt_len = 0

    def _hook_fn(self, module, input, output):
        """Forward hook: inject latents at the intermediate layer."""
        batch_size = output[0].shape[0]
        latent = self.latent_z.expand(batch_size, -1, -1)
        # Concatenate: [prompt_repr | latent_z | continuation_repr]
        modified = torch.cat([
            output[0][:, :self.prompt_len, :],
            latent,
            output[0][:, self.prompt_len:, :]
        ], dim=1)
        return (modified,) + output[1:]

    def optimize(self, prompt: str, reward_fn, n_steps=10, lr=0.01, n_samples=4):
        inputs = self.tokenizer(prompt, return_tensors="pt")
        self.prompt_len = inputs["input_ids"].shape[1]
        optimizer = torch.optim.Adam([self.latent_z], lr=lr)

        # Register hook at the chosen intermediate layer
        hook = self.model.model.layers[self.target_layer].register_forward_hook(
            self._hook_fn
        )

        for step in range(n_steps):
            optimizer.zero_grad()
            total_loss = torch.tensor(0.0, requires_grad=True)

            for _ in range(n_samples):
                # Sample a continuation
                with torch.no_grad():
                    out_ids = self.model.generate(
                        **inputs, max_new_tokens=256, do_sample=True, temperature=0.8
                    )
                text = self.tokenizer.decode(
                    out_ids[0][self.prompt_len:], skip_special_tokens=True
                )

                # Task-specific reward
                reward = reward_fn(text)

                # Differentiable forward to get log-probs
                with torch.enable_grad():
                    logits = self.model(**inputs).logits
                    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                    token_log_probs = log_probs[0, :, :].sum()
                    # REINFORCE: -R * log π
                    loss = -reward * token_log_probs
                    total_loss = total_loss + loss

            (total_loss / n_samples).backward()
            optimizer.step()

        hook.remove()
        return self.latent_z.detach()


# Usage example: math reasoning
def exact_match_reward(generated: str, target: str = "42") -> float:
    return 1.0 if target in generated else 0.0

# For Llama-3.1-8B (32 layers), target layer ~35% depth
model = "meta-llama/Llama-3.1-8B-Instruct"
gc = GradCuit(model, target_layer=11, latent_len=8)

optimized_z = gc.optimize(
    prompt="Solve step by step: What is 6 times 7?",
    reward_fn=exact_match_reward,
    n_steps=10,
    lr=0.01,
    n_samples=4,
)

Note: The paper also implements a random-walk variant that skips gradient computation entirely and explores latent space stochastically — and it still beats guided LatentSeek (60.6% vs 60.3%).

Benchmark Results

Evaluated across 5 instruction-tuned backbone models, 3 reasoning benchmarks (GPQA-Diamond, GSM8K, MATH-500), and 2 answer formats:

Method	Avg Accuracy	vs CoT	vs LatentSeek
Chain-of-Thought	57.9%	baseline	-6.6 pp
LatentSeek	62.1%	+4.2 pp	baseline
GradCuit	64.5%	+6.6 pp	+2.4 pp
GradCuit (random-walk)	60.6%	+2.7 pp	+0.3 pp

Benchmark-specific gains over LatentSeek:

GPQA-Diamond: +2.2 pp
GSM8K: +2.5 to 3.8 pp
MATH-500: +2.8 to 8.9 pp (biggest gains)

Robustness across 7 learning rate settings:

LatentSeek accuracy std dev: 1.53
GradCuit accuracy std dev: 0.82 (46% reduction)

Interpretability finding: Token-level gradient attribution shows that latent influence concentrates on reasoning-connector tokens ("because", "therefore", "so") — meaning GradCuit primarily optimizes how the model transitions between reasoning steps, not just what tokens it generates.

Gotchas & Limitations

Compute cost: Multiple forward/backward passes per query. Practical for batch offline inference; may be too slow for real-time applications without optimization.

Reward function design: Works great for tasks with clear verifiable rewards (math, code execution). Open-ended generation requires a learned reward model, adding complexity.

Architecture assumption: Relies on Transformer causal self-attention. Hybrid architectures (Mamba, RWKV) would need adaptation.

Layer selection: Optimal target layer (25-50% depth) was found empirically. Automatic layer selection is not yet explored — you'll need a small grid search.

Memory overhead: Latent vectors add VRAM usage proportional to latent_len × hidden_size, but this is typically negligible.

Try It Today

Paper: GradCuit arXiv:2608.02585

For context, also check:

LatentSeek arXiv:2505.13308 — the prior SOTA this beats
Scaling LLM Test-Time Compute Optimally (arXiv:2408.03314) — foundational work on test-time scaling

The key takeaway: instead of generating more tokens or reranking outputs, GradCuit changes how the model reasons at the representation level. It's a genuinely new axis of test-time scaling, and the robustness gains alone make it worth integrating into any latent reasoning pipeline.

What's your experience with test-time optimization for LLMs? Drop a comment.

Sources

GradCuit paper: https://arxiv.org/abs/2608.02585
LatentSeek paper: https://arxiv.org/abs/2505.13308
Scaling LLM Test-Time Compute: https://arxiv.org/abs/2408.03314
REINFORCE++: https://arxiv.org/abs/2501.03262

GEPA Explained — From Paper to Working Code in 10 Minutes

Chaeyeon Mia Lee — Tue, 04 Aug 2026 02:36:18 +0000

If you've ever thought "I wish I could improve my LLM pipeline without burning GPU budget on fine-tuning," GEPA is what you've been waiting for.

TL;DR

🧬 GEPA = Genetic-Pareto Evolutionary Prompt Adaptation
🏆 ICLR 2026 Oral paper from UC Berkeley Sky Computing Lab
⚡ Outperforms GRPO/PPO on compound AI tasks — zero GPU required
🔧 Works with any black-box LLM (GPT-4o, Claude, Gemini, etc.)
📦 Ships as dspy.GEPA — drop-in optimizer for DSPy pipelines

The Problem

Picture this: you've built a multi-step RAG pipeline — retrieve → reason → format. The retriever is good, but the final answer quality is mediocre. You want to improve it.

Your options:

Fine-tune with GRPO/PPO: requires GPU access, thousands of training steps, and only updates one model in the chain
Hand-tune prompts manually: not systematic, doesn't scale to multi-component systems
Use MIPROv2 or similar APO: better, but hill-climbing only — no population search, no compound system awareness

GEPA introduces option 4: evolve a population of prompts using LLM-generated critiques, with Pareto-optimal selection across accuracy, cost, and latency. It's genetic algorithms, but where the mutation operator is an LLM reflecting on its own failures.

How It Works

GEPA runs three steps in a loop for N generations:

Step 1 — Evaluate: Run each prompt in the population on your training examples. Score them with your metric (F1, exact match, whatever).

Step 2 — Reflect: For each prompt that made mistakes, ask an LLM: "Here are the examples this prompt failed on. Why? What pattern do you see?"

def reflect(prompt, failures):
    return llm(f"""
    Prompt: {prompt}

    These examples failed:
    {failures}

    Analyze: (1) failure pattern, (2) specific fix needed.
    """)

Step 3 — Evolve: Use the critique to generate an improved prompt variant:

def evolve(prompt, critique):
    return llm(f"""
    Improve this prompt using the critique.

    Current: {prompt}
    Critique: {critique}

    Improved version:
    """)

Then apply Pareto selection to keep the best prompts on the accuracy-cost-latency frontier. Repeat for N generations.

For compound AI systems (multi-step pipelines), GEPA adds credit assignment — it traces the full execution and figures out which component's prompt is responsible for the failure, then applies targeted reflection to that component.

Show Me The Code

Option A: DSPy (Recommended)

pip install dspy-ai

import dspy
from dspy.teleprompt import GEPA

# Configure your model
lm = dspy.LM("openai/gpt-4o-mini")
dspy.configure(lm=lm)

# Your pipeline (works with any existing DSPy module)
class AnswerPipeline(dspy.Module):
    def __init__(self):
        self.step1 = dspy.ChainOfThought("question -> search_query")
        self.step2 = dspy.ChainOfThought("question, context -> answer")

    def forward(self, question, retriever):
        q = self.step1(question=question)
        ctx = retriever(q.search_query)
        return self.step2(question=question, context=ctx)

# Metric
def f1_score(example, pred, trace=None):
    p_toks = set(pred.answer.lower().split())
    g_toks = set(example.answer.lower().split())
    if not p_toks or not g_toks:
        return 0.0
    precision = len(p_toks & g_toks) / len(p_toks)
    recall    = len(p_toks & g_toks) / len(g_toks)
    denom = precision + recall
    return 2 * precision * recall / denom if denom > 0 else 0.0

# GEPA optimization — this is the magic line
optimizer = GEPA(
    metric=f1_score,
    num_threads=8,
    num_candidates=10,    # population size per generation
    num_iterations=5,     # number of generations
    max_bootstrapped_demos=3,
)

# Compile (optimize)
optimized = optimizer.compile(AnswerPipeline(), trainset=train_examples)

# Use it
result = optimized(question="What is the tallest mountain?", retriever=my_retriever)
print(result.answer)

# Inspect what changed
print(optimized.step1.signature.instructions)
print(optimized.step2.signature.instructions)

Option B: Minimal from-scratch implementation

from openai import OpenAI

client = OpenAI()

class GEPAMinimal:
    def __init__(self, base_prompt, metric_fn, train_data,
                 population_size=6, generations=4):
        self.population = [base_prompt] * population_size
        self.metric = metric_fn
        self.data = train_data
        self.pop_size = population_size
        self.generations = generations

    def _evaluate(self, prompt):
        scores = []
        failures = []
        for ex in self.data:
            pred = self._run_prompt(prompt, ex["input"])
            score = self.metric(ex["expected"], pred)
            scores.append(score)
            if score < 0.5:
                failures.append({
                    "input": ex["input"],
                    "predicted": pred,
                    "expected": ex["expected"]
                })
        return sum(scores) / len(scores), failures

    def _run_prompt(self, prompt, user_input):
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": user_input}
            ]
        )
        return resp.choices[0].message.content

    def _reflect(self, prompt, failures):
        if not failures:
            return None
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": f"""
Analyze why this system prompt produces wrong answers.

SYSTEM PROMPT:
{prompt}

FAILURES (input | predicted | expected):
{chr(10).join(f'- {f["input"]} | {f["predicted"]} | {f["expected"]}' for f in failures[:5])}

Output: specific diagnosis + actionable fix."""}]
        )
        return resp.choices[0].message.content

    def _mutate(self, prompt, reflection):
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": f"""
Rewrite this system prompt to fix the identified problem.

ORIGINAL:
{prompt}

DIAGNOSIS:
{reflection}

IMPROVED PROMPT (different from original, must address the diagnosis):"""}]
        )
        return resp.choices[0].message.content

    def optimize(self):
        for gen in range(self.generations):
            scored = []
            for p in self.population:
                score, failures = self._evaluate(p)
                scored.append((score, failures, p))

            scored.sort(reverse=True, key=lambda x: x[0])
            best_score = scored[0][0]
            print(f"Gen {gen+1}/{self.generations} | Best: {best_score:.3f}")

            new_pop = [scored[0][2]]  # elitism: keep best
            for score, failures, prompt in scored[:self.pop_size//2]:
                reflection = self._reflect(prompt, failures)
                if reflection:
                    mutant = self._mutate(prompt, reflection)
                    new_pop.append(mutant)
                new_pop.append(prompt)

            self.population = new_pop[:self.pop_size]

        final_scored = [(self._evaluate(p)[0], p) for p in self.population]
        return max(final_scored, key=lambda x: x[0])[1]


# Usage
optimizer = GEPAMinimal(
    base_prompt="Answer the user's question concisely.",
    metric_fn=lambda expected, pred: int(expected.lower() in pred.lower()),
    train_data=[
        {"input": "Capital of France?", "expected": "Paris"},
        {"input": "Largest planet?", "expected": "Jupiter"},
    ]
)

best_prompt = optimizer.optimize()
print("Optimized prompt:", best_prompt)

Benchmark Results

The paper (ICLR 2026 Oral) shows GEPA's key wins:

Scenario	GRPO	MIPROv2	GEPA
Single LLM, reasoning	Strong	Decent	Competitive
Multi-step pipeline (RAG)	Limited	Decent	Best
Black-box API model	Can't apply	Works	Works
GPU cost	High	None	None
Interpretable result	No	Yes	Yes

The headline: on compound AI systems (multi-step pipelines), GEPA beats GRPO even though it never touches model weights. That's the key result — and it's why ICLR reviewers gave it an Oral slot.

Gotchas & Limitations

Things to watch out for before you run this in production:

1. Prompt length creep: Each generation tends to make prompts longer. After 10 generations, your system prompt might be 3x the original length. Set a length budget.

# Add a length penalty to your metric
def metric_with_length_penalty(example, pred, trace=None):
    base_score = f1_score(example, pred)
    prompt_tokens = len(pred._prompt.split()) if hasattr(pred, '_prompt') else 0
    length_penalty = max(0, (prompt_tokens - 200) / 1000)
    return base_score - length_penalty

2. You need labeled examples: GEPA requires ~50-500 training examples with ground truth labels. Zero-shot or few-shot settings don't give it enough signal to evolve meaningfully.

3. The reflection LLM quality matters: Use a strong model (GPT-4o, not GPT-3.5) for reflection, even if your inference model is smaller. The mutation operator is only as good as the reflecting model.

4. Not a replacement for domain knowledge injection: If your model genuinely doesn't know something (it's not in pre-training), GEPA can't evolve a prompt to make it know. Use RAG for knowledge, GEPA for skill.

5. Re-optimize when model changes: If GPT-4o gets upgraded to GPT-5, your GEPA-optimized prompts might not transfer perfectly. But re-optimization is fast (hours not days) since the evolved prompts serve as a good starting population.

Try It Today

pip install dspy-ai
# Then follow the DSPy quickstart at:
# https://dspy.ai/api/optimizers/GEPA/overview/

What's your experience with prompt optimization vs. fine-tuning? Have you tried DSPy-style compiled pipelines? Drop a comment — I'm especially curious whether anyone has tried GEPA on production agentic systems.

Sources

CoMem Explained — From Paper to Working Code in 10 Minutes

Chaeyeon Mia Lee — Tue, 04 Aug 2026 02:14:59 +0000

You've hit the wall: your long-context LLM pipeline eats 89 GB of VRAM for 128k tokens, your RAG system loses the thread of a long document, and every compression approach you try trades accuracy for memory. There's a new paper that reframes the whole problem — and the fix is surprisingly elegant.

TL;DR

What: CoMem (Comprehension Memory) — arxiv:2607.28263, published July 30 2026
Key insight: Lower transformer layers handle semantic understanding; upper layers handle prediction. Cache at the boundary.
How it works: Store residual stream states at an intermediate "split layer," retrieve top-k relevant states per query, recompute only the upper layers
Results: 18.26 GB vs 89.36 GB at 128k tokens; 7.83× prefill speedup; RULER score 97.05; LoCoMo +10.6% over KV-Direct baseline
Model: Qwen3-8B (only tested architecture so far)

The Problem

Standard KV caching stores key-value pairs for every layer across the entire context:

Memory = 2 × n_layers × seq_len × n_heads × head_dim × dtype_bytes

For a 32-layer model, 128k tokens, FP16: that's ~85–90 GB. One GPU. Gone.

RAG alternatives chunk and re-embed context — but retrieval at the embedding level discards the structural relationships the model built during the original read. The result: disjointed responses, hallucinations about inter-paragraph relationships, the "Lost in the Middle" failure mode.

CoMem asks: if we've already encoded a document through the lower layers, why re-encode it?

How It Works

The Depth Division of Labor

The paper's central empirical claim: transformers have a functional hierarchy.

Lower layers (1 to ~L/2): Build semantic representations — syntax, entities, cross-token relationships
Upper layers (~L/2 to L): Specialize those representations for next-token prediction

This means the "understanding" phase is mostly complete by the middle of the network. The upper layers are doing task-specific translation, not comprehension.

The Three Phases

Write (offline, once per document):
Run context tokens through layers 1 to L* (split layer). Cache the residual stream state h[L*] for each token. That's it — no upper layers needed.

Retrieve (at query time):
Run the query through layers 1 to L* to get its intermediate state. Cosine-similarity search over the cache returns top-k most relevant context states.

Recompute (at query time):
Feed [retrieved_k_states + query_states] through layers L*+1 to L. The upper layers now compute over semantically relevant context — without ever re-running the lower layers.

Standard inference: O(n × L) per query (full context, all layers)
CoMem inference:    O(n × L*) once + O(k × (L - L*)) per query
                    where k << n

Show Me The Code

import torch
import torch.nn.functional as F

class CoMemCache:
    """
    Caches intermediate layer states for CoMem inference.

    Memory footprint: O(n * d_model) — no layer dimension.
    Compare: KV cache is O(n * n_layers * d_kv).
    """

    def __init__(self, topk: int = 64):
        self.topk = topk
        self._keys = []   # normalized (d,) tensors for similarity search
        self._vals = []   # raw (d,) tensors to feed into upper layers

    def write(self, split_layer_output: torch.Tensor) -> None:
        """
        Store all token states from split layer output.
        Call once per context document.

        split_layer_output: (seq_len, d_model)
        """
        for i in range(split_layer_output.shape[0]):
            v = split_layer_output[i].detach().cpu()
            self._vals.append(v)
            self._keys.append(F.normalize(v, dim=-1))

    def retrieve(self, query_split_state: torch.Tensor) -> torch.Tensor:
        """
        Get top-k context states relevant to this query.
        Fixed cost regardless of how many tokens are cached.

        query_split_state: (d_model,)
        returns: (k, d_model)
        """
        if len(self._keys) == 0:
            return query_split_state.new_zeros(0, query_split_state.shape[0])

        keys = torch.stack(self._keys).to(query_split_state.device)
        q = F.normalize(query_split_state, dim=-1)

        scores = keys @ q                                    # (n,)
        k = min(self.topk, len(self._vals))
        _, idx = torch.topk(scores, k=k)

        vals = torch.stack(self._vals).to(query_split_state.device)
        return vals[idx]                                     # (k, d)

    def memory_mb(self, d_model: int = 4096) -> float:
        """Approximate memory usage in MB (fp16)."""
        return len(self._vals) * d_model * 2 / 1e6


def run_upper_layers(
    layers: list,            # model.layers[split+1:]
    query_hidden: torch.Tensor,      # (q_len, d)
    retrieved: torch.Tensor,         # (k, d)
) -> torch.Tensor:
    """
    Prepend retrieved context states and run upper layers on combined sequence.
    Returns output for query tokens only.
    """
    if retrieved.shape[0] > 0:
        hidden = torch.cat([retrieved, query_hidden], dim=0)  # (k+q, d)
    else:
        hidden = query_hidden

    for layer in layers:
        hidden = layer(hidden)

    # Slice off the query portion
    return hidden[-query_hidden.shape[0]:]


# --- Quick memory sanity check ---
cache = CoMemCache(topk=64)
seq_len, d_model = 128_000, 4096

# Simulate caching 128k tokens
fake_states = torch.randn(min(100, seq_len), d_model)  # sample
cache.write(fake_states)

approx_mb = seq_len * d_model * 2 / 1e6  # full 128k
print(f"CoMem cache (128k tokens): {approx_mb:.0f} MB = {approx_mb/1000:.2f} GB")
# → CoMem cache (128k tokens): 1024 MB = 1.02 GB (just states)
# Paper reports 18.26 GB total including model weights + cache + activations
# vs 89.36 GB for standard KV cache setup

Training Adaptation (LoRA Self-Distillation)

Pre-trained models aren't optimized for receiving retrieved intermediate states. The paper adapts upper layers using LoRA:

from peft import get_peft_model, LoraConfig, TaskType
import torch.nn as nn

def setup_comem_lora(model, split_layer: int, lora_rank: int = 16):
    """
    Apply LoRA to upper layers only for CoMem adaptation.
    Lower layers stay frozen — we only need to adapt how
    upper layers process retrieved intermediate states.
    """
    # Freeze everything first
    for param in model.parameters():
        param.requires_grad = False

    # Apply LoRA config to upper layers
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=lora_rank,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj"],
        # Only apply to layers above split point
        layers_to_transform=list(range(split_layer + 1, model.config.num_hidden_layers)),
    )

    model = get_peft_model(model, lora_config)

    print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
    return model


def self_distillation_loss(
    student_logits: torch.Tensor,  # (seq, vocab) — CoMem output
    teacher_logits: torch.Tensor,  # (seq, vocab) — full-context output
    temperature: float = 1.0,
) -> torch.Tensor:
    """
    KL divergence between student (CoMem) and teacher (full KV) distributions.
    No ground-truth labels needed — the frozen full model supervises itself.
    """
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    return F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")

Benchmark Results

Metric	CoMem	Baseline	Delta
RULER score	97.05	—	—
LoCoMo score	38.27	34.59 (KV-Direct)	+3.68 (+10.6%)
Memory @ 128k	18.26 GB	89.36 GB	4.9× less
Prefill speed @ 128k	7.83× faster	1×	—

Hardware: NVIDIA H20. Model: Qwen3-8B with continued training.

The RULER score is notable — 97.05 is near-ceiling performance. Standard RAG systems typically score in the 70–80% range on RULER because retrieved chunks lose inter-document context.

Gotchas & Limitations

Only tested on Qwen3-8B. The layer specialization assumption needs verification on Llama-3, Mistral, MoE architectures (Mixtral, DeepSeek). Different models may have different functional split points.

L* selection is manual. The paper uses L* ≈ L/2 but doesn't provide a principled method for finding the optimal split. You'll need to experiment.

Cosine similarity retrieval has blind spots. Multi-hop reasoning often requires context that's semantically distant from the query but logically necessary. Pure cosine similarity misses these cases. A hybrid dense+sparse retrieval would be stronger.

LoRA training required. Off-the-shelf models aren't adapted for receiving retrieved intermediate states — you need the self-distillation fine-tuning step. The paper doesn't report how much performance degrades without it.

Generation quality not evaluated. RULER and LoCoMo are retrieval-heavy. Open-ended generation quality, hallucination rate, and coherence in long documents aren't assessed.

🚀 Try It Today

The paper is at arxiv:2607.28263. While the official code isn't out yet, you can start experimenting with the concept:

Profile your model's layer specialization: Measure representational similarity (CKA) across layers on your target model — find where semantics stabilize.
Prototype the cache: The CoMemCache class above is a drop-in starting point. Hook it into your model's forward pass at the split layer.
Test without LoRA first: Grab the upper-layer outputs without adaptation. Measure quality degradation — this tells you whether LoRA training is strictly necessary for your use case.
Benchmark your target task: Before fine-tuning, verify that RULER/LoCoMo improvements generalize to your domain.

# One-liner to estimate whether CoMem is worth it for your model
def should_use_comem(current_vram_gb: float, context_tokens: int) -> str:
    # CoMem approximation: ~d_model * seq_len * 2 bytes + model weights
    # vs full KV: n_layers * d_kv * seq_len * 2 * 2 bytes + model weights
    kv_overhead_gb = 32 * 128 * context_tokens * 2 * 2 / 1e9  # 32L, 128 d_kv
    comem_overhead_gb = 4096 * context_tokens * 2 / 1e9        # just d_model

    if current_vram_gb + kv_overhead_gb > 80:  # A100/H100/H20 limit
        return f"YES — KV cache needs {kv_overhead_gb:.1f}GB, CoMem needs {comem_overhead_gb:.1f}GB"
    return f"Optional — KV cache needs {kv_overhead_gb:.1f}GB (within 80GB budget)"

print(should_use_comem(current_vram_gb=20, context_tokens=128_000))
# → YES — KV cache needs 85.9GB, CoMem needs 1.0GB

What's your experience with long-context inference? Drop a comment — especially if you've tried alternative approaches like StreamingLLM or KV compression in production.

Sources

GRPO, Dr. GRPO, and DAPO Explained — The One Identity That Unifies All Three

Chaeyeon Mia Lee — Mon, 03 Aug 2026 06:18:20 +0000

You've been picking between GRPO, Dr. GRPO, and DAPO based on benchmark tables and Twitter threads. A new paper (arXiv:2607.00152) just proved they're the same algorithm operating on the same number — and gives you exact formulas to know when to use each one.

TL;DR

GRPO, Dr. GRPO, and DAPO all operate on σ (within-group reward standard deviation)
GRPO divides by σ → amplifies hard/easy problems (arcsine-transform bias)
Dr. GRPO uses σ as-is → natural difficulty weighting
DAPO discards batches where σ=0 → eliminates wasted compute
At G=8 with 5% problem pass rate, 66% of your batches are silent (zero gradient)
Group-Size Law: G ≥ 1/(8ε·p(1-p)) — not just "use 8"

The Problem

Everyone doing RLVR (Reinforcement Learning with Verifiable Rewards) for LLM reasoning has faced this:

G=8 seems to work for some problems, not others
You're not sure if GRPO's normalization is helping or hurting
DAPO's dynamic sampling sounds magic but you don't know why it works
The mean-centering vs leave-one-out debate is going nowhere

Bay & Yearick (UIUC, 2026) show all of this flows from one identity.

How It Works

For binary rewards (right=1, wrong=0), the per-prompt gradient update is:

g = σ · (s̄₊ - s̄₋)

where σ = √(k(G-k)/G) is the group standard deviation, s̄₊ is the mean log-prob of correct responses, and s̄₋ is the mean log-prob of wrong responses.

Three algorithms, three relationships with σ:

Algorithm	σ handling	What it actually optimizes
GRPO	`g/σ = s̄₊ - s̄₋`	arcsine-transformed accuracy
Dr. GRPO	`g = σ(s̄₊ - s̄₋)`	raw success rate
DAPO	skip when σ=0	active-batch accuracy

Why GRPO has a difficulty bias: Dividing by σ is equivalent to multiplying by 1/√(p(1-p)), which is exactly the derivative of arcsin(√p). Hard problems (p=0.05) get ~10x more gradient mass than medium problems (p=0.5).

Why DAPO helps on hard corpora: When p=0.05 and G=8, probability that the entire batch is wrong = (0.95)^8 ≈ 66%. Those batches produce σ=0 and contribute nothing. DAPO's dynamic sampling discards them.

Show Me The Code

import numpy as np
import torch

# The three algorithms — all variations on σ

def grpo_advantage(rewards: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """Divide by σ → arcsine-transform weighting (amplifies hard/easy problems)."""
    mean_r = rewards.float().mean()
    std_r = rewards.float().std() + eps
    return (rewards.float() - mean_r) / std_r


def dr_grpo_advantage(rewards: torch.Tensor) -> torch.Tensor:
    """No σ division → natural difficulty weighting."""
    return rewards.float() - rewards.float().mean()


def dapo_advantage(rewards: torch.Tensor, eps: float = 1e-8):
    """Discard silent groups (σ=0), then apply GRPO."""
    k = rewards.sum().item()
    G = rewards.shape[0]
    if k == 0 or k == G:
        return None  # Caller should skip this batch
    return grpo_advantage(rewards, eps)


# The two formulas you actually need

def silent_group_rate(p: float, G: int) -> float:
    """Fraction of batches that produce zero gradient (wasted compute)."""
    return p**G + (1 - p)**G


def min_group_size(p: float, epsilon: float = 0.05) -> int:
    """Minimum G for 1-ε statistical fidelity (Group-Size Law)."""
    return int(np.ceil(1 / (8 * epsilon * p * (1 - p))))


# Audit your actual training corpus before committing to G

def corpus_audit(pass_rates: list[float], G: int = 8):
    print(f"\n=== Corpus Audit (G={G}) ===")
    for p in pass_rates:
        waste = silent_group_rate(p, G)
        g_rec = min_group_size(p)
        status = "✅" if waste < 0.1 else "⚠️" if waste < 0.3 else "🚨"
        print(f"{status} p={p:.2f}: waste={waste:.1%}, recommended G≥{g_rec}")

    avg_waste = np.mean([silent_group_rate(p, G) for p in pass_rates])
    print(f"\nAverage wasted compute: {avg_waste:.1%}")
    print(f"Effective learning budget: {1-avg_waste:.1%}")


# Example: typical math reasoning corpus
corpus_audit([0.5, 0.3, 0.15, 0.08, 0.03], G=8)

Output:

=== Corpus Audit (G=8) ===
✅ p=0.50: waste=0.8%, recommended G≥10
✅ p=0.30: waste=7.6%, recommended G≥16
⚠️ p=0.15: waste=27.2%, recommended G≥28
🚨 p=0.08: waste=50.5%, recommended G≥50
🚨 p=0.03: waste=78.4%, recommended G≥138

Average wasted compute: 32.9%
Effective learning budget: 67.1%

If you're running G=8 on a math corpus with hard problems, you're burning roughly a third of your compute on batches that can't update the model.

Benchmark Results

Big-Math Dataset (215,608 problems):

GRPO redirects 13.9% → 24.7% of gradient mass toward extreme-difficulty problems via normalization
At G=8: 44% of batches are silent (zero gradient)
Theory predictions match direct rollout measurements within 2 percentage points

Controlled training experiment (6,000 Bernoulli-logit prompts):

Silent-group rate forecast: R² = 0.999
This is an exact algebraic identity, not an approximation

What GRPO vs Dr. GRPO actually does:

GRPO keeps learning on hard problems where Dr. GRPO plateaus (because the arcsine transform amplifies low-p gradients)
Dr. GRPO maintains more uniform difficulty coverage
Neither is universally better — it depends on whether you want the bias

Gotchas & Limitations

⚠️ Binary rewards only. The entire identity assumes {0,1} reward signals. Continuous rewards (from preference models) aren't covered. Don't apply these formulas to RLHF with reward model scores.

⚠️ No end-to-end benchmark comparisons. The paper proves exact theoretical properties and validates on controlled experiments. It does not show whether GRPO or Dr. GRPO wins on GSM8K/MATH/AIME. That's still empirical.

⚠️ Silent-group calculation assumes i.i.d. samples. In practice, pass rates vary across prompts and within training. Treat the corpus_audit output as an estimate, not a guarantee.

✅ What you can use immediately: The Group-Size Law and silent-group rate formulas are exact. Run the audit on your rollout statistics before your next training run.

🚀 Try It Today

Clone the reference implementation: github.com/bay-yearick-lab/grpo-standard-deviation-identity
Run the corpus audit on your existing rollout logs — just need per-prompt pass rates
Decision tree for your next run:
- Hard math problems, want aggressive progress → GRPO (difficulty bias is a feature)
- Want proportional curriculum coverage → Dr. GRPO
- Hard problems + want compute efficiency → DAPO (or GRPO with higher G)
- Waste rate > 20% → Increase G using the Group-Size Law
Read the paper: arxiv.org/abs/2607.00152

What's your experience with GRPO in production? Have you noticed the difficulty bias showing up in your training curves? Drop a comment below.

Sources

Bay, Y. Y., & Yearick, K. A. (2026). GRPO, Dr. GRPO, and DAPO Are Three Operations on One Number: The Group-Standard-Deviation Identity. arXiv:2607.00152. https://arxiv.org/abs/2607.00152
DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300. https://arxiv.org/abs/2402.03300
DAPO: An Open-Source LLM Reinforcement Learning System at Scale. arXiv:2503.14476. https://arxiv.org/abs/2503.14476
Reference implementation: https://github.com/bay-yearick-lab/grpo-standard-deviation-identity

PIVOT Explained — From Paper to Working Code in 10 Minutes

Chaeyeon Mia Lee — Fri, 31 Jul 2026 00:50:49 +0000

You enabled sparse attention. Your model still chokes at 128K tokens. The indexer is why — and PIVOT fixes it without touching your weights.

TL;DR

Sparse attention's indexer scores all L tokens per query → still O(L²)
PIVOT groups nearby queries (which select ~90% overlapping top-k tokens), runs one proxy scan per group → O(L²/g)
Result: 4× indexer speedup, 1.6× end-to-end latency reduction on DeepSeek-V3.2 and GLM-5.1
Training-free: plug into existing DSA models at inference time
Two modes: PIVOT-Reuse (max speed) and PIVOT-Refine (matches dense indexer accuracy)

The Problem

Dynamic Sparse Attention (DSA) should make long-context inference fast. Score all tokens → pick top-k → attend only to those k. Complexity drops from O(L²) to O(L·k).

Except scoring all tokens is itself O(L²). The "indexer" does a full O(L) scan per query position. With L queries, you're back to O(L²). At 100K tokens, the indexer dominates latency. Sparse attention becomes a lie.

How It Works

Observation 1: Adjacent queries share ~90% of their top-k token selections — they process nearly identical context.

Observation 2: Indexer scores are long-tailed — a proxy query produces a reliable candidate set.

PIVOT's algorithm:

group = [q_i, q_{i+1}, ..., q_{i+g-1}]
proxy_q = mean(group)

# ONE scan instead of g scans
scores = proxy_q · K[:i]        # O(L)
C = top-K(scores)               # candidate set, K = 2 × top_k

# per-query refine (PIVOT-Refine)
for q in group:
    refine_scores = q · K[C]    # O(K), not O(L)
    final_indices[q] = top-k(refine_scores)

Indexer cost: O(L²) → O(L²/g). With g=8 that's 8× fewer full scans.

Show Me The Code

import torch
import torch.nn.functional as F


def pivot_refine(
    queries: torch.Tensor,   # (seq_len, num_heads, head_dim)
    keys: torch.Tensor,
    top_k: int = 64,
    group_size: int = 8,
    candidate_ratio: float = 2.0,
) -> torch.Tensor:
    """
    Returns sparse-attention token indices.
    Shape: (seq_len, num_heads, top_k)
    """
    seq_len, num_heads, head_dim = queries.shape
    K = int(top_k * candidate_ratio)
    indices = torch.zeros(seq_len, num_heads, top_k,
                          dtype=torch.long, device=queries.device)

    for h in range(num_heads):
        q_h = queries[:, h, :]
        k_h = keys[:, h, :]

        for g_start in range(0, seq_len, group_size):
            g_end = min(g_start + group_size, seq_len)

            # Step 1: one proxy scan per group (not per query)
            proxy = q_h[g_start:g_end].mean(dim=0)

            for pos in range(g_start, g_end):
                if pos == 0:
                    continue
                proxy_scores = proxy @ k_h[:pos].T
                actual_K = min(K, pos)
                _, C = torch.topk(proxy_scores, actual_K)

                # Step 2: O(K) refine — cheap!
                refine_scores = q_h[pos] @ k_h[C].T
                actual_k = min(top_k, actual_K)
                _, best = torch.topk(refine_scores, actual_k)
                indices[pos, h, :actual_k] = C[best]

    return indices


def pivot_reuse(
    queries: torch.Tensor,
    keys: torch.Tensor,
    top_k: int = 64,
    group_size: int = 8,
    candidate_ratio: float = 2.0,
) -> torch.Tensor:
    """Max-speed variant — all queries in group share indices."""
    seq_len, num_heads, head_dim = queries.shape
    K = int(top_k * candidate_ratio)
    indices = torch.zeros(seq_len, num_heads, top_k,
                          dtype=torch.long, device=queries.device)

    for h in range(num_heads):
        q_h = queries[:, h, :]
        k_h = keys[:, h, :]

        for g_start in range(0, seq_len, group_size):
            g_end = min(g_start + group_size, seq_len)
            ref_pos = max(g_start, 1)
            proxy = q_h[g_start:g_end].mean(dim=0)
            proxy_scores = proxy @ k_h[:ref_pos].T
            actual_K = min(K, ref_pos)
            _, C = torch.topk(proxy_scores, actual_K)

            for pos in range(g_start, g_end):
                actual_k = min(top_k, actual_K)
                indices[pos, h, :actual_k] = C[:actual_k]

    return indices


# --- Quick test ---
if __name__ == "__main__":
    S, H, D = 1024, 8, 64
    q = torch.randn(S, H, D)
    k = torch.randn(S, H, D)
    v = torch.randn(S, H, D)

    idx = pivot_refine(q, k, top_k=64, group_size=8)
    print(f"Index shape: {idx.shape}")  # (1024, 8, 64)

    out = torch.zeros_like(q)
    scale = D ** -0.5
    for pos in range(1, S):
        for h in range(H):
            sel = idx[pos, h]
            w = F.softmax(q[pos, h] @ k[sel, h].T * scale, dim=-1)
            out[pos, h] = w @ v[sel, h]
    print(f"Output shape: {out.shape}")  # (1024, 8, 64)

⚠️ Reference implementation only. Production speedups need Triton/CUDA kernels. Official code not yet released (July 2026).

Benchmark Results

Tested on DeepSeek-V3.2 and GLM-5.1 — LongBench + RULER:

Method	Indexer Speed	E2E Latency	Accuracy
Dense DSA (baseline)	1×	1×	✅ full
PIVOT-Refine	~3× faster	−28%	✅ ≈ baseline
PIVOT-Reuse	4× faster	−37.5%	⚠️ minor drop

PIVOT-Refine = dense-indexer accuracy + 3× speed. Zero retraining.

Gotchas & Limitations

DSA-only: Works out of the box on DeepSeek-V3.2 / GLM-5.1. Standard full-attention models need DSA fine-tuning first.

Group size is manual: No adaptive strategy in the paper — tune g per model/sequence length.

Batch inference untested: Single-sequence results only.

Prefill vs. decode split missing: The 1.6× E2E number doesn't break down by phase.

Try It Today

Running DeepSeek-V3.2 via vLLM? Request PIVOT integration upstream.
Prototype with the reference code above; profile vs. vanilla DSA on your target seq length.
Stack with IndexCache (arXiv:2603.12201) for cross-layer index reuse on top.
Paper: arXiv:2607.24593

Sources

Hong Liu et al., "PIVOT: Efficient Query-Group Indexing for Token-Level Sparse Attention", arXiv:2607.24593 — https://arxiv.org/abs/2607.24593
"IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse", arXiv:2603.12201 — https://arxiv.org/abs/2603.12201
"MISA: Mixture of Indexer Sparse Attention", arXiv:2605.07363 — https://arxiv.org/abs/2605.07363
"FlashMemory-DeepSeek-V4: Lookahead Sparse Attention", arXiv:2606.09079 — https://arxiv.org/abs/2606.09079
DeepSeek-V3.2 vLLM blog — https://vllm.ai/blog/2025-09-29-deepseek-v3-2

What's your experience with sparse attention in production? Drop a comment.