Улучшение качества 📈 за счет оптимального выбора весов для обучающих датасетов 📚

Датасеты бывают из разных доменов: интернет новости, википедия, код, диалоги, инструкции... И применяться одна и та же модель может на разных задачах: например это могут быть chit-chat диалоги, диалоги с детьми, диалоги, склоняющие к целевому действию - в зависимости от пользователя. В каких пропорциях брать датасеты в финальный трейн сет?

Авторы статьи (из Гугла и Стенфорда) предлагают способ выбора весов для каждого домена, который в их экспериментах:

🔹 Уменьшил лосс на каждом домене относительно бейзлайна (эвристически подобранные веса датасета Pile) - без трейдоффа, просто улучшение (рисунок 1)

🔹 Обученная их методом модель улучшила среднюю точность One-Shot’а в 5 разных задачах (рисунок 2)


Метод (рисунок 3):

1️⃣ Обучить небольшую reference Langauage Model

2️⃣ Обучить прокси модель (используя Group Distributionally Robust Optimization), чтобы получить новые веса для доменов

3️⃣ Обучить финальную большую LM с этими весами


❗️Group DRO: минимизация наибольшего лосса по всем доменам. То есть ищем модель, которая работает как можно лучше в самых тяжелых случаях.

На самом деле Group DRO не просто лосс минимизирует, а разницу лосса прокси и референс моделей (но получилось бы слишком тяжелое для понимания предложение 😁). Сам процесс обучения: поочередное изменение весов доменов и весов прокси LM’ки для решения минимакс задачи (алгоритм на рисунке 4)

В итоге:
🔸 качество строго лучше и достигается оно быстрее - относительно эвристически подобранных весов
🔸 качество такое же, но эксперименты дешевле - относительно оптимизации весов доменов под целевые задачи сразу с большой LM’кой

Статья: DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining