Trust-Region Adaptive Policy Optimization (2512.17636v1)

Abstract: Post-training methods, especially Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), play an important role in improving LLMs' (LLMs) complex reasoning abilities. However, the dominant two-stage pipeline (SFT then RL) suffers from a key inconsistency: SFT enforces rigid imitation that suppresses exploration and induces forgetting, limiting RL's potential for improvements. We address this inefficiency with TRAPO (\textbf{T}rust-\textbf{R}egion \textbf{A}daptive \textbf{P}olicy \textbf{O}ptimization), a hybrid framework that interleaves SFT and RL within each training instance by optimizing SFT loss on expert prefixes and RL loss on the model's own completions, unifying external supervision and self-exploration. To stabilize training, we introduce Trust-Region SFT (TrSFT), which minimizes forward KL divergence inside a trust region but attenuates optimization outside, effectively shifting toward reverse KL and yielding stable, mode-seeking updates favorable for RL. An adaptive prefix-selection mechanism further allocates expert guidance based on measured utility. Experiments on five mathematical reasoning benchmarks show that TRAPO consistently surpasses standard SFT, RL, and SFT-then-RL pipelines, as well as recent state-of-the-art approaches, establishing a strong new paradigm for reasoning-enhanced LLMs.

• The paper introduces TRAPO, integrating SFT and RL with trust-region clipping to stabilize and improve policy optimization.
• It employs adaptive prefix guidance and micro-group sampling to enhance exploration and sample efficiency.
• Numerical results on reasoning benchmarks demonstrate significant gains over traditional sequential SFT and RL approaches.

Trust-Region Adaptive Policy Optimization: A Technical Analysis

Motivation and Background

The prevailing post-training protocols for LLMs—namely Supervised Fine-Tuning (SFT) followed by Reinforcement Learning (RL)—are suboptimal for advanced reasoning tasks due to intrinsic incompatibilities. Specifically, SFT enforces strong imitation behavior, impeding exploration in subsequent RL phases and instigating catastrophic forgetting. RL, when initiated from an SFT-overfitted checkpoint, is forced to overcome entrenched behavioral priors, resulting in inefficient skill propagation and loss of pretrained knowledge. The paper introduces Trust-Region Adaptive Policy Optimization (TRAPO) to resolve these tensions by unifying SFT and RL within each training instance, rather than serializing them.

TRAPO Framework Overview

TRAPO entails an instance-level fusion of SFT and RL objectives, operationalized via:

• Trust-Region Supervised Fine-Tuning (TrSFT): SFT gradients are clipped within a dynamically computed trust region, determined by the probability assigned to each token by the current policy. Outside this region, optimization strength is attenuated, functionally interpolating between forward KL (mode-covering) and reverse KL (mode-seeking) regimes. The mechanism prevents distribution blending, stabilizes updates, and greatly mitigates the dilution effect of SFT in policy composition.
• Adaptive Prefix Guidance: The model’s need for expert demonstration is estimated per-prompt, allocating expert prefixes only when autonomously generated rollouts yield low return. Prefix length is dynamically increased according to micro-group sampling—ensuring minimal necessary guidance while maximizing self-exploration.

This "learn-while-practicing" paradigm augments exploration with targeted expert supervision, facilitating both knowledge distillation and policy improvement.

Algorithmic Details

TRAPO leverages expert trajectories for dual objectives:

• Internalization: Expert prefixes guide skill acquisition. TrSFT loss is applied only to prefix tokens within the trust region, ensuring robust knowledge transfer without destabilizing policy gradients.
• Exploration and RL: The suffix generated freely by the current policy is optimized via a standard RL objective (e.g., GRPO sans KL penalty). This interpolation enforces synergy rather than conflict between supervised and reinforcement objectives.

Micro-group sampling segments training into N sub-batches per prompt; each sub-batch has its own prefix length ratio and performance threshold, governing the extent of expert guidance. This lends high granularity to supervision allocation, substantially improving sample efficiency.

Main Results

Numerical Validation

TRAPO exhibits strong quantitative improvements across five mathematical reasoning benchmarks (AIME2024, AMC, MATH-500, Minerva, OlympiadBench):

• Average Gains: +6.3 points over SFT; +6.2 over pure RL (GRPO); +2.3 over conventional SFT-then-RL. The direct fusion of SFT and RL losses without trust-region clipping resulted in performance collapse (-18 points vs RL baseline), highlighting the necessity of TrSFT.
• Generalization: TRAPO delivers superior performance on ARC-c and MMLU-Pro (general-domain benchmarks), indicating that the method does not induce rigid reasoning or overfit to mathematical prompts.

Pass@$k$ analysis reveals that TRAPO not only selects optimal solutions from its pretraining solution space but also expands this space—unlike pure RL protocols, which primarily refine solution selection.

Training Dynamics

TRAPO converges to higher reward levels, produces longer and more complex outputs early in training (internalizing advanced reasoning behaviors), and maintains higher policy entropy long-term, signifying balanced refinement and ongoing openness to external guidance.

Ablation Studies

• Micro-group Sampling: Adaptive guidance alone outperforms static RL; in combination with TrSFT, results are maximized. Standard SFT or LUFFY-style offline losses degrade or offer limited improvement, confirming the theoretical advantage of trust-region clipping and instance-level guidance selection.
• Trust-Region Parameter: Optimal accuracy is achieved with moderate trust-region width ($\alpha = 0.1$), balancing safe knowledge instillation and avoidance of excessive mode blending.

Theoretical Implications

The principal theoretical contribution is diagnosing the incompatibility between forward KL (mode-covering) induced by SFT and RL’s requirement for mode-seeking behavior. TrSFT’s gradient-clipping mechanism shifts optimization towards reverse KL, ensuring the policy focuses on core expert modes and avoids void regions that induce output degeneracy.

This establishes a rigorous paradigm for hybrid model post-training, substantiated by both theoretical derivations and comprehensive empirical validation.

Practical Implications and Future Directions

TRAPO demonstrates:

• Superior reasoning acquisition: Efficient transfer and simultaneous exploration, yielding state-of-the-art performance across mathematical and general reasoning domains.
• Sample efficiency: Higher reward under matched GPU-hour budgets, with reduced dependence on expert prefixes as policy improves.
• Robust generalization: Avoids over-conditioning while reliably internalizing transferable skills.

Future research should extend TRAPO’s principles to domains requiring more intricate reasoning (multi-modal, logic-rich prompts), investigate alternative trust-region formation strategies, and optimize micro-group structures for larger LLM architectures. There is substantial opportunity to generalize adaptive guidance strategies beyond mathematical reasoning, potentially enhancing LLM robustness in open-ended dialogue and scientific synthesis tasks.

Conclusion

TRAPO, with its trust-region SFT and adaptive expert guidance, provides a principled solution to longstanding inefficiencies in LLM post-training. Its instance-level integration of supervised and reinforcement objectives, grounded in mode-seeking optimization, delivers definitive improvements in reasoning ability, sample efficiency, and generalization. The framework sets a clear direction for future research on unified, theory-aligned LLM post-training and reasoning enhancement (2512.17636).