Skip to content

LILO: Learning to Reason at the Frontier of Learnability

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LILO (Learnability Improves LLMs Optimally) is a curriculum-selection method for RL-based LLM post-training that prioritises questions with high learnability — defined as the variance of success, p̂(1−p̂) — across multiple rollouts. It provides a theoretical upper bound showing that expected policy improvement scales linearly with learnability, and empirically achieves higher final accuracy in 3× fewer training steps than standard RL baselines.

Problem

Standard policy-gradient RL methods (GRPO, PPO, VinePPO) provably gain zero gradient signal from questions that are either too easy (model always succeeds) or too hard (model always fails), because the advantage estimate collapses to zero. This forces practitioners to hand-curate datasets of "appropriate difficulty," wasting significant compute and human effort on zero-learnability training steps.

Method

LILO wraps any existing RL training loop with a rejection-sampling curriculum step. At each training iteration: 1. A candidate pool of |D| = 4×|B| questions is sampled from the dataset. 2. N_learnability = 8 rollouts are generated per candidate to estimate empirical success rate p̂. 3. Learnability p̂(1−p̂) is computed per question; the top-|B| are selected as the training batch. 4. For implementations with large N_train (e.g. VinePPO with N_train ≈ 500), the N_learnability rollouts add only ~4% sampling overhead; for small N_train implementations, the learnability samples are reused for training, adding ~4× sampling cost but no additional training cost.

The theoretical grounding (Theorem 3.1) shows that ‖∇J(θ)‖² is upper bounded by a term proportional to Var[r(s_T)] = p(1−p) for REINFORCE, RLOO, PPO, and VinePPO, directly motivating the curriculum.

Key Contributions

  • Theorem 3.1: Formal proof that expected policy improvement per gradient step is upper bounded by a quantity that scales linearly with learnability for REINFORCE, RLOO, PPO, and VinePPO.
  • LILO algorithm: A simple, algorithm-agnostic curriculum wrapper implementable in <20 lines of code on top of existing training libraries (OatLLM, VinePPO).
  • Empirical validation: First systematic evidence across 3 RL algorithms × 3 datasets × 2 base models that learnability-prioritised training consistently improves final accuracy and training efficiency.
  • Insight on generalisation gap: Shows LILO accelerates training accuracy gains but does not close the persistent train-test gap in LLM reasoning.
  • Curriculum interpretability: Demonstrates that LILO's selected questions correlate with human-interpretable difficulty proxies (gold solution step count) despite having no access to this information.

Results

  • PPO + LILO on MATH: Final test accuracy 19.1% → 21.8%; 2.5× step speed-up.
  • PPO + LILO on GSM8K: 51.1% → 53.2%; 1.9× speed-up.
  • VinePPO + LILO on MATH: 22.8% → 24.9%; 3.2× speed-up.
  • VinePPO + LILO on GSM8K: 53.2% → 55.9% (from 40.5% base); 3.3× speed-up.
  • GRPO + LILO on ORZ57K: 35.5% → 37.1% average over 5 benchmarks; 1.5× speed-up; improvement on 4/5 held-out test sets (MATH, AMC, Minerva, OlympiadBench; AIME unchanged).
  • OOD generalisation: All four PPO/VinePPO × MATH/GSM8K combinations improve on CollegeMATH and OlympiadBench holdouts.
  • Already-tuned model (Oat-Zero-1.5B): LILO finds learnable questions but yields only minimal performance improvement, suggesting saturation after heavy prior RL training.

Limitations

  • For small N_train regimes, learnability estimation adds ~4× sampling overhead, not just 4%.
  • LILO is stateless: learnability is re-estimated from scratch each step with no memory of prior estimates, which is wasteful.
  • Learnable questions become increasingly scarce as training progresses (especially on easier datasets like GSM8K), requiring larger candidate pools (|D| = 8×|B|) or dynamic compute allocation — not yet automated.
  • Learnability cannot distinguish genuinely learnable questions from aleatoric noise (e.g., coin-flip questions with inherently random outcomes).
  • No mechanism to weight questions by relevance to the test distribution; a high-learnability question irrelevant to the test set still gets prioritised.
  • Theory for GRPO's normalised advantage function (At normalised by variance) shows the learnability dependency disappears; the benefit of LILO with GRPO relies on empirical rather than theoretical grounds under that variant.
  • Large train–test accuracy gap persists even with LILO (e.g., ~95% train vs. <60% test on GSM8K); LILO does not address this generalisation failure mode.

Relevance to Agentic AI / LLM Agents

LILO directly targets the RL post-training pipeline that underlies reasoning-focused LLM agents (DeepSeek-R1, O1-style models), providing a principled, compute-efficient curriculum to accelerate capability acquisition. The connection to Unsupervised Environment Design (UED) frames each reasoning problem as an "MDP level," opening a path toward agents that autonomously discover and generate increasingly hard tasks — a prerequisite for open-ended, self-improving agents. The learnability metric is a lightweight signal requiring only rollout variance, making it easily composable with existing agentic training pipelines (GRPO, VinePPO, AsyncRLHF). The authors explicitly position future work as enabling LLM agents that "autonomously and open-endedly expand the frontier of human knowledge," aligning with scalable oversight and continual learning research directions.