Papers·1주 전
AC-ODM: 강화학습 기반 프리트레이닝 데이터 혼합 — Pythia-1B에서 MMLU 27.5% 향상
OpenDataLab 팀이 프리트레이닝 데이터 혼합을 강화학습(RL) 관점에서 접근하는 Actor-Critic Online Data Mixing (AC-ODM)을 제안했습니다. 파라미터화된 정책이 그래디언트 간섭을 최대화하는 동적 선형 대리자 역할을 하도록 설계되었으며, 프록시 모드(작은 모델에서 학습한 정책을 큰 모델로 전이)와 비프록시 모드(처음부터 직접 학습)를 모두 지원합니다. Pythia-1B에서 경쟁 기준 대비 최대 66% 적은 학습 스텝으로 최적 검증 perplexity에 도달했고, MMLU 정확도 27.5% 상대 개선, HumanEval pass@1 2.23배 향상을 보였습니다. 추가 연산 오버헤드는 스텝당 0.4%, 메모리 오버헤드는 2%에 불과합니다.
OpenDataLab이 프리트레이닝 데이터 혼합을 강화학습 문제로 재정의한 AC-ODM을 공개했습니다.
핵심 결론
- 태스크 — LLM 프리트레이닝 데이터 혼합 최적화 — RL 기반 동적 혼합 전략.
- 성능 — Pythia-1B에서 MMLU 27.5% 상대 개선, HumanEval pass@1 2.23배, 학습 스텝 최대 66% 감소.
- 오버헤드 — 스텝당 0.4% wall-clock 증가, 2% 메모리 추가로 거의 무시할 수준.
방법
- RL 프레임워크 — 데이터 혼합을 정책 π(θ)로 보고, Actor-Critic으로 그래디언트 간섭을 최대화하는 동적 선형 대리자로 수렴.
- 두 모드 — 프록시 모드: 작은 모델로 정책 학습 후 큰 모델에 전이. 비프록시 모드: 사전 지식 없이 end-to-end 학습.
- 기존 동적 혼합 대비 샘플 효율과 구조적 유연성을 모두 확보한 점이 특징입니다.
한계·조건
- 벤치 규모 — Pythia-1B 위주 실험 — 더 큰 모델(>7B)에서의 일반화는 추가 검증 필요.
- 코드 — GitHub 공개 (https://github.com/DANG-ai/AC-ODM) — 재현 가능.
- 프록시 모드의 전이 효율은 모델 크기 차이에 민감할 수 있다는 단서가 붙습니다.
편집자 한 줄
오버헤드가 거의 없으면서도 큰 폭의 성능 향상을 보여줘, 프리트레이닝 파이프라인에 실용적으로 적용해볼 만한 접근입니다.
- #pretraining
- #data-mixing
- #reinforcement-learning
- #opendatalab
OpenDataLab