Назад
287

DINO: Self-distilation with no labels

287

Введение

Трансформеры сегодня все чаще применяются в различных доменах помимо NLP. Vision Transformers (ViT) рассматриваются как альтернатива сверточным сетям (convolutional neural network, CNN) в решении задач компьютерного зрения (CV). В мае 2021 года вышла статья Emerging Properties in Self-Supervised Vision Transformers, где авторы предложили задуматься над вопросом: почему ViT так хорошо себя проявляют в задачах компьютерного зрения, и является ли предобучение на больших датасетах с разметкой тому причиной. Сомнения в этой причине связаны со следующим фактом: основой успеха трансформерной архитектуры в больших языковых моделях является этап предобучения на огромном наборе данных без разметки с помощью подхода Self-Supervised Learning (SSL).

SSL — режим обучения моделей, при котором разметка формируется автоматически из внутренней структуры данных.

Например, на этапе предобучения таких языковых моделей, как BERT или GPT, в качестве разметки могли использоваться пропущенные слова или части предложения, которые моделям необходимо было предсказывать. В SSL этот этап обучения называется “pretext task”. Он позволяет извлечь из данных наиболее общую информацию, по сравнению с обычным режимом обучения с учителем (где модель выучивает лишь часть информации об объекте, ограниченную существующей разметкой). То есть полезная информация на выходе модели — некоторые признаки / векторное представление, с которым затем работают для решения так называемых “downstream” задач (специальных задач CV — классификации, сегментации и др.).

DINO

Сегодня мы рассмотрим один из самых эффективных на сегодняшний день SSL подходов — DINO (Self-DIstilation with NO labels). Его разработчики решили проанализировать, какие признаки можно выучить с помощью self-supervised в ViT трансформере. На рисунке 1 показаны визуализации весов внимания на выходе предобученного ViT c помощью DINO.

Рисунок 1. Визуализация весов внимания c помощью DINO

Есть две особенности признаков ViT, которых нет при обучении с учителем, в трансформерной и сверточной архитектурах:

  • признаки в явном виде содержат фрагменты визуальных сцен, а именно — границы объектов;
  • выученные с помощью SSL признаки — очень хорошие входные признаки для KNN классификатора. Он достигает 78.3 % точности на ImageNet без дополнительного дообучения.

Эти свойства модели, которую наделяет данный подход в процессе обучения, плюс возможность использования огромного набора картинок без разметки (по сути весь интернет) позволяют отнести ее к так называемым “Foundation Models”. Такие модели вбирают в себя огромный пласт знаний и позволяют решать без дообучения большинство известных задач с приемлемым качеством. При этом в отличие от моделей SAM или CLIP, здесь для обучения вообще не требуется разметка. Во-первых, это делает процесс обучения более дешевым. Во-вторых, позволяет извлечь больше информации и не ограничиться существующей разметкой.

Подробнее о подходе

DINO опирается на парадигму дистилляции знаний без разметки — Self-DIstillation with NO labels. Она подразумевает такое обучение нейронной сети (НС) “студент”, при котором ее выходы стягиваются к значениям выходов другой НС — “учителя”. При этом архитектуры обеих НС одинаковы, но веса разные.

Как видно на рисунке 2, в процессе обучения на вход “студенту” и “учителю” подается одно изображение, но после разных наборов аугментаций.

Рисунок 2. Схема обучения методом DINO

Вероятностный выход P сети g нормализуется функцией softmax:

\( \LARGE {~~~~~~~~~~~~~~~~~~~~~~~~~~~~ P_s(x)=\frac{exp(g_{\Theta_s}(x)/\tau_s)}{\Sigma_{k=1}^Kexp(g_{\Theta_s}(x)^k/\tau_s)}~ } \)

где \( P_s \) — выход сети “студент” g, x — входное изображение, K — размерность выхода сети, τ — температура софтмакса, \( \Theta_s \) — параметры сети “студент”.

В качестве функции потерь используется кросс-энтропия между прогнозом “студента” и “учителя”:

\( \LARGE~~~~~~~~~~~~~~~~~~~~~~~~~~~~\underset{\Theta_s}{min}H(P_t(x), P_s(x)) \)

где

\( \LARGE~~~~~~~~~~~~~~~~~~~~~~~~~~~~~H(a,b)=-a*log(b) \)

При этом градиент считается только по сети “студент”, а веса “учителя” обновляются с помощью экспоненциального скользящего среднего (EMA):

\( \LARGE~~~~~~~~~~~~~~~~~~~~~~~~~~~~\Theta_t\leftarrow \lambda \Theta_t+(1-\lambda)\Theta_s \)

Таким образом, основное отличие от классической парадигмы “дистилляции знаний” — в DINO нет предобученного “учителя”. Он обучается на основе весов “студента”, полученных на предыдущей итерации. Подход с заморозкой весов “учителя” и последующим их обновлением с помощью EMA (по-другому momentum encoder) позволил сойтись сети в процессе обучения.

