원-토큰 롤아웃: 정책 그래디언트를 활용한 대형 언어 모델의 지도 미세 조정 가이드
One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy Gradient
September 30, 2025
저자: Rui Ming, Haoyuan Wu, Shoubo Hu, Zhuolun He, Bei Yu
cs.AI
초록
지도 미세 조정(SFT)은 대규모 언어 모델(LLM)을 적응시키는 데 주로 사용되는 방법이지만, 강화 학습(RL)에 비해 일반화 능력이 부족한 경우가 많습니다. 본 연구에서는 이러한 성능 차이가 단순히 손실 함수의 차이뿐만 아니라 더 근본적인 차이에서 비롯된다고 주장합니다: SFT는 미리 수집된 고정된 데이터셋에서 학습하는 반면, RL은 현재 정책에서 샘플링된 온-정책 데이터를 활용합니다. 이 가설을 바탕으로, 우리는 정책 경사법을 통해 SFT를 안내하는 새로운 미세 조정 알고리즘인 단일 토큰 롤아웃(OTR)을 제안합니다. OTR은 자동회귀 학습 과정을 재구성하여 각 토큰 생성을 단일 단계의 강화 학습 궤적으로 취급합니다. 각 단계에서 현재 정책의 분포로부터 여러 후보 토큰을 샘플링하여 몬테카를로 "롤아웃"을 수행합니다. 지도 데이터의 실제 토큰은 이러한 샘플에 대한 보상 신호로 사용됩니다. 정책 경사법의 지도 하에, 우리의 알고리즘은 정적이고 오프-정책인 지도 데이터를 토큰 수준에서 동적이고 온-정책인 신호로 재구성함으로써, 온-정책 학습의 일반화 이점을 포착하면서도 전체 문장 생성을 위한 비용이 많이 드는 오버헤드를 피합니다. 수학적 추론, 코드 생성, 일반 도메인 추론 등 다양한 도전적인 벤치마크에 대한 광범위한 실험을 통해, OTR이 표준 SFT를 지속적으로 능가함을 입증합니다. 우리의 연구 결과는 OTR을 LLM 미세 조정을 위한 강력하고 실용적인 대안으로 확립하고, 데이터의 온-정책 특성이 일반화의 중요한 동인임을 설득력 있게 보여줌으로써, LLM 미세 조정을 위한 새로운 방향을 제시합니다.
English
Supervised fine-tuning (SFT) is the predominant method for adapting large
language models (LLMs), yet it often struggles with generalization compared to
reinforcement learning (RL). In this work, we posit that this performance
disparity stems not just from the loss function, but from a more fundamental
difference: SFT learns from a fixed, pre-collected dataset, whereas RL utilizes
on-policy data sampled from the current policy. Building on this hypothesis, we
introduce one-token rollout (OTR), a novel fine-tuning algorithm that guides
SFT with the policy gradient method. OTR reframes the autoregressive learning
process by treating each token generation as a single-step reinforcement
learning trajectory. At each step, it performs a Monte Carlo ``rollout'' by
sampling multiple candidate tokens from the current policy's distribution. The
ground-truth token from the supervised data is then used to provide a reward
signal to these samples. Guided by policy gradient, our algorithm repurposes
static, off-policy supervised data into a dynamic, on-policy signal at the
token level, capturing the generalization benefits of on-policy learning while
bypassing the costly overhead of full sentence generation. Through extensive
experiments on a diverse suite of challenging benchmarks spanning mathematical
reasoning, code generation, and general domain reasoning, we demonstrate that
OTR consistently outperforms standard SFT. Our findings establish OTR as a
powerful and practical alternative for fine-tuning LLMs and provide compelling
evidence that the on-policy nature of data is a critical driver of
generalization, offering a promising new direction for fine-tuning LLMs.