파운데이션 모델 전 단계를 단일 우도 학습으로 통합하는 훈련법

무엇인가

이 논문이 다루는 문제는 파운데이션 모델 훈련이 사전학습, 미드트레이닝, SFT, RL로 분절돼 있다는 관행이다. 각 단계는 데이터와 학습 알고리즘이 다르고, 레시피와 인프라도 따로 발전해 왔다. 저자들은 이 분리가 본질적이지 않다고 보고, 전체 훈련을 하나의 우도 기반 학습 목표로 재구성한다. 그러면 설계의 중심은 단계별 알고리즘이 아니라 어떤 데이터를 어떤 비율로 언제 학습할지로 바뀐다. 이 관점은 단계별 지표가 아니라 최종 모델과 훈련 궤적 전체로 의사결정을 평가하게 해 준다.

어떻게 동작하나

이론의 핵심은 MaxRL을 우도 적합으로 읽는 것이다. 이진 보상에서 pθ(x)=E_{y~πθ}[r(x,y)]이고 πθ+(y|x)=πθ(y|x,r=1)라 하면, ∇θ log pθ(x)=E_{y~πθ+}[∇θ log πθ(y|x)]이다. 따라서 t 시점의 MaxRL 갱신은 현재 정책의 성공 조건부 분포 πθ_t+를 고정한 뒤 그 분포에 최대우도 적합하는 것과 같다. 이를 일반화해 각 데이터 원천 s가 t 시점에 프롬프트 분포 ρ_s,t와 목표 분포 μ_s,t를 갖고, 혼합 가중치 λ_s,t가 합 1이라고 두면 q_t=Σ_s λ_s,t(ρ_s,t∘μ_s,t)이고 통합 손실은 L(θ,q_t)=-E_{(x,y)~q_t}[log πθ(y|x)]이다. 미분 시 원천 분포와 혼합 가중치는 고정한다. 이는 원천 s를 확률 λ_s,t로 뽑아 예제를 만들고 같은 우도 손실을 적용하는 것과 같다. 기존 단계식 훈련은 한 번에 한 원천에 가중치 1을 주는 특수 사례이고, 통합 훈련은 원천들이 겹치고 비율이 변하는 것을 허용한다. 또한 현재 모델이 아닌 고정 생성 정책 ν의 성공 조건부 분포 ν+를 같은 우도로 적합할 수 있어 생성과 최적화를 분리할 수 있다. ν=πθ_t이면 on-policy MaxRL이 되고, 다른 생성 정책이면 같은 손실을 쓰되 현재 정책의 MaxRL 그래디언트와는 달라진다. REINFORCE는 같은 성공 조건부 목표를 쓰지만 유효 프롬프트 분포가 ρ_RF(x)∝ρ(x)pθ_t(x)라 성공 확률이 높은 프롬프트에 더 큰 가중을 준다. 연속 보상은 TailRL로 확장한다.

무엇과 다른가

구현은 공통 데이터 형식, 손실, 트레이너로 이뤄진다. 각 예제는 토큰 시퀀스 z와 토큰별 가중치 w의 쌍 (z,w)로 표현한다. 사전학습과 미드트레이닝은 코퍼스 토큰에 단위 가중치, SFT는 데모 토큰에 단위 가중치와 프롬프트 토큰에 0, RL은 같은 프롬프트 마스크와 보상·표본 통과율로 계산한 응답 가중치 A를 쓴다. 배치 B의 손실은 -1/T_B Σ_{(z,w)∈B} Σ_l w_l log πθ(z_l|z_<l)이고 T_B는 프롬프트가 아닌 토큰 수다. 생성기는 최근 모델 복사본 πθ_s에서 프롬프트당 N개 완성을 샘플링하고, 이진 보상이면 통과율 p̂=1/N Σ r(x,y_i)를 추정한다. 정답 완성의 응답 토큰에는 1/p̂, 오답에는 0을 준다. 이는 MaxRL의 유한 표본 형태로, 정답 완성들에 대한 평균 스코어가 되고 모두 오답이면 갱신이 0이다. 생성 정책 ν의 경우 프롬프트 가중치는 1-(1-pν(x))^N이며, N=1이면 REINFORCE의 pν(x) 가중으로 돌아가고 N이 커질수록 1에 가까워진다. 정적 데이터와 RL 생성은 공유 버퍼에서 섞이며, 목표 토큰 비율 f_s에 대해 예제당 평균 토큰 수 T_s로 f_s/T_s에 비례해 샘플링한다. 생성과 최적화는 비동기로 돌고 트레이너가 k번 갱신마다 가중치를 배포하며, 정책 지연(staleness)과 처리량은 부록에서 다룬다.

