For Distillation, Tokens Are Not All You Need....

For Distillation, Tokens Are Not All You Need
Mrigank Raman, Pranav Mani, Davis Liang, Zachary C. Lipton
Статья: https://openreview.net/pdf?id=2fc5GOPYip

Новости дистилляции. Статья с воркшопа Instruction Tuning and Instruction Following на NeurIPS 2023.

TL;DR Авторы предложили метод дистилляции LLM под названием SLIM, использующий значения топ 5% логитов для дистилляции на каждом шаге декодирования + динамическое взвешивание KL/CE лоссов. Результат лучше классической дистилляции, SFT и MiniLLM. Метод скейлится до учителей размером ~70B.

Если поглубже, то современные LLM доросли до огромных размеров и часто используются для аннотирования и генерации инструкций для файнтюнинга малых моделей. Классика файнтюнинга при этом -- supervised fine-tuning (SFT), дообучение на новых (сгенерированных) текстах как на hard labels. Мы знаем, что этот метод несёт сильно меньше информации о распределении, чем мог бы (https://t.me/gonzo_ML/1911).

Дистилляция уже используется для LLM (и многие маленькие LLM получаются из больших, та же Gemini вроде как) и с ней есть сложности. Например, маленькая модель может быть недостаточно выразительной, чтобы покрыть все моды распределения учителя.

Недавний метод MiniLLM (https://arxiv.org/abs/2306.08543) заменяет forward Kullback-Leibler divergence (KLD) на reverse KLD, чтобы предотвратить переоценку студентом низковероятностных регионов распределения учителя. Там для обучения использовался RL, что добавляло сложности. Текущая работа всё упрощает и улучшает, подход назвали SLIM (Sparse Logit Infused Modeling).

Идея проста. Давайте по нашему датасету для обучения через модель-учителя создадим датасет логитов (выходы нейронов без применения функции активации). Для каждого токена в последовательности получаем соответственно V (размер словаря) значений, это будут soft targets. Проблема с таким подходом, что это требует дофига места. Для уменьшения требований предлагается взять только 5% максимальных логитов для каждого токена, остальные считать нулями, так получаются разреженные логиты.

Дальше запускаем процедуру дистилляции (для тех, кто забыл как это работает, можно почитать тут https://t.me/gonzo_ML/118), где в качестве целей берутся soft targets и лосс, являющийся взвешенной суммой обычного кросс-энтропийного лосса (CE) и традиционного KL лосса. Вес KD лоссу даётся зависящий от соотношения логита учителя и студента, таким образом он получается адаптивным и у KL компоненты больше вклад, когда у учителя больше уверенность в предсказании, чем у студента.

Подход проверили на задачах следования инструкциям. Оценивали по Rouge-L и через фидбек от GPT-4. Сравнивались с SFT на hard labels и MiniLLM, для обучения использовали 7B модели LLaMA, LLaMA 2, MPT, учитель был 13-30B. SLIM даёт результат лучше бейзлайнов. SFT вообще хуже всех.

Затем проверили на downstream задачах: ARC, Hellaswag, MMLU, TruthfulQA. Здесь сравнивались с SFT на LLaMA 2 70B в качестве учителя и Llama2-7B в качестве студента. Тоже лучше.

Также проверили генерацию данных для предобучения. Взяли Pythia-6.9B и сгенерили ей датасет текстов + 5% логитов. Затем обучили случайно инициализированную Pythia-160M на подвыборках датасета разного размера. По графику перплексии показали, что SLIM более sample efficient, чем SFT и обычная дистилляция. Я тут не до конца понимаю, что именно они всё-таки под обычной дистилляцией подразумевают, более классический KL лосс чтоли и без топ 5%?

Anyway, просто и эффективно. Метод вроде не сильно отличается от классической дистилляции, не назвал бы его радикально новым, по мне так скорее итеративное улучшение. Интересно, как сами OpenAI, Гугл и остальные внутри себя модели дистиллируют. Есть ли большая разница.