R&B: Domänen-Neugruppierung und Datenmischungsausgleich für effizientes Training von Foundation-Modellen

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

May 1, 2025

Autoren: 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

Zusammenfassung

Datenmischstrategien haben die Kosten für das Training von Sprachmodellen erfolgreich reduziert. Obwohl vielversprechend, weisen solche Methoden zwei Schwächen auf. Erstens basieren sie auf vordefinierten Datenbereichen (z. B. Datenquellen, Aufgabentypen), die möglicherweise kritische semantische Nuancen nicht erfassen und somit Leistungspotenzial ungenutzt lassen. Zweitens skalieren diese Methoden mit der Anzahl der Bereiche in einer rechnerisch unvertretbaren Weise. Wir adressieren diese Herausforderungen durch R&B, ein Framework, das Trainingsdaten basierend auf semantischer Ähnlichkeit neu partitioniert (Regroup), um feiner abgestufte Bereiche zu schaffen, und die Datenzusammensetzung effizient optimiert (Balance), indem eine Gram-Matrix genutzt wird, die durch Domänengradienten während des Trainings induziert wird. Im Gegensatz zu früheren Arbeiten entfällt die Notwendigkeit zusätzlicher Rechenleistung, um Bewertungsinformationen wie Verluste oder Gradienten zu erhalten. Wir analysieren diese Technik unter Standard-Regularitätsbedingungen und liefern theoretische Einblicke, die die Effektivität von R&B im Vergleich zu nicht-adaptiven Mischansätzen rechtfertigen. Empirisch demonstrieren wir die Wirksamkeit von R&B auf fünf verschiedenen Datensätzen, die von natürlicher Sprache über logisches Denken bis hin zu multimodalen Aufgaben reichen. Mit nur 0,01 % zusätzlichem Rechenaufwand erreicht oder übertrifft R&B die Leistung modernster Datenmischstrategien.

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.