Разбор DINO — идеи self-supervised обучения CV моделей
В этой статье я разберу идею для self-supervised обучения моделей для обработки изображений под названием DINO (если вы не знаете, что такое self-supervised обучение, то я делала об этом ликбез в тг-канале DLStories: ссылка). Мы узнаем, что за это идея такая, а также какие интересные свойства появляются у моделей, обученных с помощью DINO.
DINO был представлен еще в 2021 году в статье "Emerging Properties in Self-Supervised Vision Transformers". Кроме самого DINO, в статье приводится наблюдение, что у self-supervised трансформеров (например, у ViT, обученного с помощью DINO) обнаруживается интересное свойство: если подать такому трансформеру на вход картинку, то ее карты внимания будут неплохо сегментировать объект на картинке. Об этой находке и о том, как с ее помощью получить неплохую zero-shot сегментацию нескольких объектов на картинке, я писала в посте в телеграме тут.
А еще совсем недавно выпустили вторую версию DINO (DINOv2). В этой статье мы разберем устройство оригинального DINO, а в следующем статье — устройство DINOv2.
Как работает DINO
Идея DINO основана на принципе knowledge distillation (KD).
В двух словах, KD – это когда у нас есть две модели — модель-учитель и модель-студент. Модель-учитель умеет хорошо решать какую-то задачу. И мы учим модель-студента с нуля решать эту же задачу таким образом: подаем на вход учителю и студенту одинаковые данные. И с помощью некоторого лосса заставляем выходы модели-студента быть похожими на выходы модели-учителя.
Получается, студент учится решать задачу, имитируя поведение учителя на входных данных. Иногда к обучению модели-студента также добавляют и стандартный лосс между выходами модели-студента и правильными ответами на входных данные.
Чаще всего в KD модель-студент меньше, легковеснее, чем модель-учитель. То есть, KD используется для того, чтобы получить более легкую модель, которая умеет решать задачу примерно так же хорошо, как большая и умная модель-учитель.
Авторы из Meta адаптировали идею KD для self-supervised обучения, и получили DINO. Вот как именно это работает:
Берем датасет картинок. Из каждой картинки датасета вырезаем два глобальных патча картинки и несколько локальных. Глобальный патч — это кусок картинки, который покрывает >50% всей картинк (в статье берут размер 224х224). Локальные патчи — куски картинки размера 96х96.
Теперь делаем две копии одной и той же модели для обработки картинок: F_s и F_t. F_s — модель-студент, F_t — модель-учитель. Это могут быть модели любого типа: ResNet, ViT, или что-то еще.
Изначально F_s и F_t не обучены. Обучаем мы их так:
В течение эпохи веса F_t заморожены. Берем картинку x из датасета, подаем на вход F_t один из ее двух глобальных патчей, а на вход F_s подаем любой рандомный патч (локальный или глобальный, но не тот же, что подали на вход F_t). Обе модели выдают на выход эмбеддинги одинкового размера. Эти эмбеддинги мы сравниваем с помощью лосса и бэкпропагейтим градиенты в модель-студент F_s:
То есть, тут идея примерно такая: в течение эпохи сеть-студент учится имитировать выходы замороженной сети-учителя. И сеть-студент получает на вход в основном локальные патчи, а сеть-учитель — только глобальные. Это как бы "учит" сеть-студента по локальным частям картинки понимать, что глобально на этой картинке изображено. Учась на такую задачу, сеть-студент может выучить что-то умное о том, как устроены картинки.
Ну и в конце каждой эпохи веса модели-учителя F_t обновляются на основе весов модели-студента F_s c помощью running mean:
Авторы замечают, что пробовали разные варианты того, как и когда обновлять веса модели-учителя. Running mean оказался лучше всего. Вообще, в последнее время я эту идею вижу много где в местах, в которых используется одновременное обучение две моделей под одну задачу. Нужно будет про это почитать подробнее и понять, есть ли объяснения, почему именно так работает лучше всего.
Как избежать collapse
Self-supervised методы часто страдают проблемой коллапса. Это когда модель в процессе обучения начинает для всех картинок датасета выдавать ограниченный набор векторов на выходе. Например, модель может начать для вообще всех картинок датасета выдавать один и тот же вектор. Это логично, так как такое поведение позволит модели лучшим образом минимизировать лоск-функцию, на которую модель обучают.
Есть много разных способов борьбы с коллапсом. Самый известный, наверное, contrastive learning — это когда модель учат не только выдавать близкие векторы на похожие входные объекты ("позитивные" пары объектов), но еще и выдавать далекие векторы на разные входные объекты ("негативные" пары объектов). Например, так можно обучать модель для распознавания лиц с помощью Triplet Loss. Подаем на вход модели лица одного и того же человека и просим модель выдать на них похожие векторы. Потом подаем на вход модели лица разных людей и просим ее выдать как можно более далекие векторы.
Вот иллюстрация идеи contrastive learning:
Из-за этой же проблемы в DINO не получится обновлять параметры обеих сетей (учителя и студента) одновременно. То, что параметры сети-учителя заморожены на протяжении эпохи обучения сети-студента, помогает избежать того, что сеть-студент и сеть-учитель одновременно быстро скатываются в коллапс. Пока параметры учителя заморожены, сеть-студент вынуждена пытаться подстроиться под выходы сети-учителя, что не позволяет ей просто начать выдавать одинаковые векторы на любой вход.
Однако в DINO contrastive learning не используется. Возможно, потому, что не совсем понятно, как лучше генерировать негативные пары картинок. Вместо этого авторы DINO предлагают использовать другие две идеи — centering and sharpening of the momentum teacher outputs. Каждая из этих идей по отдельности позволяет нивелировать один из факторов, который приводит к коллапсу, а совместное их применение позволяет DINO не впадать в коллапс совсем и успешно обучаться. Я в этих методах борьбы с коллапсом еще не разбиралась, но как разберусь, напишу об этом статью. Пока что о них можно почитать в секциях 3.1 и 5.3 статьи о DINO.
Результаты
Идея обучения DINO действительно позволяет моделям выучивать что-то полезное о картинках.
Чем это подтверждается: авторы провели много экспериментов, подтверждающих, что ViT, обученный с помощью DINO, выучивает полезную информацию о картинках. Один из них — стандартный способ сравнить "хорошесть" self-supervised модели. Он такой:
- Берем модель F, обучаем в self-supervised режиме на тренировочной части ImageNet;
- Из обученной модели F получаем эмбецдинги всех картинок тренировочной части ImageNet;
- На этих эмбеддингах обучаем простой классификатор (например, однослойную полносвязную сеть или KNN);
- Получаем из F эмбеддинги для картинок валидационной части ImageNet, прогоняем их через классификатор, получаем ответы, считаем accuracy.
Проделав такую процедуру для разных self-supervised методов, можно сравнить их "хорошесть" между собой (отмечу, что при сравнении архитектуры всех моделей должны быть одинаковы). А еще можно сравнить все эти модели с supervised подходом. Это все авторы статьи и сделали. Вот результаты:
Видно, что DINO выигрывает у многих self-supervised подходов, которые были SOTA'ми на момент выхода статьи.
В разделе 4 статьи описаны и другие сравнения ViT, обученного с помощью DINO, с другими подходами.
А здесь скажем вот что: как упоминалось в начале этого разбора, у ViT, обученного с помощью DINO, появляется интересное свойство: его self-attention карты начинают сегментировать объекты на картинке. Вот примеры того, как выглядят self-attention карты модели для разных пикселей:
Видно, что для каждого пикселя картинки self-attention карта, соответствующая этому пикселю, по сути сегментирует объект, который к этому пикселю относится.
На основе этого наблюдения можно построить zero-shot сегментацию изображений. Просто берем карту self-attention, выбираем порог и оставляем все пиксели, значение которых больше порога. Это и будет картой сегментации.
И это работает. Авторы провели такой эксперимент: взяли маленький ViT (ViT-s/8), обучили его в supervised режиме на задачу сегментации, и в self-supervised режиме с помощью DINO. Получилось, что self-attention карты DINO сегментируют объекты даже лучше, чем supervised модель! Вот результаты:
Заключение
Вот такая вот интересная идея self-supervised обучения vision моделей под названием DINO. Недавно вышла DINO-v2, ее я разберу в следующей статье.