Сравнения по качеству, скорости и размерам

Другие известные дистилляции BERT’а

(2019/03) “Distilling Task-Specific Knowledge from BERT into Simple Neural Networks”
Статья: https://arxiv.org/abs/1903.12136

В работе BERT дистиллируется в однослойную BiLSTM, получают результат сравнимый с EMLo при стократно меньшем числе параметров и в 15 раз меньшем времени инференса.

Как видно из таблицы в предыдущем посте, и DistilBERT, и TinyBERT этот результат бьют по качеству, хотя коэффициенты и сжатия, и ускорения здесь вроде выше.

(2019/08) “Patient Knowledge Distillation for BERT Model Compression”
Статья: https://arxiv.org/abs/1908.09355
Код: https://github.com/intersun/PKD-for-BERT-Model-Compression

Предложили метод под названием Patient Knowledge Distillation, который претендует на то что был первым (или среди первых), кто стал использовать при дистилляции не только выходное распределение, но и скрытые состояния учителя. Причём, ученик, похоже пытается имитировать только репрезентации для [CLS] токена.

Опять же в предыдущем посте есть сравнение и с ним, он получше DistilBERT, но слабее TinyBERT. Есть разные вариации по размеру (3 и 6), на 3-слойной ускорение около 4 раз относительно 12-слойной BERT-base.

(2019/09)“Extreme Language Model Compression with Optimal Subwords and Shared Projections”
Статья: https://arxiv.org/abs/1909.11687

Свежая статья из Гугла, засабмичена на ICLR 2020. Акцентировались на сокращении словаря, эмбеддинги которого занимают много места. Плюс также сокращают и размер скрытых состояний. Придумали метод дуального обучения, когда учитель одновременно с учеником пытается выучить сокращённый словарь. Если я не ошибаюсь, другие подходы не подразумевали, что учителя тоже надо обучать.

Сжали BERT-base в 60 раз, получили модель размером под 7Мб. Круто для мобильного. Тоже пробовали разные по размеру варианты, на самом разумном (среднем) с эмбеддингнами размером 96 сжатие получают под 20 раз (модель 22Мб), сокращение количества вычислений под 75 раз. По качеству результаты слабее TinyBERT.

Результаты BiLSTM_soft

Результаты BERT_PKD

Результаты дистиллированного BERT'а со сжатым словарём

Ну и дистилляция дистилляцией, но не забывайте, что появился ещё и весьма качественный ALBERT (тоже на ICLR 2020), который существенно легче оригинального BERT'а, но по качеству круче RoBERT'ы. И его по идее тоже можно отдистиллировать :)

Про ALBERT уже было тут: https://blog.inten.to/speeding-up-bert-5528e18bb4ea

Дистилляция: матчасть
Пора немного поговорить про истоки дистилляции и её суть.

В целом, как уже было описано, идея дистилляции подразумевает обучение более лёгкой модели-ученика, воспроизводящей поведение более сложной и тяжёлой модели-учителя.

Одной из ранних работ, на которую принято ссылаться, является:
Model Compression
Rich Caruana, Cristian Bucila, Alexandru Niculescu-Mizil
https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf

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

Поскольку неразмеченный датасет не всегда есть, авторы придумали метод (MUNGE) создания такого датасета (псевдо трейн сет), который лучше рандомного сэмплинга атрибутов или naive Bayes estimation.

В итоге на автоматически размеченных реальных или псевдо данных обучают нейросеть и получают сокращение размера модели и ускорение до 1000 раз, при этом не сильно теряя в качестве.

Затем Rich Caruana с Lei Jimmy Ba развили эту тему в работе:
Do Deep Nets Really Need to be Deep?
https://arxiv.org/abs/1312.6184

Здесь большая глубокая нейросеть заменяется на неглубокую с сохранением качества. Цель не в сокращении размера как таковом, а в попытке разобраться, так ли нужна глубина.

Model compression даёт интересное наблюдение: нельзя выучить маленькую модель с качеством как у большой, обучая её на тех же данных. Но маленькая модель в принципе способна выучить более точную функцию, и мы знаем, что функция эта не слишком сложна. Просто текущие алгоритмы оптимизации этого не позволяют. Для такого обучения почему-то требуется промежуточная сложная модель.

Соответственно этим авторы и занимаются в работе, показывая, что можно обучить неглубокие модели такого же качества, как и глубокие.

Снова используются данные, размеченные сложной моделью, transfer set (иногда это те же самые данные, на которых модель обучалась, иногда дополнительный датасет). Модель-ученик учится предсказывать логиты модели-учителя — выходы сети _ДО_ софтмакса, ибо на софтмаксе теряется много информации. Например, активации [10.0, 20.0, 30.0] и [-10.0, 0.0, 10.0] дают одинаковый результат после софтмакса, но они по сути существенно разные. Идея в том, что в этих числах закодированы тонкие взаимоотношения между классами, которые отсутствуют в разреженных hard labels, и эту информацию можно использовать для обучения ученика. Задача обучения сводится к задаче L2-регрессии на предсказание логитов.

В итоге на TIMIT и на CIFAR-10 смогли обучить неглубокие модели, приближающиеся по качеству к продвинутым моделям (CNN и ансамбль CNN) и обгоняющие DNN. Снова получилось, что имитирующие модели (они называют их mimic models) получаются более точными, чем если бы изначально были обучены на оригинальных данных.

В общем интересная история, неглубокую модель текущими методами оптимизации хорошо обучить с нуля не получается. Только через model compression (доразметку неразмеченного датасета) и/или качественную модель-учителя (считайте, дистилляцию).

Есть хорошее выступление Каруаны по данной теме:
https://www.youtube.com/watch?v=NedWl0lZw2E

