В этом параграфе вы познакомитесь с продвинутой техникой машинного обучения, получившей название дистилляции знаний. Дистилляция знаний (knowledge distillation) — это способ обучения в первую очередь нейросетевых моделей машинного обучения, направленный на передачу знаний от модели-учителя к модели-ученику.
Слишком абстрактное определение? Соглашусь, но в последние годы дистилляция знаний как поле исследований сильно разрослась и стала включать в себя множество новых и, возможно, даже неожиданных сценариев применения. Так, авторы статьи 2020 года утверждают, что смогли добиться примерной инвариантности выходов полносвязной сети к сдвигу входа-картинки с помощью дистилляции в неё знаний из сверточной сети.
Таким образом получается, что дистилляция знаний может применяться для того, чтобы передавать так называемые inductive biases от одной сети к другой. Схожие доводы встречаются и в статьях безумно популярного на момент написания данного параграфа направления трансформеров для компьютерного зрения.
Так, использование дистилляции знаний оказалась важным компонентом для получения хорошего качества предсказания на ImageNet от ViT без использования дополнительных данных. Впоследствии данный подход использовался и в других трансформерах для компьютерного зрения, например, в LeViT.
Тем не менее, среди всего разнообразия применений дистилляции знаний наиболее ярко выделяется одно — сжатие моделей.
Сжатие моделей
Задача сжатия моделей проистекает из следующего наблюдения. Неоднократно было замечено, что в широком диапазоне практически значимых задач машинного обучения точность предсказания модели существенно зависит от её размера. При этом зачастую данная зависимость выглядит достаточно тривиально: последовательное увеличение размеров модели позволяет последовательно улучшать точность её предсказаний.
Однако такой безграничный рост приводит к ряду проблем, связанных с практическим применением итоговых моделей. Сюда относятся рост времени обучения больших моделей и повышенные аппетиты таких моделей к размерам и качеству обучающей выборки. Кроме того, большие модели нередко требуют более дорогостоящего вычислительного оборудования для эффективного применения, особенно если мы говорим об обработке большого количества запросов в сжатые сроки. А для некоторых сценариев, таких как предсказание в реальном времени и/или на мобильных устройствах, применение большой модели может оказаться вовсе невозможным.
Эти проблемы породили каждая свою ветвь исследований. Так в последние годы де-факто стандартным способом обучения даже относительно компактных моделей стало использование mixed-precision training, которое позволяет ускорить обучение более или менее любых сетей на современных графических процессорах, при этом практически без потерь в итоговом качестве. Для борьбы с недостатком обучающих данных была предложена целая плеяда методов self-supervised pretraining, и новые появляются до сих пор. Сжатие моделей же концентрируется на решении проблем, связанных с этапом применения уже обученных моделей.
Как можно догадаться из названия, задача сжатия моделей заключается в том, чтобы взять большую модель, и сжать её в как можно более компактную модель при этом по возможности минимально пожертвовав качеством предсказания.
Исторически задачу сжатия моделей пытались решать множеством разных способов. Классическим примером здесь может служить прунинг, где модель обучается специальным образом (например, с использованием L2 регуляризации) так, чтобы часть весов в итоге можно было занулить и исключить из итоговой модели.
Однако методы данного семейства, как правило, страдают от двух основных проблем.
- Во-первых, простое удаление части весов каждого из слоёв обычно показывает лишь незначительное ускорение в применении итоговой модели за исключением случаев использования специализированной аппаратуры
- Во-вторых, наивный прунинг нередко приводит к существенной просадке в качестве предсказания сжатой модели, причём соотношение степени сжатия и качества итоговой модели едва ли возможно контролировать. Чтобы обойти данные ограничения, и была предложена техника дистилляции знаний.
Хинтоновская дистилляция знаний
Первой статьёй, в которой можно встретить дистилляцию знаний в современном виде является статья Хинтона и др. 2014 года.
В ней авторы рассматривают задачу классификации картинок и предлагают использовать предсказания большой заранее обученной модели-учителя в качестве новой, мягкой, разметки, которую будет пытаться повторить компактный ученик.
Формулировка
Авторы предлагают использовать дивергенцию Кульбака-Лейблера между предсказаниями учителя и ученика в качестве дополнительного слагаемого к стандартной функции потерь — кросс-энтропии между предсказанием ученика и жёсткой разметкой:
Здесь через обозначена функция потерь для дистилляции знаний. Под подразумевается число объектов, а под — классов, представленных в обучающей выборке. Через обозначена жёсткая разметка:
Через обозначены вероятности классов, предсказанные моделью-учеником, а через — моделью-учителем. Коэффициент позволяет настраивать баланс между решением исходной задачи и повторением предсказаний учителя.
В последнем переходе учтено, что логарифм частного раскладывается в разность логарифмов, после чего один из членов можно исключить, поскольку он не зависит от оптимизируемых значений . В дальнейшем для упрощения выкладок под будет подразумеваться именно последнее выражение.
Мотивация
Свой выбор функции потерь авторы мотивируют следующим образом. Широко известным фактом является то, что при классификации картинок на достаточно больших и разнообразных датасетах большие нейронные сети стабильно показывают лучшие результаты по сравнению с компактными. Однако также хорошо известно, что даже сравнительно небольшие нейронные сети способны приближать очень широкий спектр функций.
В таком случае можно предположить, что проблема обучения компактных сетей заключается не в том, что компактная модель не способна приблизить ту же функцию, что и большая, а в том, что компактная модель не способна самостоятельно выучить данную функцию из исходных данных. В таком случае потенциально мы можем подтолкнуть компактную модель к выучиванию более информативного представления путем модификации функции потерь.
Как этого добиться? Давайте возьмем заведомо более информативное представление, выученное большой моделью-учителем, и добавим в функцию потерь слагаемое, которое будет учить модель-ученика повторять его. В случае решения задачи классификации KL-дивергенция является именно таким слагаемым.
Есть и другой способ взглянуть на хинтоновскую дистилляцию знаний. Минимизация отличается от стандартного обучения, тем, что мы дополнительно минимизируем расстояние между предсказаниями ученика и учителя. От стандартной разметки такая целевая переменная отличается наличием так называемых теневых знаний (dark knowledge), которые состоят из вероятностей принадлежности объекта ко всем классам, помимо истинного.
Благодаря теневым знаниям модели-ученику во время обучения доступна дополнительная информация о взаимоотношениях между представленными в обучающей выборке классами, а также схожести отдельных объектов и классов.
Чтобы проверить данную гипотезу, авторы проводят следующий эксперимент. Сначала они обучают модель-учителя классифицировать картинки на датасете MNIST. После этого авторы обучают компактную модель-ученика с помощью ранее полученного учителя на тех же данных, но опуская при этом все картинки цифры . После этого авторы показывают, что, если исправить коэффициент сдвига для данного класса в последнем слое сети-ученика с помощью небольшой валидационной выборки, сеть способна верно определить картинок тройки, несмотря на то, что во время обучения она не видела ни одного примера.
Также косвенным подтверждением данной гипотезы можно считать и тот факт, что при использовании довольно популярной сейчас техники сглаживания разметки (label smoothing), эффективность дистилляции знаний заметно падает. Именно теневые знания на данный момент являются де-факто стандартным объяснением эффекта Хинтоновской дистилляции знаний.
Использование температуры при подсчете KL-дивергенции
В качестве дополнительной эвристики авторы также предлагают перед взятием KL-дивергенции сглаживать распределения учителя и ученика с помощью температуры , то есть вместо и считать KL-дивергенцию между и , где:
Здесь с помощью обозначены логиты классов, предсказанные моделью-учеником. Формула для выглядит аналогично.
Зачем нужна температура? Давайте рассмотрим формулу дополнительного слагаемого функции потерь. Для простоты выкладок я ограничусь функцией потерь для единственного объекта под номером , а также опущу постоянный множитель , который также не существенен для данного повествования.
Вспомним, что коэффициенты приходят из предобученной модели-учителя, а значит являются константными с точки зрения процесса оптимизации.
В таком случае несложно видеть, что мы имеем дело с чем-то очень близким к стандартной кросс-энтропийной функции потерь, но таргет — это уже не one-shot закодированные номера классов, а что-то более интересное. В таком случае компоненты предсказания ученика, которые отвечают классам, оценённым учителем, как наиболее вероятные, получат большие веса и сформируют каркас итоговой функции потерь.
В то же время все прочие компоненты получат околонулевые коэффициенты и влияния на функцию потерь практически не окажут. В какой-то степени эффект от этого может быть позитивным. Действительно, так как для преобразования предсказания нейронной сети в распределение вероятностей мы используем , итоговая модель не может предсказать строго нулевую вероятность. Поэтому типичное предсказание обученной сети содержит в себе множество практически нулевых значений.
При этом порядок между данными значениями определяется в первую очередь не похожестью объекта на представителей данных классов, а конкретным исходом стохастического процесса обучения данной модели. В таком случае нам вовсе не хотелось бы вынуждать ученика воспроизводить данный порядок, если ценой тому будет ухудшение точности предсказания истинного класса.
С другой стороны, нейронные сети являются зачастую излишне уверенными в себе классификаторами: их предсказание часто содержит ярко выраженный максимум, вероятность которого близка к единице даже в тех случаях, когда модели стоило бы усомниться. К сожалению, для нас это значит, что при дистилляции знаний из такой модели мы рискуем попасть в ситуацию, что итоговые веса настолько малы для всех классов, кроме истинного, что наше дополнительное слагаемое по сути повторяет стандартную кросс энтропию и не способно внести хоть сколь-нибудь заметный вклад в обучение модели-ученика.
Этот эффект можно нивелировать путем сглаживания предсказания учителя таким образом, чтобы сделать распределение ближе к равномерному, для чего, собственно и используется температура.
В таком случае функция потерь задается следующим образом:
Но в данную формулу незаметно закралась одна неприятная деталь. Давайте рассмотрим градиент второго слагаемого в скобках. Как и в прошлый раз, для простоты выкладок я ограничусь случаем единственного объекта под номером и опущу константный множитель :
Здесь легко узнаётся формула кросс-энтропийной функции потерь, градиент которой по логитам считается следующим образом:
Доказательство формулы для градиента кросс-энтропийной функции потерь.
Так как мы будем искать частную производную функции потерь по логитам, давайте сначала выразим через них саму функцию потерь:
Теперь мы готовы брать производную:
Поскольку для каждого фиксированного является вектором вероятностей, то , откуда мы и получаем искомую формулу.
Можно видеть, что при изменении температуры баланс между слагаемыми функции потерь (качеством решения задачи и качеством повторения предсказания учителя) нарушается. Действительно, если раньше мы настраивали его путём выбора подходящего коэффициента , то теперь мы приходим к тому, что при изменении температуры коэффициент необходимо также менять: иначе при взятии градиента одно слагаемое функции потерь будет разделено на , а другое останется неизменным.
Разумным кажется ввести множитель в формулу для явным образом. Однако прежде, чем мы сделаем это, давайте ещё раз внимательно посмотрим на получившийся градиент:
где через обозначены логиты, предсказанные моделью-учителем.
Давайте теперь устремим к бесконечности. Раскладывая экспоненты в ряд Тейлора до первого слагаемого, получаем:
Вспомним теперь, что результат применения преобразования не зависит от сдвига на константу, поэтому на выходе из нейронной сети мы можем вычитать из логитов среднее значение таким образом, чтобы . В таком случае, предыдущая формула упрощается до:
Из этой формулы следует два вывода.
- Во-первых, можно видеть, что для соблюдения баланса второе слагаемое в правильнее будет домножить не на , а на .
- Во-вторых, в данной формуле можно узнать градиент квадратичной функции потерь между векторами логитов.
То есть при стремлении температуры к бесконечности градиент второго слагаемого в стремится к градиенту квадрата нормы разности между логитами модели-ученика и модели-учителя.
Таким образом, мы приходим к финальной версии функции потерь:
Описанная выше статья произвела настоящий фурор в 2014 году. Дистилляция знаний путем минимизации KL-дивергенции между предсказаниями ученика и учителя хорошо зарекомендовала себя на практике и породила целый ряд исследований, направленных на использование и усовершенствование предложенного подхода.
Вместе с методом прижилось и понятие теневых знаний, и его довольно часто можно встретить в статьях, посвящённых данной тематике. Кроме того, зародилась традиция изучения дистилляции знаний на примере задачи классификации картинок.
Дальше по ходу параграфа мы ещё не раз столкнёмся с тем, что авторы различных методов часто прилагают результаты экспериментов на таких датасетах, как CIFAR-10, CIFAR-100, ImageNet и так далее.
Тем не менее, сети для работы с данными других модальностей тоже дистиллируют, и начнем мы с разбора статьи, которая использует предложенный метод для решения задачи языкового моделирования (language modelling).
DistilBERT как пример хинтоновской дистилляции
Одним из наиболее выдающихся примеров применения Хинтоновской дистилляции можно считать модель DistilBERT, которая сохраняет 97% качества модели BERT (согласно бенчмарку GLUE), используя при этом на 40% меньше параметров и требуя на 60% меньше времени при применении. При этом столь выдающийся результат авторы получают, используя хинтоновский подход практически без изменений.
По аналогии с тем, как это делалось для модели-учителя (в роли которого выступает BERT), авторы обучают свою модель решать задачу маскированного языкового моделирования. В дополнение к хинтоновской функции потерь использовалось ещё косинусное расстояние между итоговыми векторными представлениями токенов, полученными с помощью ученика и учителя, разворачивая представлений ученика в сторону направлений, задаваемых представлениями модели-учителя.
Ещё одна интересная деталь в этой статье — способ инициализации модели-ученика. Действительно, в качестве архитектуры для своей сети, авторы решили переиспользовать архитектуру самого BERT, но с уменьшенным вдвое числом слоёв для ускорения.
Авторы замечают, что большинство операций, которые используются в трансформерах, уже достаточно хорошо оптимизированы во всех популярных библиотеках, поэтому изменение размера внутренних представлений оказывает существенно меньшее влияние на итоговое время применения сети, нежели изменение количества слоёв. Поэтому в статья фокусировалась на сжатии модели именно в глубину, оставляя ширину неизменной.
Поскольку веса слоёв модели-ученика имеют при таком подходе такие же размерности, что и веса слоёв модели-учителя, последние можно использовать при инициализации. Ровно так авторы и поступают, копируя веса каждого второго слоя исходной модели для инициализации DistilBERT.
Может показаться, что умная инициализация не критична и наихудшим следствием использования более примитивной стратегии будет всего лишь увеличение времени, требуемого для обучения модели-ученика. Но авторы провели ablation study и выяснили, что обучение без умной инициализации приводит к потере почти процентных пункта итоговой метрики (обученная без неё модель сохраняет лишь качества модели-учителя).
Для сравнения, исключение из функции потерь кросс-энтропии между предсказанием ученика и истинной разметки приводит к потере лишь процентных пункта итоговой метрики, а исключение KL-дивергенции приводит к потере процентных пункта.
Интересно, что двумя годами позднее вышла независимая статья, авторы которой показали, что хинтоновская дистилляция — это очень сложная оптимизационная задача со множеством локальных минимумов, которые сильно усложняют поиск глобального.
Поскольку статья была написана независимо другими авторами, конкретный пример DistilBERT там не изучается, однако в целом авторы приходят к выводу, что умная инициализация может быть ключевым элементом для успеха дистилляции знаний.
Дополнительные источники знаний для дистилляции
Несмотря на широкий успех хинтоновского подхода, дистилляция знаний им не ограничивается.
Одно из наиболее очевидных направлений улучшения предложенного метода — это использование дополнительных способов передачи знаний от модели-учителя к модели-ученику. Действительно, в хинтоновской постановке единственный канал передачи знаний — это выходы с последнего слоя модели-учителя.
Однако в случае нейронных сетей это отнюдь не единственный доступный нам источник информации. Например, можно использовать веса модели-учителя для умной инициализации, как при обучении DistilBERT. К сожалению, поскольку дистилляция знаний практически всегда сопряжена со сжатием модели, не всегда получается непосредственно использовать веса учителя, и в каждом отдельном случае приходится изобретать специализированные трюки.
По этой причине DistilBERT — это единственная известная автору этого параграфа модель, в которой удалось добиться улучшения результатов благодаря использованию весов модели-учителя для умной инициализации.
Тем не менее, в нейронных сетях можно найти и другие источники информации. Хинтоновская дистилляция использует только выходы с последнего слоя сети. Почему бы нам дополнительно не использовать выходы промежуточных слоев? И действительно, исследования показывают, что использование выходов промежуточных слоев позволяет улучшить результаты дистилляции знаний.
Для получения прироста качества авторы предлагают выбрать один или несколько промежуточных слоев модели-учителя, сопоставить каждому из них промежуточный слой модели-ученика, после чего использовать квадрат нормы разности выходов итоговых пар слоев в качестве дополнительного слагаемого к хинтоновской функции потерь.
К сожалению, несмотря на кажущуюся прямолинейность данного подхода, здесь возникают две сложности.
Сложность №1
Мы явным образом предполагаем наличие заранее выбранных пар слоёв, оставляя за бортом вопрос о том, каким образом их собственно стоит выбирать. Поскольку дополнительные слагаемые функции потерь по сути обучают модель-ученика повторять промежуточные представления модели-учителя, разумным кажется сохранять порядок слоёв: слои из середины модели-учителя сопоставлять со слоями из середины модели-ученика, а слои, находящиеся ближе к концу модели-учителя, — со слоями, находящимися ближе к концу модели-ученика.
В частности, авторы оригинальной статьи просто берут средний слой в каждой из моделей и используют их в качестве своей единственной пары, однако это в большей степени связано с тем, что статья была написана в 2014 году и рассматривала достаточно маленькие по современным меркам модели. Более свежие статьи, как правило, работают с более глубокими сетями, а потому используют большее количество пар слоёв.
Так, авторы следующей работы рассматривают глубокие сверточные сети с промежуточными связями (residual connections) и предлагают разбивать каждую из моделей на группы блоков с промежуточной связью таким образом, чтобы итоговое количество групп совпало. Пример такой разбивки можно видеть на картинке ниже.
Здесь к каждой группе относится по три блока в модели-учителе и по два блока в модели-ученике. После этого выходы каждой такой группы можно сопоставить друг другу и использовать для дистилляции знаний.
После того, как пары слоёв были выбраны, перед нами может возникнуть и второе препятствие.
Сложность №2
Что, если выходы выбранных слоёв различаются по размерам? Такая ситуация запросто может случиться, ведь мы хотим, чтобы модель-ученик была поменьше, а один из способов сжатия — как раз уменьшение количества нейронов в полносвязных слоях.
В таком случае авторы оригинальной статьи предлагают использовать дополнительное преобразование, чтобы придать выходам модели-ученика нужные размеры (например, линейный слой).
Такие слои обучаются совместно с моделью-учеником, а после исключаются из сети при применении. В более поздних работах встречаются и другие, более продвинутые преобразования.
Несмотря на кажущуюся интуитивность дистилляции промежуточных выходов, практическое применение это метода, к сожалению, осложняется необходимостью выбора целого ряда гиперпараметров. Скажем, оптимальные тактики выбора пар слоёв для дистилляции или дополнительных преобразований для выравнивания размерностей выходов до сих пор являются предметами активных исследований, точно так же, как и функции потерь для оптимизации.
Иерархия методов дистилляции знаний
Выше мы рассмотрели два подхода к дистилляции знаний: хинтоновскую дистилляцию и дистилляцию промежуточных представлений. Как мы уже упоминали ранее, в последние годы область применения дистилляции знаний сильно разрослась, и новые методы появляются день ото дня.
Это породило довольно естественное желание систематизировать предложенные методы в некоторую иерархию. Мы рассмотрим две классификации методов:
- по режиму дистилляции,
- по области применения.
Режимы дистилляции знаний
Различные подходы к дистилляции знаний принято делить по так называемым режимам. Выделяют три основных режима дистилляции знаний: offline-, online- и самодистилляция.
Offline-дистилляция знаний
Все рассмотренные выше статьи так или иначе следуют некоторой общей схеме: в качестве учителя используется большая заранее обученная модель, знания из которой дистиллируются в ученика, в то время как сам учитель остается неизменным. Дистилляция в таком режиме получила название offline-дистилляции знаний.
Но что делать, если большой предобученный учитель для вашей задачи не доступен? Что если модель-учитель не помещается на доступную нам видеокарту, из-за чего обучение или вовсе невозможно, или требует в десятки раз больше времени, по сравнению с обучением желаемой модели-ученика? Что, если набор данных, описывающий вашу задачу, невелик, и большая модель может переобучиться на нём, делая дистилляцию знаний как минимум неэффективной, а возможно и вредной для итогового качества ученика?
Тут на помощь приходит online-дистилляция знаний.
Online дистилляция знаний
В качестве альтернативы авторы этой статьи предлагают брать в качестве учителя модель такой же архитектуры, что и ученик, и обучать обе модели одновременно.
То есть вместо одной модели мы случайно инициализируем две, после чего на каждом шаге обучения для каждой из моделей мы минимизируем , где в качестве учителя выступает другая модель.
В таком случае в начале обучения градиент дистилляционного члена не будет иметь какого-то чёткого направления, а обучение обеих моделей будет происходить преимущественно за счет минимизации обычной функции потерь. На поздних же этапах обучения в дело включится и KL-дивергенция, что позволит дополнительно повысить качество каждой из моделей.
Почему данный подход работает? Широко известно, что в ряде задач ансамблирование нескольких одинаковых нейронных сетей, одинаково обученных на одних и тех же данных, но из разных случайных инициализаций, дает прирост в итоговой метрике.
Этот факт подталкивает нас к выводу о том, что в зависимости от инициализации одна и та же нейронная сеть вычленяет из данных разные закономерности. Опираясь на данный вывод, авторы вышеупомянутой статьи утверждают, что в предложенной постановке каждая из моделей в процессе обучения может воспользоваться информацией, которая иначе была бы доступна только модели, стартовавшей из другой инициализации.
Авторы проводят ряд экспериментов с моделями разных размеров, обучая их на датасетах CIFAR-100 и Market-1501, и показывают, что использование даже одной дополнительной модели позволяет добиться заметного улучшения в качестве предсказаний обучаемой модели.
Так на датасете CIFAR-100 совместное обучение ансамбля из двух моделей практически во всех экспериментах дает прирост в процентных пункта к итоговой точности предсказания, причем метод позволяет достигнуть положительного эффекта даже для самой большой из рассмотренных моделей при её совместном обучении с самой малой моделью. Кроме того, авторы проводят ряд экспериментов, в которых сравнивают offline-дистилляцию большей модели в меньшую с их совместным обучением и показывают, что предложенный метод позволяет добиться лучших результатов.
Online-постановка естественным образом обобщается на случай большего числа моделей в обучаемом ансамбле. В таком случае в качестве дистилляционного слагаемого авторы предлагают минимизировать среднее значение KL-дивергенций от текущей модели до предсказаний каждой из других моделей в ансамбле, поскольку минимизация KL-дивергенции до усредненных вероятностей приводит к худшему результату.
При этом авторы в своих экспериментах показывают, что увеличение числа моделей в ансамбле приводит к улучшению результатов обучения. Кроме того авторы отмечают, что для ускорения обучения можно достаточно эффективно использовать несколько видеокарт, поскольку на каждом шаге между устройствами передавать необходимо только результаты предсказания.
Самодистилляция
В качестве отдельного режима дистилляции знаний принято выделять также самодистилляцию (self distillation), при которой учитель и ученик являются одной и той же моделью. Самодистилляция включает в себя две основные группы методов.
Первая группа методов направлена на использование информации, которая накапливается в модели во время обучения, для дополнительного улучшения качества предсказаний той же самой модели. Методы данной группы являются как бы продолжением идей online дистилляции знаний, поскольку учитель и ученик обучаются одновременно.
Хороший пример метода из данной группы можно найти в этой статье, где авторы пытаются заставить представления менее глубоких слоёв быть эквивалентными представлениям более глубоких слоёв. А именно, авторы предлагают разделить сеть на несколько частей ( в статье) и после каждой такой части добавить небольшую предсказательную голову. Все такие головы обучаются путем минимизации суммы трёх слагаемых:
- кросс-энтропии с истинной разметкой;
- KL-дивергенции с предсказаниями полной сети;
- квадратичной функции потерь между промежуточными представлениями данной головы и выходом последней части сети.
Таким образом авторы добиваются от ResNet50 точности предсказания на тестовой выборке CIFAR-100 с минимальным замедлением обучения.
Для сравнения, стандартное обучение такой же сети позволяет добиться лишь точности предсказания, а дистилляция из ResNet152 (которая, в свою очередь, показывает точность предсказания в ) позволяет улучшить данный показатель лишь до .
При этом обучение в предложенном режиме занимает часов (обычное обучение занимает часа), а дистилляция из ResNet152 занимает уже часов без учета обучения модели учителя (что требует дополнительных часов).
Вторая группа методов по сути заключается в offline дистилляции из обученной модели в новую модель такой же архитектуры. То есть мы выбираем некоторую архитектуру нейронной сети, обучаем одну модель стандартным образом, а затем обучаем точно такую же модель из новой случайной инициализации с использованием хинтоновской дистилляции из ранее обученной модели.
Стоит заметить, что с хинтоновской точки зрения данное действие едва ли способно улучшить итоговое качество модели. Действительно, будучи точно такой же моделью, ученик обладает идентичной способностью к обучению, а значит учитель едва ли может предоставить ему какую-либо дополнительную информацию во время обучения.
Поэтому такая самодистилляция изначально была предложена как метод изучения процесса хинтоновской дистилляции знаний, поскольку в такой постановке у задачи минимизации KL-дивергенции гарантированно есть глобальный минимум, причем мы даже знаем точку, в которой он достигается. В частности именно с помощью данного метода авторы ранее упомянутой статьи демонстрируют, что хинтоновская дистилляция знаний является сложной оптимизационной задачей.
Тем удивительнее, что авторы статьи 2018 года обнаружили, что самодистилляция в предложенной выше постановке позволяет получить прирост в обобщающей способности итоговой модели. Так, они проводят ряд экспериментов с моделями DenseNet и Wide-ResNet на датасете CIFAR-100 и показывают, например, что самодистилляция DenseNet-112-33 позволяет повысить точность предсказания на тестовой выборке с до .
Вопрос об источнике прироста качества в данном случае до сих пор в значительной степени открыт. Авторы статьи приписывают данный эффект комбинации умного сглаживания разметки и внесения в обучение информации о взаимоотношении классов в датасете. Но на наш взгляд эксперименты, которые предъявляют в статье в качестве доказательства этих гипотез, едва ли можно назвать убедительными.
Возможно, здесь в очередной раз проявляется то, что одинаковые модели могут вычленять из одних и тех же данных разные закономерности в зависимости от случайной инициализации, и дистилляция одной такой модели в другую позволяет ученику увидеть ранее недоступные ему связи.
Также хочется обратить внимание на интересную статью, вышедшую в 2020 году. В ней показывается, что в случае обучения с L2-регуляризацией предложенная выше самодистилляция производит неявный отбор признаков.
Ну и раз мы проходили мимо самодистилляции, здесь никак нельзя не упомянуть статью 2019 года, которая в течении практически года держала почетный статус SOTA на датасете ImageNet.
Её авторы предлагают подход, который во многом очень близок к описанному выше. Они обучают модель на исходном наборе данных, после чего используют её для разметки новых данных, взятых в данном случае из стороннего обширного набора данных JFT-300M (закрытый набор данных, который нередко упоминается в статьях от Google).
После этого авторы отбрасывают картинки, для которых модель дает неуверенные предсказания, чтобы избежать данных out-of-domain. Кроме того, они выравнивают размеры классов, чтобы избежать связанных с этим спецэффектов (согласно авторам статьи, модели меньшего размера показали себя более чувствительными к данной оптимизации).
Таким образом, авторы получают большое количество дополнительных шумно размеченных данных, на которых, совместно с основным набором, они обучают новую модель. Эту модель, в свою очередь, можно использовать для получения новой разметки для дополнительных данных, с помощью которых обучается следующая модель, и такие итерации можно продолжать произвольное количество раз.
В качестве разметки авторы предлагают использовать мягкую разметку, задаваемую моделью-учителем, но и бинаризованная разметка показывает схожие результаты на данных in-domain. Ключевая деталь здесь — что предсказание на новых данных производится без аугментаций, в то время как ученик учится воспроизводить разметку уже с высоким уровнем аугментации данных, а также с применением других техник регуляризации, таких как dropout и stochastic depth.
Авторы утверждают, что ученик обучается лучше переносить свои знания на новые данные. Предложенный метод позволил авторам добиться от модели EfficientNet-L2 точности предсказания в на тестовом наборе данных ImageNet, существенно улучшив результат исходной модели в и обновив мировой рекорд.
Области применения дистилляции знаний
Сжатие генеративных состязательных сетей
Подавляющее большинство рассмотренных выше статей так или иначе ограничиваются задачей классификации картинок. Такой выбор, хоть и не является случайным, всё же несёт больше исторический, нежели практический характер.
На самом деле, многие предложенные выше методы достаточно тривиально могут быть обобщены и на другие задачи машинного обучения. Например, метод дистилляции промежуточных представлений по сути вовсе никак не зависит от природы итоговых выходов модели, а потому может использоваться при сжатии практически любой модели.
В частности данный метод может быть использован для сжатия генеративных состязательных сетей. Так авторы довольно популярной статьи в данной области демонстрируют -, - и даже -кратное ускорение для ряда популярных pix2pix генеративных сетей, при этом не теряя в качестве генерации.
Как уже упоминалось ранее, авторы используют дистилляцию промежуточных представлений: модель-ученик учится минимизировать L2-расстояние между своим промежуточным представлением и промежуточным представлением модели-учителя. Но, так как данные представления имеют различное количество каналов (ученик выучивает более сжатое представление) авторы используют дополнительную свертку 1х1 над представлением ученика для сопоставления тензоров друг с другом.
Помимо дистилляции промежуточных представлений, авторы также пользуются наличием модели-учителя для того, чтобы получить парную картинку в случае обучения на неспаренных данных (как это происходит, например, в CycleGAN).
Парная картинка используется для минимизации L1 нормы разности с предсказанием модели. Кроме того, авторы во время обучения минимизируют и стандартную для генеративных состязательных сетей функцию потерь, при этом для модели-ученика используется такой же дискриминатор, что и для модели-учителя, что позволяет авторам инициализировать веса с помощью весов оригинального дискриминатора.
Таким образом авторы предлагают следующий рецепт для сжатия генеративных состязательных сетей. Сначала необходимо обучить модель-учителя. После этого нужно сконструировать сжатый генератор-ученика, скопировать (вместе с весами) дискриминатор и обучить получившуюся систему с помощью минимизации взвешенной суммы трёх функций потерь:
- Стандартной функции потерь генеративных состязательных сетей.
- L1-расстояния между предсказанием генератора и парной картинки. При этом если в данных парной картинки нет, вместо неё используется результат генерации моделью-учителем.
- L2-расстояния между промежуточными представлениями двух генераторов.
Этот метод хорошо себя зарекомендовал на практике и получил широкое распространение в своей нише.
Дистилляция знаний при квантизации моделей
Ещё одно интересное применение дистилляции знаний — улучшение результатов квантизации моделей. Техника квантизации нейронных сетей заключается в том, чтобы перевести часть весов или даже всю сеть из полной точности (как правило, float32) во float8 или даже float4.
Помимо очевидной экономии памяти, такое представление нередко позволяет использовать специальные ядра современных графических ускорителей или специальные регистры процессоров для достижения существенного ускорения при применении квантизованных моделей.
К сожалению, бесплатный сыр бывает только в мышеловке. Сжатое представление на то и называется сжатым, что является менее богатым, нежели полная точность. Поэтому большинство весов сети приходится изменять при сжатии, чтобы они попали на более грубую сетку. Разумеется изменения каждого отдельного веса может показаться незначительным, однако когда все веса сети незначительно изменяются, итоговый результат подсчетов может оказаться вовсе неузнаваемым.
Чтобы смягчить данный эффект, модель принято доучивать после квантизации. И вот здесь на помощь приходит дистилляция знаний: например, из сети полной точности в квантизованного ученика. Ровно к такой схеме приходят авторы этой статьи.
Итоговая схема выглядит следующим образом: мы обучаем сеть в полной точности, квантизуем ее веса и доучиваем ее в квантизованном виде с использованием дистилляции знаний из сети в полной точности.
Дистилляция знаний за пределами сжатия моделей
Хочется отметить, что сжатие моделей — это хоть и основное, но всё же не единственное применение дистилляции знаний. Так, раньше в этом параграфе уже упоминалась самодистилляция, которая позволяет получать прирост в обобщающей способности обучаемой модели без использования моделей большего размера.
Самодистилляцией, однако, примеры применения дистилляции знаний без сжатия модели не ограничиваются. Так, в 2020 году был предложен метод BYOL-предобучения без учителя, основанный на дистилляции знаний.
Метод BYOL направлен на предобучение моделей компьютерного зрения и основан на идее так называемого контрастивного предобучения (contrastive pretraining). Суть методов данного семейства заключается в том, чтобы обучать модель выдавать схожие представления для различных аугментаций одной и той же картинки.
Действительно, случайные патчи, вырезанные из фотографии автомобиля, скорее всего также являются фотографиями автомобиля. При этом, если в наших данных присутствует достаточное количество фотографий различных автомобилей, мы можем надеяться на то, что модель выучит некоторое общее понимание концепта автомобиля даже несмотря на то, что мы можем вовсе не знать, на каких конкретно картинках автомобили присутствовали, а на каких - нет.
Однако если мы хотим добиться от такой модели осмысленных представлений сначала нам необходимо преодолеть проблему коллапса представлений.
Действительно, у предложенной выше задачи есть тривиальное решение, в котором выход обучаемой сети не зависит от ее входа. В таком случае представления для произвольных аугментаций любой картинки будут совпадать, то есть функция потерь окажется нулевой. Тем не менее сами представления при этом окажутся совершенно бесполезными.
Поэтому различные методы контрастивного обучения отличаются в первую очередь как раз способами борьбы с проблемой коллапса представлений. Так, авторы метода BYOL часто сравнивают свои результаты с довольно свежим на момент написания статьи методом SimCLR, в котором предлагается обучать модель одновременно минимизируя расстояния между парами представлений для различных аугментаций одной картинки и максимизируя расстояния между представлениями для различных картинок.
При этом, для повышения эффективности такого обучения, во время генерации батча данных авторы сначала выбирают некоторое количество картинок из набора данных, затем для каждой картинки производят две различные случайные аугментации, после чего полученные картинки используются для создания одной позитивной пары, расстояние между представлениями которой будет минимизироваться, а также для создания множества негативных пар с аугментациями других картинок в батче, расстояния между представлениями которых будут максимизироваться.
Авторы BYOL подвергают данный подход критике, показывая, что для эффективного обучения SimCLR требует большого размера батча, а также довольно агрессивных аугментаций. В противном же случае качество обучаемых представлений заметно падает. Авторы BYOL объясняют данный эффект тем, что сам подход использования негативных пар является субоптимальным, поскольку требует аккуратного выбора негативных примеров. Поэтому свой метод авторы конструируют таким образом, чтобы модель обучалась только на позитивных парах картинок.
В таком случае каким образом авторам удается решить проблему коллапса представлений? Для этого, вместо минимизации расстояния между представлениями обучаемой сети для двух аугментаций одной картинки, авторы обучают свою модель минимизировать расстояние между представлением обучаемой (online) сети для одной аугментации и представлением для второй аугментации, которое задается уже другой, целевой (target) сетью. То есть в некотором смысле здесь происходит дистилляция знаний из целевой сети в обучаемую.
Последней важной деталью является природа целевой сети. Авторы BYOL замечают, что даже использование произвольно инициализированной сети в качестве целевой для предложенного выше метода обучения приводит к выучиванию обучаемой моделью осмысленных представлений.
Подробнее, линейный классификатор, обученный на основе выученных таким образом представлений картинок из набора данных ImageNet показывает тестовой точности предсказания, в то время как использование представлений задаваемых самой целевой сетью позволяет добиться лишь точности. Мотивированные данным наблюдением, авторы предлагают в качестве целевой использовать такую же сеть, что и обучаемая.
При этом:
- градиенты не текут через целевую модель, и она не обновляется на шаге градиентного спуска;
- обучаемая модель заканчивается дополнительным двухслойным перцептроном, который используется для преобразования её финальных представлений в представления целевой модели;
- веса целевой модели не меняются на шаге градиентного спуска, а вместо этого они обновляются между шагами с помощью экспоненциального сглаживания весов обучаемой модели:
где, следуя обозначениям из статьи, мы обозначили через и веса целевой и обучаемой моделей соответственно, а — вещественный параметр.
Предложенный метод позволяет авторам добиться уже тестовой точности от линейного классификатора на наборе данных ImageNet, заметно превосходя предложенные ранее методы self-supervised предобучения, и практически преодолевая разрыв между self-supervised и supervised обучением классификаторов на основе ResNet.
Стоить заметить, что сценарии применения дистилляции знаний отнюдь не ограничиваются выше рассмотренными. На данный момент уже существует множество различных подходов и алгоритмов, так или иначе связанных с дистилляцией знаний, и их количество растет день ото дня.
Данный параграф не ставит своей целью полный обзор таких методов. Вместо этого всем заинтересовавшимся я рекомендую обратить внимание на довольно исчерпывающий обзор от 2020 года. Здесь можно найти множество ссылок на актуальные к тому моменту статьи, в числе которых присутствует и большинство статей, упомянутых в этом параграфе.
Открытые проблемы
Завершим параграф упоминанием открытых проблем в области дистилляции знаний.
Действительно, несмотря на впечатляющие результаты, дистилляция знаний всё же не является идеальным методом, и ряд вопросов до сих пор остаются без ответа.
Например, с ростом популярности дистилляции знаний выяснилось, что использование учителя с большей обобщающей способностью не всегда приводит к улучшению обобщающей способности ученика. В какой-то степени данный эффект можно списать на то, что компактная модель-ученик упирается в пределы своего качества предсказания, и тогда использование более умного учителя уже не приносит дополнительной пользы.
Но это не объясняет, почему использование более точной модели в качестве учителя может приводить даже к ухудшению итоговой точности модели-ученика. В чём причина данного эффекта и как выбрать оптимального учителя для фиксированного ученика, до сих пор открытый вопрос.
И как уже упоминалось ранее, дистилляция знаний из одной сети в точно такую же нередко приводит к росту обобщающей способности ученика по сравнению со своим учителем. С точки зрения хинтоновской теории, которая является де-факто стандартным способом объяснения дистилляции знаний, это звучит абсурдно.
Модель-ученик гарантированно способна приблизить ту же функцию, что и модель-учитель. Тем не менее, этого не происходит, а модель-ученик выучивает свое собственное представление, которое нередко качественно превосходит представление учителя. Данный факт уже сложно объяснить в парадигме передачи знаний от учителя к ученику, потому что здесь ученик оказывается в состоянии получить больше знаний, нежели учитель способен передать. Несмотря на то, что на данную тему написана уже не одна статья, исчерпывающего объяснения пока нет.
Так или иначе, дистилляция знаний неоспоримо работает и является основным практическим подходом к сжатию нейросетевых моделей на данный момент.