AMFT: Aligning LLM Reasoners by Meta-Learning the Optimal
  Imitation-Exploration Balance
Authors: name: Lixuan He; name: Jie Feng; name: Yong Li
Published: 2025-08-09 11:40:54

Abstract: Large Language Models (LLMs) are typically fine-tuned for reasoning tasks
through a two-stage pipeline of Supervised Fine-Tuning (SFT) followed by
Reinforcement Learning (RL), a process fraught with catastrophic forgetting and
suboptimal trade-offs between imitation and exploration. Recent single-stage
methods attempt to unify SFT and RL using heuristics, but lack a principled
mechanism for dynamically balancing the two paradigms. In this paper, we
reframe this challenge through the theoretical lens of implicit
rewards, viewing SFT and RL not as distinct methods but as complementary
reward signals. We introduce Adaptive Meta Fine-Tuning (AMFT), a novel
single-stage algorithm that learns the optimal balance between SFT's implicit,
path-level reward and RL's explicit, outcome-based reward. The core of AMFT is
a meta-gradient adaptive weight controller that treats the SFT-RL
balance as a learnable parameter, dynamically optimizing it to maximize
long-term task performance. This forward-looking approach, regularized by
policy entropy for stability, autonomously discovers an effective training
curriculum. We conduct a comprehensive evaluation on challenging benchmarks
spanning mathematical reasoning, abstract visual reasoning (General Points),
and vision-language navigation (V-IRL). AMFT consistently establishes a new
state-of-the-art and demonstrats superior generalization on out-of-distribution
(OOD) tasks. Ablation studies and training dynamic analysis confirm that the
meta-learning controller is crucial for AMFT's stability, sample efficiency,
and performance, offering a more principled and effective paradigm for LLM
alignment.Our codes are open-sourced via https://github.com/hlxtsyj/AMFT.