Ранее мы разобрали, что за зверь такой SSL и какой линейкой его можно померять. Также окунулись в историческое прошлое SSL методов глубокой древности (4 года назад ?) и выяснили их проблемы. В дальнейшем мы уже не будем придерживаться хронологического таймлайна, но попробуем проследить за направлениями мысли научного сообщества. В этой статье мы подробнее рассмотрим методы, строящие задачу исключительно на сравнении изображения и его аугментированной версии.
Напомню, что это третья статья из цикла про SSL в Computer Vision.

Выбор аугментаций

Шатаем собачку аугментациями

Среди основных аугментаций можно выделить 2 группы:

  1. Геометрические трансформации (повороты, кропы, скейлинг, флипы, эластичные преобразования)

  2. Цветовые трансформации (яркость, контрастность, добавление шума)

Как и при обучении в supervised манере, правильный выбор аугментаций важен для SSL. Базово аугментации для SSL выбираются более сложными, чем в supervised подходе. Подробный анализ и список принятых аугментаций приводится в статье о SimCLR.

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

AAE (Augmented autoencoders)

?M. Sundermeyer, Z. Marton, M. Durner et al. Augmented Autoencoders: Implicit 3D Orientation Learning for 6D Object Detection (Февраль 2019)
В данной работе авторы предлагают использовать автоенкодеры для восстановления исходного изображения после применения аугментации. В качестве лосса используют попиксельную L2 норму. Изначально работа описывает фреймворк для обучения на задачу определения ориентации объекта в пространстве.

Joint embedding architectures

Ниже немного поговорим про методы SSL, которые используют совместное представление эмбеддингов от 2х различных аугментаций одного и того же изображения. Они во многом опираются на идеи сиамских сетей и минимизируют расстояние между схожими изображениями или представлениями одного изображения.
Основная проблема, возникающая при такой постановке - как избежать скатывания в тривиальное решение = вне зависимости от входа всегда возвращать константу. В самом деле, если на каждое изображение сеть будет стабильно выдавать нулевой вектор, то основное условие будет удовлетворено. Собственно, последующие 3 статьи будут посвящены уловкам, которые позволяют избежать тривиального решения.

Пытаемся минимизировать расстояние между эмбеддингами похожих изображений

AMDIM (Augmented Multiscale Deep InfoMax)

?P. Bachman, D. Hjelm, W. Buchwalter. Learning Representations by Maximizing Mutual Information Across Views(Июнь 2019)
Метод заключается в максимизации совместной информации выходов с различных слоёв сети. Грубо говоря, максимизируется пропускная способность нейронной сети сохранять сигнал от входного изображения после свёрточных преобразований (это и помогает избегать тривиального решения здесь, ведь сеть должна сохранять информацию о входном изображении). Использование выходов с различных слоёв помогает проучивать сеть на всю глубину. Также максимизируется информация между аугментированными представлениями изображений, поскольку аугментация не вносит сильное искажение в семантику изображения, и эти представления должны быть близки в пространстве эмбеддингов.

Гораздо подробнее про этот метод можно почитать здесь, очень интересно! Также есть небольшой обзор методов по оптимизации совместной информации.

BYOL (Bootstrap Your Own Latent)

?J.-B. Grill, F. Strub, F. Altché et al. Bootstrap Your Own Latent A New Approach to Self-Supervised Learning(Июнь 2020)

Основная идея метода заключается в том, что эмбеддинг экстрактора можно сравнивать с эмбеддингом из соседней сети. В сеть подаются 2 различные аугментации одного изображения. При этом, обучение сети с набором параметров (верхняя на картинке) и (нижняя на картинке) происходит следующим образом:

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

В данной схеме в лосс с верхней сети подаются не сырые эмбеддинги, а прошедшие через блок predictor. Этот MLP блок наследуются из архитектуры SImCLR (будет разобрана дальше). Для дальнейшего повествования важно будет отметить наличие BatchNorm(BN) слоя в них.

Влияние BN на качество BYOL

В статье производится анализ влияния блока Batch Normailazion на обучение BYOL (можно распространить и на SimSiam далее). Эксперименты показали, что отключение BN блока приводят к деградации обучения на уровне рандомной инициализации.

