R&B: Regrupamiento de Dominios y Equilibrio de Mezcla de Datos para el Entrenamiento Eficiente de Modelos Fundacionales

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

May 1, 2025

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

Resumen

Las estrategias de mezcla de datos han logrado reducir los costos involucrados en el entrenamiento de modelos de lenguaje. Aunque prometedoras, estos métodos presentan dos deficiencias. En primer lugar, dependen de dominios de datos predeterminados (por ejemplo, fuentes de datos, tipos de tareas), lo que puede no capturar matices semánticos críticos, dejando rendimiento sobre la mesa. En segundo lugar, estos métodos escalan con el número de dominios de una manera computacionalmente prohibitiva. Abordamos estos desafíos mediante R&B, un marco que redivide los datos de entrenamiento basándose en la similitud semántica (Regroup) para crear dominios más granulares, y optimiza eficientemente la composición de los datos (Balance) aprovechando una matriz Gram inducida por los gradientes de dominio obtenidos durante el entrenamiento. A diferencia de trabajos anteriores, elimina la necesidad de computación adicional para obtener información de evaluación, como pérdidas o gradientes. Analizamos esta técnica bajo condiciones de regularidad estándar y proporcionamos insights teóricos que justifican la efectividad de R&B en comparación con enfoques de mezcla no adaptativos. Empíricamente, demostramos la efectividad de R&B en cinco conjuntos de datos diversos que abarcan desde lenguaje natural hasta tareas de razonamiento y multimodales. Con tan solo un 0.01% de sobrecarga computacional adicional, R&B iguala o supera el rendimiento de las estrategias de mezcla de datos más avanzadas.

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.