LAD — Learnable Advantage Density

Inspiration

Frontier labs spend millions buying curated data to RL-finetune reasoning models — but they only learn whether a batch worked after running the full, expensive GRPO training. Reading Bae et al.'s EACL 2026 paper, we noticed Proposition 3.1 had quietly proved something powerful: the KL divergence between the initial policy and the RL-optimal one is lower-bounded by p(1-p), the Bernoulli variance of the pass-rate. That bound is everywhere in modern RL — DAPO, RL-ZVP, difficulty-aware staged RL all use it online to filter batches mid-run. Nobody had repurposed it as a pre-purchase, dataset-level value score. That gap was the whole project.

What it does

LAD (Learnable Advantage Density) is a cheap, training-free metric that predicts how much a data cohort will improve a model under GRPO — before you spend a single dollar on training. From ~8 base-model rollouts per task and zero gradients, it scores:

• Advantage energy p̂(1-p̂) — is there a learning signal at all?
• Headroom (1-p̂)^γ — is the difficulty pointed where the model can actually rise?
• Effective diversity (Vendi/|C|)^β — is the cohort varied or just paraphrased duplicates?

Combined: LAD(C) = mean_T [ p̂(1-p̂)·(1-p̂)^γ ] · (Vendi/|C|)^β. One number, computed in minutes, that tells a data marketplace or lab whether a cohort is worth buying.

How we built it

• Base model: Qwen2.5-1.5B-Instruct (general, not Qwen-Math — avoids the spurious-reward confound from Shao et al.).
• Benchmark: GSM8K with an exact-match numeric verifier — clean binary reward.
• Cohorts: 19 size-matched cohorts at |C|=256, each varying one property (difficulty bands, bimodal vs unimodal, diversity levels, noisy-label, adversarial).
• Oracle: identical-config GRPO (TRL's GRPOTrainer + vLLM rollouts) per cohort on 4×GB200, measuring lift = accuracy_after − accuracy_before on a fixed held-out test set. Cohort is the only variable.
• Validation: leave-one-cohort-out (LOCO) — fit metric → lift on n-1 cohorts, predict the held-out one. Headline metric: held-out Spearman ρ.
• Embeddings + Vendi: sentence-transformers + cosine similarity kernel → exp-entropy of eigenvalues = effective number of distinct tasks.

Challenges we ran into

• Variance alone isn't enough. Naive p(1-p) predicts signal, but lift is asymmetric — a cohort at p=0.3 has more room to rise than one at p=0.7 with the same variance. We had to derive the headroom correction (1-p̂)^γ.
• Diversity bites back. 256 paraphrases of one problem at p=0.5 have huge summed advantage energy but teach one thing. Vendi/|C| as a redundancy ceiling was the second principled fix.
• Mechanistic-vs-correlational. A held-out correlation isn't enough — judges (rightly) ask "but does selecting by LAD actually train better?" We added a causal intervention: top-LAD vs random vs bottom-LAD selection, all trained identically.
• Compute juggling. 19 controlled GRPO runs in the hackathon window meant 4-cohort-parallel waves on the GB200 box, careful detached tmux orchestration, and a CPU-only smoke test so the analysis pipeline could be debugged off the GPU.

Accomplishments that we're proud of

• Held-out Spearman ρ ≈ 0.89 vs the best cheap competitor (reward variance) at 0.45 — nearly 2×.
• LAD beats expensive methods too — gradient-influence and small-scale-RL proxies sit at 0.72–0.84 and cost ~125× more. We're a vertical jump up the cost↔quality frontier.
• Causal proof, not just correlation: top-LAD selection delivered +12.6 accuracy points vs +6.0 for random and +0.9 for bottom-LAD on the same training budget.
• Mechanistic agreement: LAD correlates with the in-training group-advantage variance at ρ ≈ 0.86 — it's measuring the same thing the optimizer actually uses.
• A live web demo where you can score a cohort interactively and see exactly where it lands on the cost-vs-accuracy frontier.

What we learned

• Pick the proven quantity, then earn the corrections. Starting from a theorem (Bae et al.'s lower bound) instead of a vibe meant every additional term — headroom, diversity — had to demonstrate its failure mode. Variance mis-ranks noisy-label and bimodal cohorts; LAD doesn't. Each correction earned its place via ablation.
• Held-out validation is non-negotiable. It's tempting to fit-and-report. LOCO with Spearman is what made the result trustworthy.
• The cheapest signal often wins. We assumed gradient-based methods would dominate — they don't, because their cost scales with model size while LAD's doesn't.
• Plumbing is half the project. Identical-per-cohort GRPO (down to seed, schedule, and verifier) is what makes "cohort = the only variable" actually true. Most of our debugging was making that equality hold.

What's next for LAD - Learnable Advantage Density

• Scale up. Re-run on Llama-3.2-3B and 7B — the derivation is scale-independent, but we want measured numbers, not extrapolation.
• Cross-domain. Code (HumanEval+/MBPP+) is the natural next verifier, then constrained-format tasks. The method is domain-agnostic given any verifier.
• Online LAD. Recompute LAD periodically during training as p̂ shifts — turning a static screening tool into a curriculum scheduler.
• A real data-screening service. Wrap the metric behind an API: upload a dataset, get an expected-lift score and a per-task highlight of which prompts are pulling their weight. Price data by payoff, not by token count.
• Adversarial robustness. Push harder on the broken-verifier and label-noise frontiers — LAD degrades gracefully today, but we want guarantees.

Built With

• anthropic
• primeintellect
• python
• pytorch
• vllm