Основные выводы авторов о причинах такого поведения заключаются в том, что BN исполняет роль Contrastive Loss в отсутствие явных негативных примеров. Так при отключении этого блока, модель коллапсирует эмбеддинг к одной размерности [1, 0, 0, ...]. В то время как включение BN предотращает такое поведение, поскольку изображение нормируется относительно остальных семплов в батче, которые и играют роль негативных примеров.

Said another way, with batch normalization, BYOL learns by asking, "how is this image different from the average image?". The explicit contrastive approach used by SimCLR and MoCo learns by asking, "what distinguishes these two specific images from each other?"

Этим же можно объяснить влияние размера батча на финальное качество модели.

SimSiam (Simple Siamese)

?X. Chen, K. He. Exploring Simple Siamese Representation Learning(Ноябрь 2020)
Статья развивает идею сиамских сетей и BYOL в частности, отказываясь от скользящего среднего при обновлении соседней сети (просто шарить веса). Вместо этого они предлагают запретить обновлять градиенты, которые идут не через MLP блок predictor. Казалось бы, при такой постановке задачи сеть должна скатиться к тривиальному решения - обучить MLP как тождественную функцию, а эмбеддинги с различных аугментаций подтянуть к одинаковому значению. И более того, при определённых конфигурациях обучения такое вырождение действительно происходит. Однако не во всех!

Авторы предоставляют подробные ablation study для различных конфигураций (наличие Batch Normalization слоя, размер батча, заморозка/разморозка обучения MLP) и находят конфигурации, при которых всё хорошо работает (включен BN и разморожено обучение MLP). В качестве объяснения, почему модель не скатывается в тривиальное решение, авторы выдвигают гипотезу, что такая ассимметричная архитектура близнецов напоминает Expectation–maximization (EM) алгоритм, качели которого не дают решению выродиться.
Дополнительно можно почитать здесь

DINO (self-DIstillation with NO labels)

?M. Caron, H. Touvron, I. Misra et al. Emerging Properties in Self-Supervised Vision Transformers(Март 2021)

DINO развивает подход с сиамскими сетями с обновлением одной из веток по скользящему среднему (momentum encoder) и привносит видение этого процесса как Knowledge distilation. Соотвественно, дистиляция происходит от учителя (teacher) к ученику (student) через кросс-энтропию распределений вероятностей. При этом обновление весов студента происходит через механизм бекпропа, а учителя - через скользящее среднее от весов ученика. BatchNorm авторы заменяют 2мя процедурами - центрирование и sharpening, что по сути есть тоже нормализация, только параметр сдвига в центрировании обновляется через скользящее среднее. Авторы так описывают эффект: centering prevents one dimension to dominate but encourages collapse to the uniform distribution, while the sharpening has the opposite effect. Авторы применяют мультикроп (из SwAV) и показывают, что увеличение количества локальных патчей изображений улучшает качество.
Также они используют этот метод для обучения трансформера (ViT) и обнаруживают интересный эффект - прямо из коробки сеть без какой-либо разметки начинает сегментировать изображение! При этом делает она это лучше, чем при supervised обучении.

Карта сегментации она для всех, даром. И пусть никто не уйдёт обиженным.

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

Self-classifier

?E. Amrani, L. Karlinsky, A. Bronstein Self-Supervised Classification Network(Март 2021)

Авторы рассуждают следующим образом: при построении сиамской сети с обычной кросс-энтропией решение вырождается к одному классу, но что будет, если принять гипотезу о равномерном распределении всех классов? Ответ такой - запасясь теоремой Байеса, можно сделать новый лосс и новую СОТУ. Очень простая постановка задачи, зато очень эффективная.

Barlow Twins

?J. Zbontar, L. Jing, I. Misra. Barlow Twins: Self-Supervised Learning via Redundancy Reduction(Март 2021)
Метод, как и предыдущие, предлагает учить сеть на аугментированных представлениях одного изображения, однако предлагает новый лосс. Имея батч определённого размера, эмбеддинги по каждой из размерностей образуют определённую выборку. Тогда мы можем потребовать от сети, чтобы у нас были скоррелированы выборки для каждой из размерностей и декоррелированы выборки для всех соседей. Этот принцип называется принцип сокращения избыточности, его предложил в своих работах исследователь нейронаук Horace Barlow, в честь которого и назван метод. Иными словами, для выхода из 2х сетей близнецов мы можем посчитать матрицу кросс-корреляций и приближать её к единичной матрице. Тут также используется MLP слои projector.

