DoReMi: Оптимизация смесей данных ускоряет предварительное обучение языковых моделей

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

May 17, 2023

Авторы: Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, Adams Wei Yu

cs.AI

Аннотация

Пропорции смешения доменов данных предварительного обучения (например, Википедия, книги, веб-тексты) существенно влияют на производительность языковой модели (LM). В данной статье мы предлагаем метод Domain Reweighting with Minimax Optimization (DoReMi), который сначала обучает небольшую прокси-модель с использованием групповой оптимизации, устойчивой к распределению (Group DRO), по доменам для получения весов доменов (пропорций смешения) без знания о целевых задачах. Затем мы повторно выбираем набор данных с этими весами доменов и обучаем более крупную, полноразмерную модель. В наших экспериментах мы применяем DoReMi к прокси-модели с 280 миллионами параметров для определения весов доменов, чтобы более эффективно обучить модель с 8 миллиардами параметров (в 30 раз больше). На наборе данных The Pile DoReMi улучшает перплексию во всех доменах, даже когда он снижает вес какого-либо домена. DoReMi повышает среднюю точность на нескольких целевых задачах на 6,5% по сравнению с базовой моделью, обученной с использованием стандартных весов доменов The Pile, и достигает точности базовой модели за в 2,6 раза меньше шагов обучения. На наборе данных GLaM DoReMi, который не имеет информации о целевых задачах, даже соответствует производительности использования весов доменов, настроенных на целевых задачах.

English

The mixture proportions of pretraining data domains (e.g., Wikipedia, books, web text) greatly affect language model (LM) performance. In this paper, we propose Domain Reweighting with Minimax Optimization (DoReMi), which first trains a small proxy model using group distributionally robust optimization (Group DRO) over domains to produce domain weights (mixture proportions) without knowledge of downstream tasks. We then resample a dataset with these domain weights and train a larger, full-sized model. In our experiments, we use DoReMi on a 280M-parameter proxy model to find domain weights for training an 8B-parameter model (30x larger) more efficiently. On The Pile, DoReMi improves perplexity across all domains, even when it downweights a domain. DoReMi improves average few-shot downstream accuracy by 6.5% over a baseline model trained using The Pile's default domain weights and reaches the baseline accuracy with 2.6x fewer training steps. On the GLaM dataset, DoReMi, which has no knowledge of downstream tasks, even matches the performance of using domain weights tuned on downstream tasks.