사전학습과 강화학습을 같은 손실로 잇는 통합 학습

배경과 문제

파운데이션 모델은 보통 사전학습, 중간학습, 지도식 미세조정, 강화학습의 순서로 학습된다. 각 단계는 데이터와 학습 알고리즘이 다르고, 전문적인 학습 레시피와 인프라에 묶여 별개의 학습 문제로 취급되어 왔다. 예를 들어 사전학습과 중간학습은 빈 프롬프트에 대해 일반 말뭉치나 선별된 말뭉치의 텍스트를 맞추고, 지도식 미세조정은 지시문과 데모 응답 쌍을 맞춘다. 반면 강화학습은 보상 피드백을 이용해 모델이 생성한 완성문을 최적화하므로 언뜻 최대우도 추정과 다른 문제처럼 보인다.

이런 단계 분리는 과학적 질문도 쪼갠다. 사전학습 스케일링 법칙은 데이터·계산·모델 성능의 관계를, 강화학습 스케일링 연구는 사후학습의 효율을 따로 연구한다. 그러나 이전 단계의 학습 결정이 이후 학습에 영향을 준다는 연구가 쌓이면서, 전체 학습 과정에서 내린 결정이 최종 모델 능력을 어떻게 함께 결정하는지에 대한 이해는 부족하다. 이 논문은 단계별 지표 대신 최종 모델에 미치는 효과로 결정을 판단할 수 있는 통합된 관점을 제안한다.

기존 접근의 한계

기존 파이프라인은 단계마다 별도의 학습 레시피를 설계해야 한다. 사전학습, 중간학습, 지도식 미세조정은 모두 최대우도 추정을 쓰지만 데이터만 다르다. 강화학습은 정책 최적화로 정식화되어 이 패턴을 깨는 것처럼 보인다. 예를 들어 REINFORCE는 샘플링된 완성문에서 기대 보상의 기울기를 추정하고, PPO나 GRPO는 어드밴티지 추정, 보상 정규화, 정책 업데이트 제약 등을 추가한다.

데이터 투입 시점을 단계별로 정하면 중간학습이 최종 강화학습을 얼마나 돕는지 따로 살피기 쉽다. 논문은 중간학습 데이터 양이나 시점을 최종 성능과 함께 평가하자고 제안한다. 통제 실험에서는 단계적 학습의 중간학습 노출을 최적점 이상으로 늘렸을 때 최종 과제 정확도가 떨어졌다. 같은 자료라도 언제 얼마나 재사용하는지가 중요하다는 관찰이다.

연구방법을 단계별로

첫 단계는 MaxRL의 우도 해석을 확장하는 것이다. 이진 보상에서 프롬프트 x의 성공 확률 pθ(x)를 정의하고, 성공 조건부 분포 πθ+(y|x)를 생각하면 MaxRL 기울기는 성공 조건부 분포에서 로그 우도의 기울기 기댓값과 같다. 따라서 반복 t에서 MaxRL 업데이트는 현재 정책의 성공 조건부 분포 πθ+t를 고정한 채 최대우도 적합을 하는 것과 동일하다. 이로써 강화학습도 사전학습·중간학습·지도식 미세조정과 같은 우도 적합 형태로 놓인다.

다음으로 각 데이터 소스를 프롬프트 분포 ρs,t와 목표 분포 µs,t로 표현하고, 혼합 가중치 λs,t를 곱해 단일 결합 분포 qt를 만든다. 통합 학습은 이 혼합 분포에 동일한 우도 손실 L(θ; qt) = −E(x,y)∼qt[log πθ(y|x)]를 적용한다. 소스 분포와 혼합 가중치는 미분할 때 고정된다. 기존 단계적 학습은 모든 혼합 가중치를 한 번에 한 소스에만 두는 특수 사례다.