Поскольку для оценки лосса используется статистика по батчу, то метод чувствителен к выбору размера батча. Авторы на ablation study показывают, что оптимальный размер батча у них = 1024, при этом есть падение качества как при уменьшении, так и при увеличении. Также существенно влияет размерность эмбеддинга - увеличение ведёт к улучшению качества. Этот эффект можно объяснить тем, что такой метод обучения ведёт к ортогонализации признаков по размерностям эмбеддинга, следовательно смешанные признаки расходятся каждый на свой базис. Также метод показывает устойчивость при отсутствии BN слоя. Возможно, его роль исполняет кросс-корреляция, которая рассчитывается по всему батчу.

VICReg (Variance-Invariance-Covariance Regularization)

?A. Bardes, J. Ponce, Y. LeCun. VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning(Май 2021)

Метод предлагает оптимизировать эмбеддинги сразу по нескольким информационным критериям:

  • Invariance - критерий схожести 2х эмбеддингов аументированных представлений изображения

  • Variance - используется hinge loss для дисперсии каждого эмбеддинга изображения. Сеть не штрафуется, если дисперсия эмбеддинга превышает некоторый порог, т.е. он содержит информацию значимо отличающуюся от нулевого вектора.

  • Covariance - аналогично Barlow Twins, уменьшает ковариацию между различными компонентами векторов-эмбеддингов.

Отличителной особенностью метода является то, что 2 сети не обязаны быть одинаковыми, а веса не обязаны быть общими. Кроме того, авторы распространили VICReg лосс на BYOL и SimSiam и показали, что использование этого лосса улучшает метрики и убирает необходимость вводить Batch Normalization в сеть.

Вместо итога

Мы упомянули основные аугментации, применяемые в SSL, обучение автоенкодера и подробно рассмотрели методы, которые используют изображения и их аугментации для обучения. Большинство этих методов основано на сиамских сетях и идее, что аугментированные представления одного и того же изображения должны быть близки в пространстве эмбеддингов. Однако одного критерия близости недостаточно, чтобы избежать коллапса сети в тривиальное решение, поэтому были предложены дополнительные информационные критерии (AMDIM, Barlow Twins, VICReg), либо дополнительные трюки при обучении (BYOL, SimSiam).
Если в предыдущей статье мы рассматривали качество методов в конце статьи, то здесь мы сделаем ещё хитрее - мы рассмотрим результаты в последней статье цикла после рассмотрения ещё пачки методов. В качестве небольшого спойлера - они примерно отсортированы в порядке улучшения метрик. Если вам больше интересны практические апекты, то предлагаю сразу перейти к последней статье цикла.
В следующей статье мы разберём методы на основе contrastive подхода - как получать выгоду не только из похожести, но и из различия!

Список статей в цикле

  1. SSL. Проблематика и постановка задачи

  2. SSL. Метрики и первые pretext tasks

  3. SSL. Обучение на изображении и его аугментациях

  4. SSL. Contrastive learning

  5. SSL. Кластеризация как лосс

  6. SSL. Результаты и основные фреймворки

Дополнительные ссылки

  1. Whitening-MSE - ?A. Ermolov, A. Siarohin, E. Sangineto et al. Whitening for self-supervised representation learning(Июль 2020)

  2. CorInfoMax - ?S. Ozsoy, S, Hamdan, S. Ö. Arik et al. Self-Supervised Learning with an Information Maximization Criterion(Сентябрь 2022)

  3. RandAugment - ?E. D. Cubuk, B. Zoph, J. Shlens et al. RandAugment: Practical automated data augmentation with a reduced search space(Сентябрь 2019)

  4. OBoW - ?S. Gidaris, A. Bursuc, G. Puy et al. OBoW: Online Bag-of-Visual-Words Generation for Self-Supervised Learning(Декабрь 2020)

  5. Decorr - ?T. Hua, W. Wang, Z. Xue et al. On Feature Decorrelation in Self-Supervised Learning(Май 2021)

  6. TWIST - ?F. Wang, T. Kong, R. Zhang et al. Self-Supervised Learning by Estimating Twin Class Distributions(Октябрь 2021)