Изначально momentum encoder был придуман для контрастивного обучения. Но в этом подходе не используется контрастивная функция потерь, и роль “учителя” больше похожа на концепцию Mean Teacher, где также веса “учителя” усредняются с помощью экспоненциального сглаживания. Обучение “учителя” можно представить как способ ансамблирования моделей методом усреднения весов Поляка-Руперта, когда текущее значение весов вычисляется как среднее значение весов на предыдущих шагах оптимизации. Отличие только в способе усреднения — здесь это EMA.

Как именно реализуется SSL подход в DINO: входное изображение x пропускается через набор аугментаций, и на выходе формируется 2 глобальных кропа исходного изображения и несколько локальных. Базовая конфигурация DINO следующая:

  • 2 глобальных кропа размером 224 х 224, захватывающие более 50 % изображения;
  • 6 локальных кропов 96 х 96, включающие менее 50 % изображений.

Через “студента” пропускаются все кропы, а через “учителя” — только глобальные. Таким образом мы учим “студента” обобщать информацию о всем изображении на основе набора локальных кропов.

Подробнее об архитектуре

И еще несколько слов об архитектуре в подходе DINO. В качестве сетей “учителя” и “студента” могут использоваться как ViT, так и ResNet — это так называемые бэкбоны, веса которых (полученные на этапе предобучения) мы будем использовать для наших даунстрим задач, например, классификации, сегментации, детекции и др.

Также на этапе предобучения на выходе каждого бэкбона стоит 3-х слойная полносвязная сеть — projection head, число выходов которой составляет 2048. Выходы именно этих полносвязных “голов” нормируются с помощью softmax и после этого используются в функции потерь.

Одна из проблем обучения с помощью SSL подхода — коллапс модели. Коллапс может заключаться либо в том, что все выходы модели выдают равномерное распределение, либо один из выходов всегда доминирует. Для избежания доминирования одного из выходов модели используют центрирование (centering): добавление смещения к выходу сети. Но центрирование способствует другому виду коллапса — равномерному распределению по всем выходам. Для борьбы с этим в функции softmax используют температуру. Таким образом, баланс между центрированием и температурой softmax позволяет избежать оба вида коллапса.

Оценка результатов

В таблице ниже представлены результаты сравнения метрики accuracy на датасете ImageNet сетей ResNet-50 и ViT-small в режиме обучения с учителем и в SSL режиме с линейным и KNN классификаторами на выходе. Выбор архитектур для сравнения связан с тем, что они близки по числу параметров и метрике в режиме supervised.

Рисунок 3. Таблица с результатами сравнения метрик DINO и других SSL подходов

DINO выигрывает другие SSL подходы и совсем немного проигрывает supervised моделям. При этом с трансформерным бэкбоном разница между SSL и supervised меньше, чем с ResNet.

Также стоит отметить интересный факт: метрики линейной модели и KNN довольно близки (77.0 vs 74.5), если в качестве признаков используются выходы трансформера. Это еще раз подтверждает тот факт, что DINO позволяет извлекать более разнообразные и сильные признаки именно из трансформеров, а не из классических сверточных сетей.

Свойства признаков на выходе ViT

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

Результаты ImageNet классификации с помощью KNN говорят о потенциале признаков предобученного ViT для получения векторного представления изображения с целью поиска похожих. Это подтверждают хорошие метрики на известных бенчмарках — Oxford and Paris image retrieval datasets. Также представления хорошо себя проявили на бенчмарках поиска копий изображений.

На рисунке 4 приведены примеры карт весов внимания ViT трансформера, обученного с помощью подхода DINO и просто на разметке. По сути DINO позволяет получить готовые признаки для семантической сегментации. На бенчмарке DAVIS 2017 (семантическая сегментация объектов на видео) предобученная трансформерная модель по метрике существенно превосходит результаты, полученные с помощью других SSL подходов. Она сопоставима с результатами supervised обучения.

Еще один интересный факт: карты внимания разных голов трансформера “подсвечивают” различные объекты или части объектов на изображении, как на рисунке 4.

Предобученные с помощью SSL трансформеры позволяют улучшить метрики на 1-2% на главном бенчмарке компьютерного зрения — ImageNet. Следовательно, если с нуля обучать трансформер на ImageNet, то метрики будут хуже, чем если добавлять этап предобучения с использованием DINO.

Рисунок 4. Карты весов внимания ViT трансформера, обученного с помощью DINO / на разметке

DINOv2

Через 2 года после выхода статьи про DINO вышла новая статья, где продолжилось исследование влияния self-supervised подхода на признаки трансформеров — DINOv2: Learning Robust Visual Features without Supervision.

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

Рисунок 5. Пайплайн подготовки данных в DINOv2

