Introduction
대규모 Language Model(LLM)을 훈련할 때, 일반적으로 전체 토큰에 대해 일률적으로 next-token prediction(Causal Language Modeling, CLM) 방식으로 학습해왔다. 그러나 이 논문은 이러한 접근에 의문을 제기하며, 훈련 코퍼스에 있는 모든 토큰이 학습에 똑같이 중요한 것은 아니다라는 문제 제기를 한다. 즉, 노이즈가 많거나 의미가 불분명한 토큰까지 모두 학습에 포함하는 것이 최적이 아닐 수 있다는 것이다.
이 논문에서는 그 문제 상황을 동기 삼아, Selective Language Modeling(SLM)이라는 새로운 접근을 제안한 RHO-1 모델을 소개한다. RHO-1은 모든 토큰을 대상으로 하기보다는, 중요하다고 판단되는 토큰에 대해서만 언어 모델의 로스를 계산해 학습한다. 이를 통해, pretraining 데이터 중에서도 노이즈가 심하거나 학습 효율이 낮은 토큰을 자동으로 거르며, 결과적으로 데이터 효율성과 성능을 동시에 높이는 데 성공했다.
문제 상황: Not all tokens are equally important
일반적으로 LLM 훈련에는 웹 크롤링 등으로 대규모 텍스트 데이터를 수집하는데, 아무리 정제(document-level filtering)를 잘해도, 토큰 단위로 보면 여전히 학습에 부적합한 부분(깨어진 문법, 무의미한 토큰 등)이 다수 존재한다.
- 모든 토큰에 대해 똑같은 로스를 부여하는 Causal LM 방식은, 이런 불필요한 토큰까지 학습 대상으로 삼아, computation 낭비와 성능 저하를 야기한다.
- 게다가 웹 텍스트 토큰은 LLM에게 꼭 필요한 분포가 아닐 수도 있으며, 이미 잘 학습된 토큰이나 끝내 못 배울 무작위 토큰까지 동일한 가중으로 학습한다는 비효율이 발생한다.
Training Dynamics 분석
저자들은 우선 LLM을 실제로 학습시키면서, 각 토큰별 loss가 어떻게 변해가는지 관찰했다. 그 결과, 토큰을 크게 네 부류로 나눌 수 있었다:
- H→L: 처음엔 loss가 높았지만(H) 최종적으로 낮아지는(L) – 학습할 가치가 있는 토큰
- L→L: 처음부터 끝까지 loss가 낮은(L) – 이미 잘 학습돼 있는 토큰
- H→H: 처음부터 끝까지 loss가 높은(H) – 끝내 수렴 안 하는 토큰
- L→H: 낮았다가(O) 오히려 높아지는(H) – 이상하게 언젠가 어렵게 돼버린 토큰
실험적으로, 유의미하게 loss가 떨어지는(H→L) 토큰은 소수(약 26%)에 불과하며, 많은 토큰(51%)은 L→L 이미 습득되었거나 쉽고, 약 11% 토큰(H→H)은 계속 어렵기만 하고 수렴되지 않았다. 또한 어떤 토큰들은 학습 과정 중에도 loss가 들쭉날쭉하며 안정적으로 배우거나, 버려지지 못하는 것으로 나타났다.
이런 분석은 모든 토큰을 똑같이 학습할 필요가 없다는 직관을 강하게 뒷받침한다.
Selective Language Modeling (SLM)
위 문제의 해결책으로 저자들은 SLM 기법을 제시한다. 요약하자면 다음과 같다:
- Reference Model 준비
- 우선 소량의 ‘고품질’ 텍스트 데이터(수억~수십억 단위)를 모아 Reference LM(RM)을 학습한다.
- 이 Reference Model은 말 그대로 “토큰별로 얼마나 RM이 잘 맞추는가”를 측정하기 위한 기준.
- Excess Loss 계산
- 본격적인 대규모 pretraining 코퍼스를 한 번 pass하며, Reference Model이 각 토큰에 대해 얼마나 큰 loss를 가지는지($L_{RM(x_i)}$) 계산한다.
- 동시에 현재 학습 중인 모델(θ)의 loss L_θ(xi)도 관찰해, excess loss = $[L_{\theta(x_i)} - L_{RM(x_i)}]$를 구한다.
- 곧, Reference 대비 현재 모델이 특정 토큰을 얼마나 더 학습 필요로 하는지를 수치화한 것.
- Top k% 토큰 선별
- excess loss가 높은 토큰들은 아직 제대로 못 배우고 있고, 동시에 Reference에 비해 학습하면 이득이 될 만한 토큰이라고 볼 수 있다.
- 반면 excess loss가 낮은 토큰들은 이미 충분히 학습됐거나(no improvement), 혹은 Reference와 괴리가 없는 토큰이므로 학습 효율이 낮다.
- 따라서 Top k% (예: 60~70%) 토큰에 대해서만 로스를 계산하고 backpropagation을 적용한다. 그 외 토큰은 loss 계산에서 제외한다.
구현의 장점
- 추가 비용이 거의 없다: 한 미니배치 안에서 Reference LM의 log prob만 미리 구해두면, 로스 계산 시 토큰마다 excess loss 순위를 매길 수 있다. 이는 GPU 상에서도 크게 오버헤드를 일으키지 않는다.
- 정제율(k%)은 하이퍼파라미터로 실험적 튜닝이 가능하다.
- 토큰 단위 노이즈 제거: 문서 레벨 필터링으로 제거하기 애매했던 잡음(HTML 태그 파편, 이상한 토큰 등)을 학습 과정에서 동적으로 제외할 수 있다.
실험 결과
1. 수학 도메인(Math) Continual Pretraining
수학 텍스트(OpenWebMath) 약 14~15B tokens를 활용해, Tinyllama-1B 및 Mistral-7B 모델을 추가 학습(continual pretraining)하는 실험을 했다.
- Baseline: 일반 causal language modeling으로 단순히 이어서 학습(CT).
- RHO-1: 동일 코퍼스를 동일 step만큼 학습하되, SLM으로 상위 60~70% 토큰만 남김.
(A) Few-shot CoT 성능
- GSM8K, MATH, SVAMP 등 여러 수학 태스크에서, RHO-1(1B)은 baseline 대비 평균 16%p, (7B)는 10%p 이상의 정확도 향상을 보였다.
- 또, SLM은 baseline 수준에 도달하는 데 5~10배 더 적은 pretraining token만 사용해도 되었다(학습 효율 극대화).
- RHO-1(7B)는 math토큰 단 15B로, 500B를 쓴 DeepSeekMath(7B) 수준 성능에 도달.
(B) Tool-Integrated Fine-Tuning
- 위 모델들을 math reasoning + external tools를 사용하는 포맷으로 fine-tune(ToRA)한 결과, MATH dataset 정확도가 7B에서 51.8%로 향상, 이는 동급 모델 대비 최고 수준. 1B 모델도 40.6% 달성해, 같은 크기 모델 중 SOTA.
2. 일반 도메인(General) Continual Pretraining
Tinyllama-1B 모델을 일반 텍스트 80B tokens(코드 + 텍스트 + 소량 math)로 추가 학습 시, RHO-1이 baseline 대비 15개 태스크 평균 6.8% 향상, 특히 코드/수학 문제에서 +10% 이상 점프했다.
3. Self-Reference
Reference Model을 만들 때 반드시 별도 고품질 데이터가 필요한 것은 아니며, 동일 코퍼스 내부에서 self-supervision(“self-reference”)을 통해서도 유의미한 향상을 달성했다.
- 이는 정말 중요/유용한 토큰과 노이즈 토큰을 구분해줄 수 있는 Reference Model 만 있으면 되기 때문.
추가 분석
- 선택된 토큰(Selected Tokens) vs 미선택 토큰(Unselected Tokens)의 loss
- 실험에서, Selected Tokens의 loss를 줄이면 downstream(수학 태스크)까지 성능이 크게 향상됨을 시각적으로 확인(선형적 혹은 파워 법칙적 관계).
- 반면 Unselected Tokens에 대한 학습은 downstream에 별 효과가 없다.
- SLM이 실제로 어떤 토큰을 선택하는가?
- 예시를 보면, 올바른 수학 공식, 문맥상 유의미한 수학 텍스트 등은 많이 선택되고, 의미 없는 HTML/로그 태그, 형식만 맞지 않는 심볼 등은 잘 배제된다.
- Token Selection Ratio
- 너무 적게 선택하면 데이터가 부족해지고, 너무 많이 선택하면 필터링 효과가 떨어진다. 대략 60~70%가 적절.
결론 및 요약
RHO-1는 Not all tokens are equally important라는 문제 인식에서 출발해, Selective Language Modeling (SLM)을 제안했다. 핵심 아이디어는 (1)참조 모델을 만들어 토큰별 loss를 측정하고, (2)현재 모델과의 차이(excess loss)가 큰 토큰만을 선별해 학습하는 것이다. 이로써 토큰 레벨 노이즈 제거와 데이터 효율 개선을 동시에 달성했으며, 수학·일반 도메인에서 큰 성능 향상과 학습 비용 절감을 증명했다.
향후에는 SLM을 훨씬 큰 스케일의 LLM 혹은 multimodal 데이터로 확장하거나, token re-weighting, token-level RL, multi-reference 등 다양한 확장이 가능할 것으로 기대된다.
SLM이 단순히 안 맞는 토큰을 버리는 수준이 아니라, LLM 훈련 과정에서 학습 가치가 높은 토큰에 집중하게 만드는 하나의 실행 가능한 전략이라는 점에서, 차세대 LLM pretraining에 중요한 통찰을 제공한다고 볼 수 있다.