Ну и наконец ещё одна работа, на которую чаще всего ссылаются:
Distilling the Knowledge in a Neural Network
Geoffrey Hinton, Oriol Vinyals, Jeff Dean
https://arxiv.org/abs/1503.02531

Авторы развивают работы Каруаны и коллег, предлагая метод дистилляции для нейросетей.

Идея всё та же, в распределениях вероятностей классов при предсказании моделью-учителем хранится ценная богатая структура, которую и надо использовать для клонирования поведения модели учителя. Для этого можно использовать как предсказания на оригинальном датасете, так и на любом другом.

Для этого предлагается метод дистилляции, более общий, чем у Каруаны, за счёт повышения температуры на выходном софтмаксе. Версия Каруаны с логитами показывается частным случаем этой процедуры.

Технически задача дистилляции сводится к минимизации кросс-энтропии между предсказаниями учителя и ученика. При обучении ученика в качестве soft labels берутся предсказания после софтмакса, но взятого с какой-то температурой выше 1 (чтобы получить менее жёсткое распределение). Когда ученик обучен, он использует для предсказания уже обычный софтмакс.

В работе есть ещё одно интересное ответвление про обучение на очень больших датасетах ансамблей моделей-специалистов, которые фокусируются на своём подмножестве классов. Это чем-то напоминает Mixture of Experts (MoE), но легче него по вычислениям и параллелизации.

Дистиллируем дальше
Дистилляция моделей — это ещё не всё.

Data Distillation: Towards Omni-Supervised Learning
Ilija Radosavovic, Piotr Dollár, Ross Girshick, Georgia Gkioxari, Kaiming He
Статья: https://arxiv.org/abs/1712.04440

Авторы из FAIR исследуют подход под названием omni-supervised learning, специальный вариант semi-supervised, когда в дополнение к размеченным данным используются неразмеченные надыбанные на масштабах интернета.

Для этого авторы предложили метод под названием data distillation, который в отличие от дистилляции моделей обучает студента не на ансамбле моделей, а на ансамбле предсказаний, выполненных одной моделью на базе множества трансформаций данных.

Процесс выглядит так:
1) Обучаем модель на размеченных данных
2) Генерируем аннотации для неразмеченных данных, делая предсказания на различных трансформациях данных и ансамблируя результаты
3) Обновляем (переобучаем) модель на увеличенном наборе данных

Подход с ансамблированием преобразований похож на метод повышения точности предсказаний на инференсе, когда со входной картинкой делалось несколько преобразований (повороты, кропы и т.п.), а затем результат предсказаний усреднялся (multi-transform inference).

По наблюдению авторов, агрегация предсказаний генерирует новое знание и модель в принципе может использовать эту информацию для self-training.

В работе генерировали hard labels (то есть собственно метки классов), а не soft labels (вектор вероятностей классов). Дальше эти метки использовались как настоящие. В обучении делали так, чтобы в минибатче всегда была смесь настоящих меток и вот этих вот сгенерированных.

Применили для нескольких задач комп.зрения: определение ключевых точек и детекция объектов. Трансформации традиционные: скейлинг и горизонтальный флип.

Подход работает, позволяет получать более точные модели, чем без такого метода. До точности моделей, обученных на изначально размеченном датасете такого же размера, не дотягивает, но где в реальной жизни такие датасеты взять, кейс нереальный. Зато неразмеченного добра при желании насобирать можно много, так что вполне применимый подход.

Продолжаем дистиллировать данные.

Dataset Distillation
Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, Alexei A. Efros
Статья: https://arxiv.org/abs/1811.10959
Страница: https://ssnl.github.io/dataset_distillation/
Код: https://github.com/SsnL/dataset-distillation

Теперь заход с другой стороны — что если модель не трогать, а большой датасет отдистиллировать в маленький? Например, превратить 60000 обучающих примеров MNIST’а в 10? Причём не как раньше, когда пытались оставлять наиболее значимые примеры, а прям вот создать синтетические примеры (возможно не из того многообразия, где обитают сами исходные данные), на которых модель хорошо и быстро обучится.

Оказывается, можно.

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

В прямолинейном варианте подход требует знания начальных весов сети. В более интересном варианте достаточно знать распределение, из которого сэмплили начальные веса. Для моделей, предобученных на другие задачи (AlexNet), метод может найти дистиллированный датасет для быстрого файнтюнинга.

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

На MNIST получилось отдистиллировать весь датасет в 10 примеров (по одному на класс!), на которых LeNet обучается до почти 94% точности за 2000 шагов градиентного спуска с фиксированной инициализацией. С рандомной инициализацией 88.5%. На CIFAR10 тоже интересные результаты.

Неплохо также работает адаптация между MNIST/USPS/SVHN.

Или, например, получается адаптация претренированной AlexNet на PASCAL-VOC и CUB-200 датасеты с одной картинкой на категорию.

Метод заточен на конкретные архитектуры и инициализации, есть куда расширять. Но выглядит очень интересно. Например, можно существенно сократить вычисления на оценку моделей в neural architecture search (помните эти страшные цифры? https://www.technologyreview.com/s/613630/training-a-single-ai-model-can-emit-as-much-carbon-as-five-cars-in-their-lifetimes/)

Статья подана на ICLR 2020.

Дистиллированные картинки выглядят интересно. Есть над чем подумать.

Отдельное направление для применения этой работы — data poisoning attack. Из сабмита на ICLR эту часть убрали :) Наверняка готовят отдельную статью, это обещает быть сильно.

Хитрый data poisoning с помощью дистилляции позволяет сгенерить небольшое число примеров, которые всего за одну итерацию обучения сильно портят одну конкретную категорию классификации, не трогая остальные. Без доступа к точным значениям весов модели!

MNIST: 10 images train test accuracy from 12.9% to 93.8% (with a fixed known initialization)