Barlow Twins: Self-Supervised Learning via Redundancy Reduction
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, Stéphane Deny
Статья:
https://arxiv.org/abs/2103.03230 Код (PyTorch):
https://github.com/facebookresearch/barlowtwins Мартовская работа от команды FB, включая Яна Лекуна, сильно перекликающаяся с BYOL. Тоже про self-supervised learning (SSL) без негативных пар.
Общий фреймворк аналогичен — мы ищем репрезентации, которые инвариантны относительно различных искажений (аугментаций) входного объекта. Проблема здесь, как уже упоминалось в посте про BYOL, это константные вырожденные репрезентации, и различные современные методы обходят это по-разному, часто это какая-то особенность имплементации (асимметричность архитектуры ветвей или же обновления весов), не дающая туда скатиться, но такое решение по-прежнему остаётся валидным решением задачи минимизации заданного лосса. В текущей работе авторы предлагают новый лосс, который by design устраняет такие решения, так что всякие хитрые уловки больше не нужны, и архитектуру можно сделать симметричную.
Предложенный метод вдохновлён нейронаукой, а точнее работой Барлоу (
https://en.wikipedia.org/wiki/Horace_Barlow) 1961 года под названием “Possible Principles Underlying the Transformations of Sensory Messages” (
https://www.cnbc.cmu.edu/~tai/microns_papers/Barlow-SensoryCommunication-1961.pdf), где он предложил гипотезу о кодировании сенсорной информации в нервной системе. В соответствии с гипотезой, цель обработки сенсорной информации — избавиться от избыточности и перекодировать сигнал в factorial code (факториальный код?
https://en.wikipedia.org/wiki/Factorial_code) — код, компоненты которого статистически независимы.
Отсюда следует основная инновация авторов текущей работы, их хитрая функция потерь. Она работает поверх матрицы кросс-корреляции для выходов двух идентичных сетей, обрабатывающих каждая свою аугментацию входного объекта.
Функция потерь в свою очередь состоит из двух компонентов, первый, т.н. invariance term, отвечает за близость диагонали матрицы к единицам, что даёт инвариантность к искажениям входа. Второй компонент, redundancy reduction term, представляет собой сумму квадратов недиагональных элементов и стремится их все занулить, тем самым декоррелируя компоненты вектора. Авторы также предлагают взгляд на эту функцию через призму теории информации и Information Bottleneck принципа Нафтали Тишби.
Двухкомпонентная функция на первый взгляд напоминает двухкомпонентный же лосс в InfoNCE (стандартный лосс в контрастивных методах, где первая компонента про позитивные пары, а вторая про негативные). Но есть концептуальные различия. InfoNCE максимизирует изменчивость репрезентаций, максимизируя попарные расстояния между разными сэмплами, а BT достигает этого декоррелированием компонент вектора репрезентации. То есть BT по факту не зависит от большого количества примеров для негативных пар и нормально работает на небольших батчах (как вы помните, негативные пары набираются внутри одного батча). В этом поведении BT похож на BYOL. Дополнительно BT может работать с репрезентациями очень большой размерности, что часто плохо умеют другие методы.
Вся остальная имплементация в целом очень похожа на конкурирующие методы. Всё так же есть две сети, но теперь они полностью идентичны и бэкпроп идёт через обе. Энкодер это снова ResNet-50, сеть проектор — это трёхслойный MLP с 8192 юнитами (это больше, чем обычно). Аугментации и сам процесс оптимизации аналогичен BYOL.
Всё это добро обучалось на ImageNet на 32 V100 в течение 124 часов (батч 2018). На той же конфигурации BYOL обучался 113 часов (батч 4096).