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 — фреймворка, который перераспределяет обучающие данные на основе семантического сходства (Regroup) для создания более детализированных доменов и эффективно оптимизирует состав данных (Balance), используя матрицу Грама, индуцированную градиентами доменов, полученными в процессе обучения. В отличие от предыдущих работ, он устраняет необходимость в дополнительных вычислениях для получения оценочной информации, такой как потери или градиенты. Мы анализируем эту технику в рамках стандартных условий регулярности и предоставляем теоретические обоснования, подтверждающие эффективность 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.