한 토큰 롤아웃: 정책 그라디언트로 지도 학습 미세조정 혁신

초록

본 논문은 고정된 지도 학습 데이터가 일반화에 한계를 보이는 원인을 ‘온‑폴리시 데이터 부재’로 규정하고, 토큰 수준에서 정책 그라디언트를 적용한 ‘One‑Token Rollout(OTR)’ 알고리즘을 제안한다. OTR은 현재 모델의 토큰 분포에서 여러 후보를 샘플링하고, 정답 토큰을 보상으로 사용해 토큰별 손실을 구성함으로써, 정적인 오프‑폴리시 데이터에 온‑폴리시 학습 신호를 부여한다. 실험 결과, 수학, 코딩, 일반 추론 등 다양한 벤치마크에서 기존 SFT 대비 일관된 성능 향상을 입증한다.

상세 분석

본 연구는 대형 언어 모델(LLM)의 지도 미세조정(Supervised Fine‑Tuning, SFT)이 강화학습(RL) 대비 일반화 능력이 떨어지는 현상을 ‘데이터의 온‑폴리시 여부’라는 근본적인 차이로 설명한다. SFT는 사전에 수집된 고정 데이터셋을 사용해 오프‑폴리시 방식으로 파라미터를 업데이트하는 반면, RL은 현재 정책이 생성한 데이터(온‑폴리시)를 반복적으로 수집해 학습한다는 점을 강조한다. 이러한 관점에서 저자들은 토큰 하나를 독립적인 ‘단일‑스텝 강화학습 트래젝터리’로 간주하고, 정책 그라디언트 식을 토큰 수준으로 축소한다. 구체적으로, 현재 상태 sₜ(프롬프트와 이전 토큰)에서 정책 π′₍θ₎(·|sₜ) (온도 κ>1을 적용해 탐색을 촉진)로 K개의 후보 토큰 a′ₜ,ⱼ를 샘플링한다. 이후 정답 토큰 xₜ와 비교해 보상 R(a′ₜ,ⱼ, xₜ)를 부여한다. 보상은 정답이면 1, 그 외에는 β(=−0.1)으로 설정해, 정답 토큰이 샘플링될 경우 손실이 SFT와 동일하게 로그우도에 비례하도록 하고, 잘못된 토큰이 샘플링될 경우 로그우도에 β 가중치를 곱해 억제한다. 손실식은

Lₜ^{OTR}(θ)=−(N_gt/K)·log π_θ(xₜ|sₜ)−β·∑_{a′ₜ,ⱼ≠xₜ}log π_θ(a′ₜ,ⱼ|sₜ)

이며, 전체 시퀀스 길이 T에 대해 평균한다. 이 식은 (1) 정답 토큰이 현재 정책에서 얼마나 자주 선택되는가에 따라 동적으로 가중치를 부여해 학습 효율을 높이고, (2) 정책이 잘못된 토큰에 과도히 높은 확률을 부여하는 것을 억제하는 정규화 역할을 수행한다.

알고리즘적 장점은 두 가지이다. 첫째, 온‑폴리시 신호를 얻기 위해 전체 문장을 생성·평가하는 비용이 크게 감소한다. 기존 RL‑HF(PPO) 방식은 수천 토큰 길이의 샘플을 생성하고 보상을 계산해야 하지만, OTR은 토큰 하나당 K개의 후보만 샘플링하므로 계산 복잡도가 O(K·T)이며, 실제 구현에서는 K=4~8 정도로 충분히 성능이 향상된다. 둘째, 기존 SFT와 동일한 학습 파이프라인(AdamW, 동일한 배치·학습률 등)을 유지하면서도 온‑폴리시 효과를 얻으므로, 기존 인프라를 그대로 활용할 수 있다.

실험에서는 OpenR1‑Math‑220k 데이터에서 5k 샘플을 추출해 2 epoch 학습을 수행했으며, Qwen2.5‑3B/7B, Qwen3‑4B/8B, Olmo3‑7B 등 다양한 모델에 적용했다. 수학 벤치마크(GSM8K, Olympiad, AIME 등)에서는 평균 정확도가 SFT 대비 1~4%p 상승했으며, 특히 큰 모델(Qwen3‑8B)에서는 2%p 이상 개선을 기록했다. 코드 생성(HumanEval+, MBPP+)과 일반 추론(SuperGPQA, MMLU‑Pro)에서도 비슷한 패턴이 관찰돼, OTR이 특정 도메인에 국한되지 않고 전반적인 일반화 능력을 향상시킴을 확인했다. 또한, 일부 경우(예: Qwen3‑4B)에서는 SFT가 성능 저하를 보였음에도 OTR이 이를 회복하거나 초과하는 결과를 보여, 온‑폴리시 시뮬레이션이 모델의 사전 지식을 보존하면서 새로운 태스크에 적응하도록 돕는다는 가설을 뒷받침한다.

한계점으로는 보상 설계가 단순(정답 = 1, 그 외 = β)하므로, 복합적인 품질 지표(예: 코드 실행 성공률, 논리적 일관성)를 직접 반영하기 어렵다는 점이다. 또한, 토큰 수준 샘플링이 문맥 전체의 상호작용을 충분히 포착하지 못할 가능성도 존재한다. 향후 연구에서는 베타 값을 동적으로 조정하거나, 문장 수준 롤아웃과 토큰 수준 OTR을 혼합하는 하이브리드 방식을 탐색할 여지가 있다.

전반적으로 OTR은 ‘데이터‑중심’ 관점에서 SFT와 RL 사이의 격차를 메우는 실용적인 접근법이며, 기존 SFT 파이프라인에 최소한의 수정만으로 온‑폴리시 효과를 도입할 수 있다는 점에서 LLM 미세조정 연구에 중요한 전환점을 제공한다.