Subclass Distillation

Rafael Müller, Simon Kornblith, Geoffrey Hinton

Статья: https://arxiv.org/abs/2002.03936



Что-то давно мы про дистилляцию не писали. А тут как раз свежая работа от Хинтона и ко. Или от ко и Хинтона. Как посмотреть.



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



А если классов мало, то и дополнительной информации мало. Надо с этим что-то делать.



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



В качестве модельного примера можно взять CIFAR-10 и превратить его в CIFAR-два-по-пять (CIFAR-2x5), где мы группируем наборы по 5 классов в два новых класса (1: airplane, automobile, bird, cat, deer; 2: dog, frog, horse, ship truck).



Далее модифицируется учитель, чтобы у него было 10 логитов на выходе (по числу всех подклассов), которые объединяются в два класса простым суммированием вероятностей после софтмакса.



Учитель обучается на задаче классификации двух классов, не 10, это важно. Подклассами его явно не кормят, он должен сам присвоить какие-то вероятности входящим в класс подклассам. Финальный кросс-энтропийный лосс смотрит только на итоговое качество бинарной классификации.



Ну на самом деле, чтобы стимулировать учителя выучивать нетривиальные подклассы (а то начнёт, скажем, в один подкласс всё запихивать, заполняя остальные нулями), есть ещё один стимулирующий auxiliary loss, подталкивающий сетку раскидывать разные примеры по разным подклассам (даже если они принадлежат одному).



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



Как альтернатива дистилляции подклассов есть ещё дистилляция предпоследнего слоя, где тоже содержится больше информации, чем в последнем. Но по идее в случае подклассов вся нерелевантная классификации информация убирается, чего не скажешь про предпоследний слой, там может быть всякое.



Метод дистилляции подклассов работает. Качество дистилляции повышает. Это во-первых. То есть та самая бинарная классификация на искусственно собранных двух классах CIFAR-2x5 работает лучше. Заодно и классификация на 10 классов тоже работает ощутимо лучше рандома, хоть её специально не обучали. А во-вторых, если глазами посмотреть на придуманные сетью подклассы, то они выглядят вполне разумно (хоть и неидеально).



Из менее модельных случаев проверяют метод на CelebA, тоже работает. И ещё на кликовом датасете Criteo, и там работает. Что интересно, на этом датасете и обычная дистилляция и предложенная дают качество выше учителя. Это если обучать на полном датасете (он огромен). Но дистилляция подклассов обучается гораздо быстрее и делает это лучше обычной в случае ограниченных данных.



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