어떻게 쓰나

통제 실험은 124M 파라미터 모델을 처음부터 학습한다. 사전학습 코퍼스는 20B 토큰, 미드트레이닝 코퍼스는 510M 토큰으로 고정하고, Reasoning Gym 스타일의 검증 가능한 자연어 추론 퍼즐을 쓴다. 단계식은 사전학습→미드트레이닝→RL 순서이고, 통합식은 미드트레이닝 데이터를 사전학습 마지막 구간과 RL 전 구간에 섞는다. Base는 미드트레이닝에 나온 퍼즐 크기, Extended는 더 큰 퍼즐을 RL에 추가하고 그보다 큰 held-out 퍼즐로 일반화를 본다. MaxRL을 Base 6000회, Extended 12000회 갱신하고 RL 결과는 두 시드 평균, 지표는 pass@1과 pass@16이다. 저자들이 보고한 결과는 최선의 통합 구성이 Base와 Extended, 그리고 미드트레이닝·RL 어디에도 없던 더 큰 퍼즐 일반화에서 단계식을 능가했다는 것이다. 미드트레이닝 토큰을 사전학습에 더 많이 섞을수록 좋았고, 2.14B 혼합에서는 Base 최고 정확도가 중간 정도의 RL 혼합에서 나왔으며 Extended와 일반화는 RL 중 미드트레이닝을 더 많이 섞을수록 대체로 좋아졌다. 통합은 미드트레이닝 코퍼스를 2~4에포크 재사용해도 성능이 나았지만, 단계식은 미드트레이닝 예산을 늘리면 최종 과제 정확도가 떨어졌다. 총 연산을 맞춘 비교에서는 단계식이 중간 예산의 정확도에서 유리했고, 통합식은 더 큰 예산과 가장 작은 예산에서 더 높은 정확도를 보였으며 일반화에서는 거의 모든 예산에서 최고 정확도와 커버리지를 냈다. 망각 비교에서는 사전학습-미드트레이닝을 섞은 다섯 통합 구성 모두 사전학습만 한 모델과 비슷한 held-out 손실과 다섯 일반 벤치마크 평균을 유지했고, 테스트한 단계식 여섯 구성 중 망각이 가장 적은 구성도 held-out 손실이 상당히 높고 평균 성능이 낮았다. 원문 서술에는 Fig. 4~6의 절대 수치나 개선폭이 제시되지 않아 백분율은 확인할 수 없다.

전제와 한계