이 정식화는 다른 생성 정책 ν가 만든 성공 응답 분포 ν+도 같은 최대우도로 맞출 수 있다. 다만 ν가 현재 정책과 다르면 일반적으로 현재 정책의 MaxRL 기울기와는 다르다. 이진 보상의 REINFORCE는 같은 성공 조건부 목표를 쓰되 성공 확률이 높은 프롬프트에 더 큰 가중치를 준다. 연속 보상의 TailRL도 보상 임계값별 성공 조건부 분포를 섞는 목표로 표현된다. 이런 확장이 모든 강화학습 알고리즘의 동일성을 뜻하지는 않는다.

실제 구현에서는 각 학습 예를 토큰 시퀀스 z와 토큰별 가중치 벡터 w의 쌍으로 표현한다. 사전학습과 중간학습은 말뭉치 토큰에 단위 가중치를, 지도식 미세조정은 데모 토큰에 단위 가중치를 주고 프롬프트 토큰은 0으로 마스킹한다. 강화학습은 같은 프롬프트 마스크를 쓰고 각 완성문에 보상과 샘플 통과율로 계산한 응답 가중치 A를 부여한다. 손실은 배치 내 비프롬프트 토큰 총수로 정규화한 토큰 평균 음의 로그 우도다.

데이터 생성에서는 프롬프트당 N개의 완성문을 최근 모델 복사본에서 샘플링하고, 이진 보상의 통과율 p̂를 추정한다. 정답 완성문은 응답 토큰에 가중치 1/p̂를 받고 오답은 0을 받는다. 이 가중치는 MaxRL의 유한 샘플 형태로, N개 샘플에 대한 가중 스코어 함수 평균은 정답 완성문들의 평균 스코어가 된다. 모든 완성문이 오답이면 업데이트가 0이 된다. 생성과 최적화는 비동기로 진행되어 생성기가 최근 가중치로 데이터를 만들어 공유 버퍼에 추가하고, 트레이너가 버퍼에서 샘플링해 학습한 뒤 주기적으로 가중치를 공개한다.

결과와 숫자의 의미

통제된 추론 퍼즐 실험은 124M 파라미터 모델과 공통 20B 토큰 사전학습 자료, 약 510M 토큰의 고정 중간학습 자료를 사용했다. 통합 학습은 중간학습 자료를 사전학습 후반과 강화학습 중에 섞는다. Base는 8~9인 퍼즐, Extended는 더 큰 10~11인 퍼즐까지 학습하며, 일반화 평가는 학습하지 않은 12인 퍼즐이다. 각 시험집합은 1024문항이다. 문항당 64응답에서 pass@1과 pass@16을 추정하고 두 강화학습 시드의 결과를 평균했다. 최고의 통합 구성은 두 설정의 최종 정확도와 크기 일반화에서 이점을 보였다.

그림 5는 사전·중간·강화학습과 응답 생성을 포함한 추정 FLOP별 관측 최선의 체크포인트를 연결한다. 단계적 학습이 일부 중간 예산에서 정확도가 높았고, 통합 학습은 큰 예산과 일반화 평가에서 유리했다. 학습 토큰은 대략 6NT, 생성 토큰은 2NT로 계산하며 N은 모델 파라미터 수다. 평가·패딩·활성값 재계산·미사용 생성 등의 비용은 제외했다. 이는 시험한 설정의 관측 경계이며 모든 학습 비용을 잰 결과나 예측적 스케일링 법칙은 아니다.

강화학습 전 일반 능력 보존도 비교했다. 다섯 통합 초기화는 사전학습만 한 모델과 비슷한 보류 데이터 손실과 일반 벤치마크 성능을 유지했다. 반면 유지 성능이 가장 좋은 단계적 구성도 보류 손실이 높고 벤치마크 평균이 낮았다. 최종 강화학습 정확도가 가장 높은 단계적 구성의 유지 성능은 더 나빴다. 이 결과는 124M 통제 조건에서 자료를 사전학습 중에 섞는 것이 망각을 줄일 수 있음을 시사한다.

