R&B: 효율적인 파운데이션 모델 학습을 위한 도메인 재그룹화 및 데이터 혼합 균형화

R&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training

May 1, 2025

저자: Albert Ge, Tzu-Heng Huang, John Cooper, Avi Trost, Ziyi Chu, Satya Sai Srinath Namburi GNVV, Ziyang Cai, Kendall Park, Nicholas Roberts, Frederic Sala

cs.AI

초록

데이터 혼합 전략은 언어 모델 훈련에 소요되는 비용을 성공적으로 절감해 왔다. 이러한 방법은 유망하지만 두 가지 결점을 가지고 있다. 첫째, 이들은 미리 정의된 데이터 도메인(예: 데이터 소스, 작업 유형)에 의존하는데, 이는 중요한 의미론적 뉘앙스를 포착하지 못해 성능을 제한할 수 있다. 둘째, 이러한 방법은 도메인 수에 따라 계산적으로 비효율적인 방식으로 확장된다. 우리는 이러한 문제를 R&B 프레임워크를 통해 해결한다. R&B는 의미론적 유사성을 기반으로 훈련 데이터를 재분할(Regroup)하여 더 세분화된 도메인을 생성하고, 훈련 과정에서 얻은 도메인 그래디언트로 유도된 그램 행렬을 활용하여 데이터 구성을 효율적으로 최적화(Balance)한다. 기존 연구와 달리, R&B는 손실이나 그래디언트와 같은 평가 정보를 얻기 위한 추가 계산을 필요로 하지 않는다. 우리는 이 기법을 표준 규칙성 조건 하에서 분석하고, R&B의 효과를 비적응적 혼합 접근법과 비교하여 이론적으로 입증한다. 실험적으로, 우리는 자연어 처리부터 추론 및 다중모달 작업에 이르는 다섯 가지 다양한 데이터셋에서 R&B의 효과를 입증한다. 단 0.01%의 추가 계산 오버헤드만으로도 R&B는 최신 데이터 혼합 전략의 성능을 능가하거나 동등한 수준을 달성한다.

English

Data mixing strategies have successfully reduced the costs involved in training language models. While promising, such methods suffer from two flaws. First, they rely on predetermined data domains (e.g., data sources, task types), which may fail to capture critical semantic nuances, leaving performance on the table. Second, these methods scale with the number of domains in a computationally prohibitive way. We address these challenges via R&B, a framework that re-partitions training data based on semantic similarity (Regroup) to create finer-grained domains, and efficiently optimizes the data composition (Balance) by leveraging a Gram matrix induced by domain gradients obtained throughout training. Unlike prior works, it removes the need for additional compute to obtain evaluation information such as losses or gradients. We analyze this technique under standard regularity conditions and provide theoretical insights that justify R&B's effectiveness compared to non-adaptive mixing approaches. Empirically, we demonstrate the effectiveness of R&B on five diverse datasets ranging from natural language to reasoning and multimodal tasks. With as little as 0.01% additional compute overhead, R&B matches or exceeds the performance of state-of-the-art data mixing strategies.