Technical Summary: LARK (Learnability-Grounded Trajectory Selection)

Problem Statement

The paper addresses the challenge of reasoning distillation, a paradigm where a smaller student Large Language Model (LLM) is fine-tuned to imitate reasoning trajectories generated by larger teacher models. While existing data selection methods often operate at the question level or rely on heuristics such as trajectory quality, correctness, or model confidence, they frequently overlook a critical factor: learnability.

Specifically, for a single question, multiple candidate reasoning trajectories may exist (generated by different teachers, sampling runs, or styles). Even if these trajectories lead to the same correct final answer, they provide vastly different training signals. A trajectory that appears "high-quality" or "natural" is not necessarily the one from which the student model can learn most efficiently. The central question posed is: Can we design a principled trajectory selection criterion that ensures distilled reasoning trajectories are learnable by the student model?

Methodology: LARK

The authors propose LARK (Learnability-grounded Anchor-time Ranking), a method that selects trajectories based on an optimization-theoretic definition of learnability rather than static quality metrics.

1. Learnability as an Optimization Objective

LARK frames trajectory selection as a policy optimization problem. It defines a learnability factor $\rho(\theta; q)$, which characterizes the rate at which the student's training loss decreases. This factor is defined as the ratio of the squared gradient norm to the loss:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
where $q$ is the weight vector over $K$ candidate trajectories. Under the Polyak–Łojasiewicz (PL) condition, a higher $\rho$ implies faster convergence.

2. Anchor-Time Approximation

Directly optimizing the time-integrated learnability objective over the entire fine-tuning trajectory is intractable. LARK introduces the anchor-relative condition, which posits that for over-parameterized models in the "lazy training" regime, the learnability rate remains proportional to its value at the initial anchor parameters $\theta_{ref}$. This allows the objective to be approximated by maximizing $\rho(\theta_{ref}; q)$.

3. Local Linearization and Forward-Pass Proxy

Maximizing $\rho(\theta_{ref}; q)$ directly is problematic because it is a quasi-convex function that collapses to a degenerate "one-hot" solution (selecting only a single trajectory). To avoid this, LARK:

• Local Linearization: Expands $\rho$ around the uniform distribution $p$ (standard SFT without selection) using a first-order Taylor expansion.
• Forward-Pass Surrogate ($\hat{g}_k$): Instead of computing exact gradients (which require backward passes for every candidate), LARK derives a computable proxy $\hat{g}_k$ using only forward passes. This proxy estimates the local gradient of the learnability rate based on the student's token-level prediction residuals and cross-entropy loss.
  • Specifically, $\hat{g}_k$ favors trajectories where the student is "confidently wrong" in a structured way: the loss $\ell_k$ is high (indicating poor prediction), but the residual mass is concentrated on a few "near-miss" tokens (providing a clear correction signal).

4. $\chi^2$-Regularized Selection Policy

To prevent the selector from "hacking" the learnability criterion by drifting too far from the original data distribution (which harms generalization), LARK imposes a $\chi^2$-regularization term. This leads to a closed-form, budget-parameterized selection rule:

• Trajectories are ranked by their score $\hat{g}_k$.
• The top $B$ trajectories are selected.
• Unlike hard top-$B$ truncation, LARK assigns soft weights proportional to the score margin above the $(B+1)$-th trajectory. This weighting rule is derived to maximize the lower bound of the learnability improvement while controlling the $\chi^2$ divergence from the uniform prior.

Key Contributions

1. Formulation: The authors formulate reasoning trajectory selection as a learnability-grounded policy optimization problem, introducing the anchor-time learnability rate $\rho$ as a principled objective.
2. Algorithm: They derive a closed-form, budget-parameterized selection rule that uses a forward-pass surrogate $\hat{g}_k$ and $\chi^2$-regularization. This avoids expensive per-trajectory backward passes while theoretically guaranteeing bounds on estimation error and generalization.
3. Empirical Validation: The paper demonstrates that LARK consistently outperforms existing baselines (including heuristic, quality-based, and student-side scoring methods) across multiple base models (Qwen-2.5-7B/1.5B, Llama-3.2-3B) and reasoning benchmarks (AIME, AMC, GPQA, MATH-500).

Experimental Results

• Performance: LARK achieved the highest average Acc@5 across all tested model-budget settings. For example, on Qwen-2.5-7B with a budget of $B=3$, LARK improved performance by 7.07 percentage points over the strongest baseline (RSR).
• Efficiency: LARK is a forward-pass-only method, matching the computational cost of the lowest-cost student-side baselines, whereas exact gradient-based methods would require significantly more computation.
• Diagnostic Analysis:
  • Loss Decay: Trajectories selected by LARK induced the fastest reduction in SFT loss during training.
  • Predictive Validity: The LARK score $\hat{g}_k$ showed a strong positive correlation (Pearson $r=0.69$) with downstream accuracy when aggregated by teacher model.
  • Token-Level Behavior: Selected trajectories exhibited lower correct-token probabilities but higher "wrong-token concentration," confirming that LARK selects examples where the student is wrong but the error is structured and correctable.

Significance and Claims

The paper claims that learnability is a distinct and effective principle for efficient reasoning distillation, separate from trajectory quality or difficulty. By grounding selection in the optimization dynamics of the student model, LARK identifies trajectories that provide the most efficient learning signal.

The authors emphasize that their method:

• Preserves Generalization: Through $\chi^2$-regularization, it prevents the selection policy from overfitting to a narrow set of "easy" or "hacked" trajectories.
• Is Theoretically Grounded: The selection rule is derived from first principles of optimization theory (PL conditions, gradient flow) with explicit error bounds.
• Is Practical: It requires no hyperparameter tuning for the selection rule itself (the regularization strength is determined by the budget $B$) and avoids the computational overhead of backward passes.

The work suggests that future extensions could apply this learnability-grounded perspective to noisier candidate pools and broader reasoning domains, but the current scope is strictly limited to reasoning distillation in mathematical problem solving.