0.76B 실험은 이미 354.53B 토큰 사전학습과 17.73B 토큰 중간학습을 마친 모델에 수학 강화학습을 적용했다. 추가 240M 토큰 말뭉치를 강화학습 중 섞거나 강화학습 전에 사용했다. 표 16의 5% 혼합은 단계적 학습 대비 SciQ 65.10% 대 53.00%, 보류 손실 변화 −0.1474 대 +0.0472 nats/token이었다. 하지만 HumanEval+ pass@1은 10.90% 대 15.51%, MBPP+는 20.97% 대 25.74%로 낮았다. 표 15의 수학 정확도도 일괄 상승하지 않았다. Math와 Coding은 각각 두 벤치마크를 동일 가중 평균한 값이다.

수학 학습 뒤 모든 모델에 같은 코딩 강화학습을 4000회 추가하자 결과가 달라졌다. 표 22에서 5% 혼합 출발 모델의 HumanEval+ pass@1은 24.96%, MBPP+는 32.69%로 단계적 모델의 22.94%, 29.81%보다 높았다. 모든 통합 모델의 최종 두 코딩 정확도와 MBPP+ 커버리지가 높았지만, HumanEval+ pass@16은 일부 혼합률에서 단계적 모델보다 낮았다. 5% 혼합 모델은 네 지표 모두에서 가장 높았다. 이는 수학 학습 직후 코딩 성능과 후속 적응 능력을 구분해야 함을 보여준다.

구체적인 예시

통제 실험의 추론 퍼즐은 사람들을 서로 다른 층에 배정하는 문제다. 프롬프트에는 점유된 층 목록, 각 사람의 가장 친한 친구, 자연어 단서들이 주어진다. 단서는 차이, 합, 비교, 홀짝, 정수 배수로 층 번호를 제한한다. 중간학습 말뭉치는 기호 오라클이 만든 풀이 추적으로 구성되며, 'Ava의 가장 친한 친구'나 '가장 높은 층에 사는 사람' 같은 참조 패턴을 시연한다.

강화학습은 중간학습에 없던 참조 조합을 도입한다. 8인 평가 예에서 정답 배정을 보면 가장 높은 층의 거주자는 Nora이고, Nora의 가장 친한 친구는 Mia다. 따라서 ‘가장 높은 층 거주자의 친구’는 Mia이며, 거기서 친구 관계를 한 번 더 따라가야 Felix가 된다. 원문의 그림 10은 정답을 사용해 두 합성을 설명하지만 이 배정은 모델 입력에 주어지지 않는다. 구성 연산을 아는 모델이 새로운 조합까지 배울 수 있는지 살피는 설계다.

실제 데이터 실험의 혼합 비율은 사전 추정한 목표 토큰 분율이다. 예를 들어 0.5% 혼합은 중간학습 토큰 수를 중간학습 토큰과 강화학습 응답 토큰의 합으로 나눈 값이 0.5%가 되도록 설정한다. 실제 응답 길이가 달라지면 실현 분율은 달라질 수 있다. 일정 실험의 '절반 소멸 2%'는 초기 2%에서 시작해 학습 중간에 0이 되고 이후 유지되는 일정이다.

한계와 오해하기 쉬운 점

이 논문은 적합한 강화학습 목적의 우도 표현을 제안했지만, PPO나 GRPO를 포함한 모든 강화학습 알고리즘이 완전히 동일하다고 주장하지 않는다. 경험적 어드밴티지, 정책 비율, 클리핑, KL 페널티는 고정 목표 우도 정식화를 보존하지 않을 수 있다. 또한 유한 샘플에서는 정답 응답의 가중치가 1/추정 성공률이고 오답은 0이므로, 정답이 한 번도 나오지 않은 프롬프트는 업데이트되지 않으며 성공률이 낮은 프롬프트가 덜 반영된다.