Разработчики DINO понимали: нельзя просто увеличивать датасет за счет неразмеченных данных без дополнительной предобработки. Эксперименты с первой версией DINO проводились на небольших подготовленных датасетах, которые были почищены от шумов и дубликатов. Поэтому в следующей версии разработчики решили сделать акцент на пайплайне подготовки данных (см. рисунок 5).

Но сначала несколько слов об источниках данных.

Под curated data понимается подготовленный, отфильтрованный датасет, а под uncurated — любые картинки из интернета. В качестве curated data использовались датасеты ImageNet-22k, Google Landmarks и еще несколько датасетов хорошего качества. В качестве uncurated data — выгруженные из интернета картинки после дедубликации с помощью PCA, фильтрации неприемлемого контента (специальная нейронная сеть), c заблюренными лицами. В результате было собрано 1.2 млрд. уникальных картинок.

После предварительной фильтрации были получены эмбеддинги всех изображений (рисунок 5), по которым провели повторную дедубликацию, отобрали наиболее похожие картинки в сравнении с датасетом ImageNet-22k. Сделано это было с помощью подсчета косинусного расстояния между эмбеддингами и кластеризацией. Таким образом, на выходе пайплайна получился датасет размером 172 млн. изображений.

Подход к обучению тоже изменился в сравнении с первой версией DINO. Для обучения 2-й версии использовали дискриминантный self-supervised подход, который представляет собой комбинацию лосс-функций DINO и iBOT. Что это значит: считается кросс-энтропия как между выходами “учителя” и “студента” (image level), так и между отдельными фрагментами изображения после предварительного маскирования (patch level — iBOT).

В результате авторам удалось существенно превзойти по метрикам все современные state of the art  self-supervised модели на датасете ImageNet-1k. И сравниться по метрикам с моделями weakly supervised (различные модификации CLIP).

Рисунок 6. Cравнение метрик weakly supervised и self-supervised моделей

Применение DINO

Уникальные свойства DINO вдохновили многих разработчиков на создание собственных решений на базе этого метода.

Одно из таких решений — STEGO (Self-supervised Transformer with Energy-based Graph Optimization). Это подход к семантической сегментации без разметки, основанный на обучении сегментационной головы с помощью специальной контрастивной лосс-функции. При этом в качестве бэкбона для получения векторного представления картинок используют трансформер, предобученный с помощью DINO подхода.

По сути STEGO — способ получения попиксельного эмбеддинга входного изображения, после кластеризации которого можно получить качественную семантическую маску.

Рисунок 7. Подход STEGO

На рисунке 7 слева показан способ вычисления функции потерь, где F — косинусная мера между эмбеддингами, полученными с помощью замороженного бэкбона (ViT), а S — косинусная мера между эмбеддингами на выходе сегментационной головы.

Обучение сводится к минимизации  скалярного произведения двух косинусных мер — F и S. Лосс функция считается для трех вариантов пар изображений: изображение и его ближайший сосед, изображение и его копия, изображение и случайная картинка из выборки. Получается, это контрастивная функция потерь, где первые 2 пары положительные, а третья пара отрицательная.

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

Рисунок 8. Сравнение STEGO и других подходов к семантической сегментации без разметки

Помимо очевидного применения DINO для задач компьютерного зрения, есть примеры успешной реализации данного подхода в другом домене — распознавании речи. В статье Leveraging in-the-wild data for effective self-supervised pretraining in speaker recognition приводится пример предобучения методом DINO на большом датасете WenetSpeech и успешного дообучения полученных весов на датасете CNCeleb.

Заключение

Self-supervised подход для предобучения ViT трансформеров позволяет добиться качества, сопоставимого с обычными сверточными сетями (которые были обучены в supervised режиме).

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

Но главный вывод после анализа результатов DINO заключается в следующем: self-supervised подход может быть ключевым фактором при создании BERT-like моделей на базе ViT трансформеров. Таким образом, DINO — это инструмент для создания больших vision моделей, которым можно скормить все картинки мира и получить базовую модель для решения большинства CV-задач.

Источники

  1. Статья Deep School про Self-supervised learning
  2. Emerging Properties in Self-Supervised Vision Transformers
  3. DINO: Emerging Properties in Self-Supervised Vision Transformers
  4. DINOv2: Learning Robust Visual Features without Supervision
  5. Unsupervised semantic segmentation by distilling feature correspondences
Старт 13 мая
Computer Vision Rocket

Приходите на наш курс CV Rocket, чтобы погрузиться в продвинутый Computer Vision: от сложностей и корнер-кейсов в «обычных» задачах до мультимодальных моделей и дизайна CV-систем!

Телеграм-канал

DeepSchool

Короткие посты по теории ML/DL, полезные
библиотеки и фреймворки, вопросы с собеседований
и советы, которые помогут в работе

Открыть Телеграм

Увидели ошибку?

Напишите нам в Telegram!