⚗️ Что такое дистилляция и как она применяется в LLM — часть I

Чем больше модель, тем сложнее ее инферить и дороже обучать. Решить проблему призвана, в том числе, дистилляция — передача знаний от тяжёлой модели («учителя») более лёгкой («ученика»). Расскажем, какие типы дистилляции существуют и как их используют.

Классический способ предложил Джеффри Хинтон в статье 2015 года. Учёный выдвигает гипотезу, что распределение классов, которые модель предлагает в качестве ответа, само по себе содержит немало знаний. Поэтому имеет смысл тренировать «ученика» не на ответах, а на распределении классов «учителя», используя Softmax с температурой. В качестве лосса использовали кросс-энтропию между двумя распределениями — ответами учителя и ученика.

Одна из первых моделей, которую дистиллировали на претрейне, — DistilBERT. Результат получился впечатляющим: language understanding удалось сохранить на 97%, а скорость по заявлению авторов выросла на 60%. Интересно, что дистиллировали веса, а в архитектуре модели изначально было вдвое меньше энкодер-блоков, чем у базовой BERT — 6 против 12. В основе обучения — перекрестная энтропия ответов «учителя» и «ученика», MLM и L cos — косинусная близость между эмбеддингами на скрытых слоях. Идеи DistilBERT позднее применяли, например, в DistilGPT.

Самый простой из современных методов — имитация модели. Его суть — добиться, чтобы небольшая модель копировала поведение крупной. Для этого «учителя» просят генерировать ответы на разные запросы, а потом на них обучают «ученика».

Маленькие модели отлично подражают большим, но не развивают собственные навыки. Поэтому «ученики» не получают новые знания, зато неплохо справляются с тем, чтобы извлекать имеющиеся. Этот метод подходит, когда нужно натренировать модель под конкретные задачи, например, для суммаризации или разметки данных.

Для дистилляции знаний в «младшую» модель можно использовать метод Chain-of-Thought Prompting. Суть: просить LLM давать не только ответ, но и описывать цепочку рассуждений, которые к нему привели. Как показывают исследования, такой подход существенно увеличивает качество ответов на некоторых датасетах.

К примеру, авторы статьи Distilling Step-by-Step! попросили «ученика» предсказывать не только ответы «учителя», но и обоснования, чередуя запросы. Так маленькая модель тренируется думать как большая LLM, а не просто копирует ответы и поведение — на некоторых датасетах этот подход даёт отличный результат.

Кроме того, можно использовать датасет, составленный по reward-модели. В этом случае «ученик» будет тренироваться не на всех ответах «учителя», а только на тех, которые reward-модель считает хорошими, что тоже может улучшить результаты.

Наконец, можно расширить датасет, на котором учится младшая модель, с помощью генерации с разными параметрами вроде температуры или seed. Набор данных по одному промту получится более разнообразным, а поведение «ученика» в теории должно больше походить на поведение «учителя».

На этом всё. Спасибо, что прочитали! Делитесь опытом и впечатлениями от поста в комментариях! А во второй части текста мы разберём другие методы дистилляции и, конечно, затронем MiniLLM. Оставайтесь на связи!

Разбор помог подготовить ❣ Сергей Воробьев

@stuffyNLP