실제 데이터 실험은 이질적인 미드트레이닝 코퍼스와 수학·코딩 RL을 쓴다. Math는 MATH-500과 GSM8K 평균, Coding은 HumanEval+와 MBPP+ 평균이며 벤치마크 가중치는 같다. 미드트레이닝 비율을 0~5%로 고정해 늘리면 수학 정확도는 큰 변화 없이 아주 조금 낮아지고, 가장 큰 비율이 가장 높은 커버리지를 보였다. 망각 억제 효과는 더 분명해서 0이 아닌 모든 혼합이 held-out 미드트레이닝 손실을 베이스 모델보다 낮췄고 비율이 클수록 손실이 더 낮았으며, 단계식은 오히려 손실을 높였다. SciQ 정확도도 대체로 혼합과 함께 개선됐다. 반면 미드트레이닝 데이터를 많이 섞을수록 Coding 정확도는 낮아졌지만 Coding 커버리지는 베이스와 단계식보다 높았다. 즉 이질적 미드트레이닝 혼합은 망각을 줄이지만 과제 성능에서는 정확도와 커버리지의 trade-off가 나타난다. 평균 0.5% 혼합에서 상수, 선형 감소, 선형 증가, 중간에 0으로 감소하는 스케줄을 비교하면, 중간에 섞기를 끝내는 스케줄은 상수 혼합보다 최종 Coding과 SciQ 정확도가 높았고 SciQ 이득도 유지됐지만 held-out 코퍼스 손실 감소는 덜했다. 수학 정확도와 커버리지는 네 스케줄에서 비슷했다. 연속 학습에서는 0.5, 1, 2, 5% 통합 모델이 같은 코딩 RL을 받았을 때 초기 코딩 정확도는 낮았지만 최종 정확도와 MBPP+ 커버리지가 단계식보다 높았고, 5% 모델이 네 지표 모두 가장 좋았다. HumanEval+ 커버리지는 작은 혼합에서 단계식과 비슷하거나 낮았다. 2% 혼합에서 REINFORCE는 Math 정확도가 높고 MaxRL은 Math 커버리지가 높았으며, REINFORCE가 held-out 손실과 SciQ, ARC-Challenge 정확도에서도 유리했다. KL 정규화는 held-out 손실 악화를 줄였지만 강한 패널티는 수학 정확도와 커버리지를 크게 떨어뜨려 통합 훈련의 조합을 재현하지 못했다.

개발자 관점에서 이 논문의 실용적 의미는 훈련 스택을 단계별 알고리즘 묶음이 아니라 가중치가 붙은 토큰 시퀀스를 받는 단일 트레이너로 설계할 수 있다는 점이다. 코퍼스, 데모, 보상으로 걸러진 생성물을 같은 손실 형식으로 표현하고, 단계 전환은 데이터 비율과 시점을 바꾸는 스케줄링 문제로 다룰 수 있다. 다만 구현체가 갖춰야 할 것은 토큰별 손실 가중치, 프롬프트 마스크, 토큰 비율 기준 혼합, 비동기 생성 버퍼와 정책 지연 관리다. 평가도 과제 정확도만 보지 말고 커버리지, 이전 능력 보존, 이후 과제로의 연속 학습 출발점을 함께 봐야 한다. 논문은 특정 하이퍼파라미터나 범용 레시피를 확정하지 않았으므로, 실무에서는 데이터 혼합 실험을 빠르게 돌리는 설계 패턴으로 받아들이는 편이 안전하다.

저자들이 밝힌 한계는 분명하다. 어떤 혼합과 스케줄이 보편적으로 좋은지 찾지 못했고, 정확도·커버리지·보존 사이의 trade-off만 드러냈다. 연산 연구도 예측력 있는 스케일링 법칙을 세운 것이 아니라 평가한 구성의 특성을 기술한 수준이다. 데이터 스케줄과 혼합 설계에 따라 최적화 방법이 어떻게 달라져야 하는지는 깊게 다루지 않았고 향후 과제로 남겼다. 실험 규모도 작아서 대부분 한 번에 하나의 RL 과제만 최적화했으며, 여러 과제를 동시에 훈련할 때의 경쟁 학습 역학과 불균등한 진행 속도는 반영하지 못했다. 또한 새 과제에 대한 다운스트림 적응은 결국 별도 단계가 필요하므로 단계 개념을 완전히 없앴다고 보기는 어렵다. 연속 학습 결과는 이전 훈련 과정이 적응을 도울 수 있다는 초기 증거이며, 더 긴 과제 순서에서 보존과 어떻게 상호작용하는지는 추가 연구가 필요하다.