RegMix-D: 프록시 학습 궤적(Proxy Training Trajectories)을 통한 동적 데이터 믹싱

데이터 믹스(Data mixture) 선택은 거대 언어 모델(Large Language Model, LLM) 사전 학습(pretraining)에 있어 매우 중요합니다. RegMix와 같은 기존 방법들은 소규모 프록시 실행(proxy runs)에 대해 회귀 모델(regression model)을 적합(fitting)시켜 단일 정적 믹스(single static mixture)를 선택합니다. 우리는 RegMix를 동적 믹싱(dynamic mixing)으로 확장한 간단한 방식인 RegMix-D를 제안합니다. 우리의 핵심 관찰 결과는 프록시 실행이 최종 손실(endpoint losses)뿐만 아니라 전체 손실 궤적(full loss trajectories)을 생성하며, 이를 데이터 믹스를 더욱 개선하는 데 사용할 수 있다는 점입니다. 이러한 궤적을 바탕으로 회귀 모델을 학습함으로써, 우리는 여러 학습 단계에서의 최적의 믹스를 예측할 수 있습니다. RegMix-D는 두 가지 배포 모드를 지원합니다: 타겟 학습(target training) 전에 전체 믹스 스케줄(mixture schedule)을 생성하는 오프라인 변형(offline variant)과, 관찰된 손실(observed loss)을 사용하여 학습 중에 믹스를 조정하는 온라인 변형(online variant)입니다. 1B 파라미터 타겟 모델로 Pile 데이터셋의 25B 토큰에 대해 실험한 결과, RegMix-D는 프록시 효율성(proxy-efficient)을 유지하면서도 13개의 다운스트림 태스크(downstream tasks) 전반에서 RegMix 및 DoReMi보다 일관되게 우수한 성능을 보였습니다. 특히 RegMix-D는 단 128개의 프록시 모델(RegMix 프록시 연산 예산의 25%)만으로도 RegMix를 능가합니다.