실험은 124M과 0.76B의 소형 모델에 한정되며 주로 한 번에 하나의 강화학습 과제를 최적화했다. 계산 제약으로 0.76B에서는 사전학습 중 통합 학습을 시험하지 않았다. 여러 과제의 장기 연속 학습이나 큰 모델로의 확장은 아직 검증되지 않았다. 새 과제에 적응하는 별도 단계도 남아 있고, 최적의 혼합 비율·일정이나 예측적 스케일링 법칙도 확립하지 않았다.

헬드아웃 손실은 24개 구성요소를 동일 가중 평균한 3,608,122 scored tokens의 평가이며, 실제 훈련 코퍼스 전체 토큰 평균이 아니다. 기본 모델 대비 변화는 base 1.4352 nats/token 기준이다. 0.76B 모델의 pass@1은 문항당 16응답으로 추정되고, 수학 커버리지는 별도 256응답에서 pass@32/64/256, 코딩 커버리지는 32응답에서 pass@32로 추정된다. 후속 코딩의 pass@16은 16응답이며 같은 지표로 혼동하지 말아야 한다.

0.76B 비교는 12000회 업데이트를 맞췄지만 총 계산량은 맞추지 않았다. 통합 구현은 비동기 생성과 보상 기준값 차감 없는 가중치, 단계적 구현은 현재 정책 생성과 기준값 차감을 사용한다. 같은 통합 인프라의 0% 혼합 대조군도 중요하다. 보고된 7.8배는 0% 혼합의 3.7초 대 단계적 학습의 28.8초라는 업데이트 처리량이다. 각각 GPU 24대와 16대를 썼고 통합 구현은 응답을 재사용할 수 있어, 같은 장비의 비용 절감이나 새 응답 생성 속도로 해석할 수 없다.

가중치는 네 업데이트마다 생성기에 전달하지만, 소비되는 데이터가 네 단계 이내라는 보장은 없다. 일부 생성기 버퍼 스냅샷의 정책은 12~44업데이트 뒤처졌으며 실제 학습에 소비된 예제의 정책 버전 분포는 기록하지 않았다. 비동기화 효과와 데이터 혼합 효과를 완전히 분리했다고 보기 어려운 이유다.

실제로 어떻게 활용할까

공유 버퍼와 가중 토큰 손실은 데이터 혼합 비율과 일정을 같은 트레이너 안에서 비교할 수 있게 한다. 다만 이를 실제 시스템에 적용할 때는 업데이트 처리량뿐 아니라 GPU 수, 새 응답 생성량, 데이터 재사용과 지연, 최종 과제 성능을 함께 측정해야 한다. 이 구분은 논문의 서로 다른 구현 조건과 측정 범위에서 도출되는 활용 관점이다.

데이터 혼합 비율과 일정은 정확도, 커버리지, 유지 성능 사이의 트레이드오프를 조절한다. 예를 들어 중간학습 데이터를 강화학습 내내 일정하게 섞으면 유지 성능이 좋아지지만 코딩 정확도가 낮아질 수 있고, 학습 중간에 혼합을 끝내면 코딩 정확도가 개선되지만 말뭉치 손실 감소는 줄어든다. 따라서 실제 적용에서는 목표 과제와 보존할 능력의 우선순위에 따라 일정을 선택해야 한다.

지속 학습에서는 현재 과제 직후의 점수만으로 다음 과제의 적응력을 판단하기 어렵다. 이 연구에서 수학 강화학습 중 정적 자료를 섞은 모델은 코딩 단일 응답 정확도가 낮아졌지만, 동일한 후속 코딩 학습 뒤에는 더 높은 정확도에 도달했다. 이는 이미 사전·중간학습을 마친 0.76B 모델의 관찰이며, 사전학습부터 모든 단계를 통합하면 항상 같은 이득이 난다는 증거는 아니다.