R&B : Regroupement de domaines et équilibrage des mélanges de données pour un entraînement efficace des modèles de fondation

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

May 1, 2025

Auteurs: 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ésumé

Les stratégies de mélange de données ont permis de réduire avec succès les coûts associés à l'entraînement des modèles de langage. Bien que prometteuses, ces méthodes présentent deux défauts majeurs. Premièrement, elles s'appuient sur des domaines de données prédéfinis (par exemple, sources de données, types de tâches), ce qui peut ne pas capturer des nuances sémantiques critiques, laissant ainsi des performances potentielles inexploitées. Deuxièmement, ces méthodes évoluent de manière prohibitive en termes de calcul avec le nombre de domaines. Nous abordons ces défis via R&B, un cadre qui repartitionne les données d'entraînement en fonction de la similarité sémantique (Regroup) pour créer des domaines plus granulaires, et optimise efficacement la composition des données (Balance) en exploitant une matrice de Gram induite par les gradients de domaine obtenus tout au long de l'entraînement. Contrairement aux travaux précédents, il élimine le besoin de calculs supplémentaires pour obtenir des informations d'évaluation telles que les pertes ou les gradients. Nous analysons cette technique sous des conditions de régularité standard et fournissons des insights théoriques qui justifient l'efficacité de R&B par rapport aux approches de mélange non adaptatives. Empiriquement, nous démontrons l'efficacité de R&B sur cinq ensembles de données diversifiés, allant du langage naturel aux tâches de raisonnement et multimodales. Avec seulement 0,01 % de surcharge de calcul supplémentaire, R&B atteint ou dépasse les performances des stratégies de mélange de données les plus avancées.

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.