Learning to Reason at the Frontier of Learnability (2502.12272v3)

Abstract: Reinforcement learning is now widely adopted as the final stage of LLM training, especially for reasoning-style tasks such as maths problems. Typically, models attempt each question many times during a single training step and attempt to learn from their successes and failures. However, we demonstrate that throughout training with two popular algorithms (PPO and VinePPO) on two widely used datasets, many questions are either solved by all attempts - meaning they are already learned - or by none - providing no meaningful training signal. To address this, we adapt a method from the reinforcement learning literature - sampling for learnability - and apply it to the reinforcement learning stage of LLM training. Our curriculum prioritises questions with high variance of success, i.e. those where the agent sometimes succeeds, but not always. Our findings demonstrate that this curriculum consistently boosts training performance across multiple algorithms and datasets, paving the way for more efficient and effective reinforcement learning with LLMs.

• The paper introduces Sampling for Learnability (SFL), which targets LLM reasoning by prioritizing questions with high success variability.
• It adapts reinforcement learning curricula using techniques from Unsupervised Environment Design to optimize training via dynamic difficulty levels.
• Empirical results demonstrate faster convergence and improved accuracy on MATH and GSM8K datasets, underscoring SFL's practical advantages.

Learning to Reason at the Frontier of Learnability

The paper "Learning to Reason at the Frontier of Learnability" by Foster and Foerster addresses a crucial and often overlooked challenge in the training of LLMs geared towards reasoning tasks, such as mathematics problem-solving. Reinforcement Learning (RL) is an increasingly favored training regime for such models, as the complexity of reasoning tasks benefits from iterative learning methods. However, this paper reveals an inefficiency in conventional training approaches, which either solve a question in every attempt or fail consistently, yielding sparse and uninformative gradients for policy updates.

This work presents an adaptation of a known technique from the RL domain, termed Sampling for Learnability (SFL), applied to fine-tune LLMs. Unlike traditional training patterns that indiscriminately sample questions, SFL emphasizes questions showing a high variance in the success rate, thereby prioritizing questions on the cusp of being solved correctly. Such a strategy theoretically supplies a richer informational signal for model adaptations, guiding models more effectively along the frontier of learnability.

Methodological Insights

The authors draw on concepts from Unsupervised Environment Design (UED), repurposing them for LLM fine-tuning. By treating the training environment as an underspecified Markov Decision Process (MDP), they employ SFL to dynamically sculpt a curriculum that evolves with the model's capabilities. This approach leverages UED's autotelic parallelism to define the curriculum design as a two-player zero-sum game, allowing for adaptive levels of complexity without explicit supervision.

The main innovation lies in evaluating question difficulty based on its success variability per attempt, as sampled across multiple trajectories (or rollouts). This embeds a statistical rationale into curriculum design exploiting the principle of maximal variance (i.e., learnability), optimizing interactions between model state and problem difficulty.

Empirical Results

Empirical evaluations are conducted using Proximal Policy Optimization (PPO) and VinePPO across two datasets: MATH, featuring competition-grade questions, and GSM8K, comprising simpler arithmetic problems. The paper outlines improved performance metrics in both training speed and final accuracy with SFL, demonstrating enhanced generalization capabilities on out-of-distribution data such as College-Math and OlympiadBench datasets.

Significant findings include that SFL promotes a more efficient RL finetuning cycle: models attain desired solution rates more rapidly, achieving superior train and test accuracies. The authors present data suggesting that SFL's incremental computational overhead (arising from additional rollouts) is justifiably offset by decreased model training time, making it a valuable addition to existing LLM training regimes.

Implications and Future Directions

The paper suggests implications that are both practical and theoretical. On a practical level, the integration of SFL into LLM reinforcement learning frameworks aids in overcoming initial sampling inefficiencies, reducing computational expenses over exhaustive training regimes. Theoretically, by redistributing focus towards medium-difficulty questions, SFL may encourage LLMs to adopt more generalized reasoning schemas over implicit rote-learning patterns common in simpler paradigms.

Potential paths for future research include applying SFL to other RL algorithms such as preference-based optimization or hierarchical learning frameworks to observe its impacts on decision-making models that operate over non-binary reward matrices. Incorporating SFL's principles into varied domains and advocating for different problem distributions might further inform methodology, leading to robust learning models capable of complex reasoning tasks.

Thus, this paper orchestrates an intelligent intersection of curriculum design and reinforcement learning paradigms, arguing convincingly for dynamically learnability-focused sampling strategies. As LLM complexity scales, such considerations might become cornerstone strategies, contributing materially to the field of machine reasoning and adaptive learning.