VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning
Adrien Bardes, Jean Ponce, Yann LeCun
Статья: https://arxiv.org/abs/2105.04906
Неофициальный код: https://github.com/vturrisi/solo-learn (с разными методами SSL, включая VICReg)
Библиотека для SSL от FAIR (где можно ожидать появления VICReg, но пока нет): https://github.com/facebookresearch/vissl + https://vissl.ai/
Май 2021, продолжение развития Barlow Twins от Яна Лекуна и ко.
Проблема всё та же — хочется избежать коллапса репрезентаций, но желательно не разными [работающими, заметим] хаками и имплементациями, а каким-нибудь красивым методом, который был бы обоснован и интерпретируем.
Архитектура сиамская симметричная, обе ветви одинаковые и с расшаренными весами, ResNet-50 энкодер + MLP проектор, в каждую ветвь подаётся своя аугментация входного изображения. Вся инновация снова в лосс-функции.
В данной работе предлагается новая трёхкомпонентная целевая функция под названием VICReg (Variance-Invariance-Covariance Regularization), самое новое из которого — это Variance.
Разберём по частям.
Invariance regularization — в целом общее место во многих методах, отвечающее за одинаковость репрезентаций аугментаций одного и того же объекта. Здесь реализуется через mean-squared euclidean distance.
Covariance regularization — идея аналогичная Barlow Twins для того, чтобы получить независимые фичи, декоррелируя их. Реализована в виде суммы квадратов внедиагональных элементов матрицы ковариации, которые данный лосс стремится занулить. Но разница в том, что в BT была кросс-корреляция векторов из двух ветвей сиамской сети, а здесь это ковариация между фичами внутри каждой ветви сети, делающая фичи независимыми друг от друга локально в пределах ветви. Заодно избавились от стандартизации, которая была в BT (и приводила к численной нестабильности, когда эмбеддинг уезжал в константу), здесь вместо неё работает Variance regularization.
Variance regularization — метод регуляризации дисперсии фич, защищает от коллапса. Технически это hinge loss на стандартное отклонение фич проекции по измерению батча. В лоссе фигурирует параметр гамма, который является целевым для стандартного отклонения по каждой фиче, в работе был установлен в 1.
То есть по сути в работе декомпозировали различные цели SSL в отдельные компоненты лосса. Они замешиваются в итоговый лосс с гиперпараметрами, дающими вес каждому компоненту.
Из новых свойств метода то, что он (вроде как первый) не требует какой-либо отдельной нормализации фич.
Эксперименты снова на ImageNet в режимах linear classification (обучают линейный классификатор поверх эмбеддингов полученных в self-supervised режиме) и semi-supervised (когда дополнительно обучают с учителем на 1% или 10% разметки). Метод примерно в одной группе с BYOL, SwAV и Barlow Twins. Аналогично на других картиночных датасетах и задаче детекции объектов.
Есть довольно подробные абляции чтобы понять что именно и как работает. Заодно показывают, что регуляризацию дисперсии можно использовать и с другими архитектурами (например, BYOL или SimSiam) и это даёт некоторое улучшение.
Метод снова не особо чувствителен к размеру батча и нормально работает на небольших батчах.
Потенциальный недостаток метода — необходимость вычисления матрицы ковариации, а это квадратичная сложность по размеру фич (но всё равно не настолько ужасно как в методах, требующих вычисления обратных матриц, что используется в методе W-MSE на основе whitening). Собираются в будущем поискать более эффективные решения, не требующие вычисления полной матрицы.
Adrien Bardes, Jean Ponce, Yann LeCun
Статья: https://arxiv.org/abs/2105.04906
Неофициальный код: https://github.com/vturrisi/solo-learn (с разными методами SSL, включая VICReg)
Библиотека для SSL от FAIR (где можно ожидать появления VICReg, но пока нет): https://github.com/facebookresearch/vissl + https://vissl.ai/
Май 2021, продолжение развития Barlow Twins от Яна Лекуна и ко.
Проблема всё та же — хочется избежать коллапса репрезентаций, но желательно не разными [работающими, заметим] хаками и имплементациями, а каким-нибудь красивым методом, который был бы обоснован и интерпретируем.
Архитектура сиамская симметричная, обе ветви одинаковые и с расшаренными весами, ResNet-50 энкодер + MLP проектор, в каждую ветвь подаётся своя аугментация входного изображения. Вся инновация снова в лосс-функции.
В данной работе предлагается новая трёхкомпонентная целевая функция под названием VICReg (Variance-Invariance-Covariance Regularization), самое новое из которого — это Variance.
Разберём по частям.
Invariance regularization — в целом общее место во многих методах, отвечающее за одинаковость репрезентаций аугментаций одного и того же объекта. Здесь реализуется через mean-squared euclidean distance.
Covariance regularization — идея аналогичная Barlow Twins для того, чтобы получить независимые фичи, декоррелируя их. Реализована в виде суммы квадратов внедиагональных элементов матрицы ковариации, которые данный лосс стремится занулить. Но разница в том, что в BT была кросс-корреляция векторов из двух ветвей сиамской сети, а здесь это ковариация между фичами внутри каждой ветви сети, делающая фичи независимыми друг от друга локально в пределах ветви. Заодно избавились от стандартизации, которая была в BT (и приводила к численной нестабильности, когда эмбеддинг уезжал в константу), здесь вместо неё работает Variance regularization.
Variance regularization — метод регуляризации дисперсии фич, защищает от коллапса. Технически это hinge loss на стандартное отклонение фич проекции по измерению батча. В лоссе фигурирует параметр гамма, который является целевым для стандартного отклонения по каждой фиче, в работе был установлен в 1.
То есть по сути в работе декомпозировали различные цели SSL в отдельные компоненты лосса. Они замешиваются в итоговый лосс с гиперпараметрами, дающими вес каждому компоненту.
Из новых свойств метода то, что он (вроде как первый) не требует какой-либо отдельной нормализации фич.
Эксперименты снова на ImageNet в режимах linear classification (обучают линейный классификатор поверх эмбеддингов полученных в self-supervised режиме) и semi-supervised (когда дополнительно обучают с учителем на 1% или 10% разметки). Метод примерно в одной группе с BYOL, SwAV и Barlow Twins. Аналогично на других картиночных датасетах и задаче детекции объектов.
Есть довольно подробные абляции чтобы понять что именно и как работает. Заодно показывают, что регуляризацию дисперсии можно использовать и с другими архитектурами (например, BYOL или SimSiam) и это даёт некоторое улучшение.
Метод снова не особо чувствителен к размеру батча и нормально работает на небольших батчах.
Потенциальный недостаток метода — необходимость вычисления матрицы ковариации, а это квадратичная сложность по размеру фич (но всё равно не настолько ужасно как в методах, требующих вычисления обратных матриц, что используется в методе W-MSE на основе whitening). Собираются в будущем поискать более эффективные решения, не требующие вычисления полной матрицы.
GitHub
GitHub - vturrisi/solo-learn: solo-learn: a library of self-supervised methods for visual representation learning powered by Pytorch…
solo-learn: a library of self-supervised methods for visual representation learning powered by Pytorch Lightning - vturrisi/solo-learn
Not All Memories are Created Equal: Learning to Forget by Expiring
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, Angela Fan
Статья: https://arxiv.org/abs/2105.06548
Code: https://github.com/facebookresearch/transformer-sequential (плюс там же код для недавно обсуждавшегося Feedback Transformer)
Прикольная работа от хороших авторов в тему трансформеров с памятью.
Много свежей движухи в стане трансформеров направлено на увеличение attention span (все эти Big Bird, Longformer, Reformer и т.п.) и/или добавление памяти (Compressive Transformer и другая экзотика). Но, как известно, не все йогурты одинаково полезны. Мало увеличивать span и добавлять больше памяти, надо уметь и забывать ненужное и освобождать память.
Данная работа посвящена именно этому и идейно во многом является продолжением линий Adaptive Attention Span (тех же авторов, https://news.1rj.ru/str/gonzo_ML/99) и Compressive Transformer (DeepMind, https://news.1rj.ru/str/gonzo_ML/165).
Adaptive Attention Span позволял адаптивно определять этот самый span для каждой головы внимания независимо и не считать внимание по большим окнам, где это не надо. Compressive transformer перегонял в долгую память репрезентации, уходящие из обычного поля зрения, но делал это довольно топорно.
Суть текущей работы, Expire-Span, в том, что давайте считать внимание не по всем элементам какого-то окна (attention span), а только по тем, которые “не забылись”. Тогда и окно можно сделать сильно больше, если в нём не так много ключевой информации, которую надо помнить. В данной работе attention span доходит до 128К токенов, что, кажется, рекорд. Фокусируются на декодерах.
Технически это работает так. Для каждой репрезентации токена (hidden state) мы имеем предиктор (по сути однослойная сетка с сигмоидой), дающий значение expiration value — как долго (сколько шагов по времени) эту репрезентацию нужно помнить. Как только элемент выходит из этого окошка, он “забывается” и не участвует далее в вычислении весов внимания. Количество элементов, участвующих в этом вычислении, называется размером памяти (C).
Окошко памяти сделано по аналогии с работой про adaptive attention span без резкого перепада, а как soft mask, в нём есть участок длины R (гиперпараметр), где вес элемента линейно спадает от единицы до нуля. Это чтобы был нормальный участок, откуда бы поступали градиенты на обучение этого механизма.
Лосс функция содержит дополнительный терм, уменьшающий средний размер используемой памяти C.
Механизм работает независимо для каждого слоя, что потенциально даёт возможность каждому слою специализироваться на своём временном масштабе (а можно, кстати, и заэнфорсить явный bias в этом месте, интересно, что получилось бы).
В работе указано много тонкостей имплементации. Среди них, например, есть абзац про регуляризацию памяти, когда при обучении для каждого батча случайным образом определяется число l<L (максимальный span) и память дальше этого l зануляется. Это чтобы модель не рассчитывала, что память всегда будет содержать специфическую информацию. Есть чувство, что это можно было бы назвать memory dropout и посвятить ему отдельную статью. А тут скромно абзац внутри текста.
Для сравнения бейзлайнами взяты Transformer-XL, Adaptive-Span и Compressive Transformer. Expire-Span демонстрирует очень хорошие результаты.
На задачке с коридором (когда в начале коридора есть цветное пятно, говорящее, в каком из поворотов в конце коридора будет награда, а коридор может быть длинным, так что помнить надо долго) ощутимо бьёт XL даже при малом объёме памяти (что логично, для такой задачи её как раз нужно мало). Есть задача навигации по нескольким комнатам с порталами (надо помнить, какой портал где правильный) и задача про длинные текстовые инструкции для игры, которые надо запоминать и в нужном месте использовать. Тут тоже лучше, чём XL.
Есть задача копирования последовательности, когда идёт сколько-то символов A, потом много B, и потом надо скопировать то самое количество A. Чем больше span, тем лучше работает (до 128К).
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, Angela Fan
Статья: https://arxiv.org/abs/2105.06548
Code: https://github.com/facebookresearch/transformer-sequential (плюс там же код для недавно обсуждавшегося Feedback Transformer)
Прикольная работа от хороших авторов в тему трансформеров с памятью.
Много свежей движухи в стане трансформеров направлено на увеличение attention span (все эти Big Bird, Longformer, Reformer и т.п.) и/или добавление памяти (Compressive Transformer и другая экзотика). Но, как известно, не все йогурты одинаково полезны. Мало увеличивать span и добавлять больше памяти, надо уметь и забывать ненужное и освобождать память.
Данная работа посвящена именно этому и идейно во многом является продолжением линий Adaptive Attention Span (тех же авторов, https://news.1rj.ru/str/gonzo_ML/99) и Compressive Transformer (DeepMind, https://news.1rj.ru/str/gonzo_ML/165).
Adaptive Attention Span позволял адаптивно определять этот самый span для каждой головы внимания независимо и не считать внимание по большим окнам, где это не надо. Compressive transformer перегонял в долгую память репрезентации, уходящие из обычного поля зрения, но делал это довольно топорно.
Суть текущей работы, Expire-Span, в том, что давайте считать внимание не по всем элементам какого-то окна (attention span), а только по тем, которые “не забылись”. Тогда и окно можно сделать сильно больше, если в нём не так много ключевой информации, которую надо помнить. В данной работе attention span доходит до 128К токенов, что, кажется, рекорд. Фокусируются на декодерах.
Технически это работает так. Для каждой репрезентации токена (hidden state) мы имеем предиктор (по сути однослойная сетка с сигмоидой), дающий значение expiration value — как долго (сколько шагов по времени) эту репрезентацию нужно помнить. Как только элемент выходит из этого окошка, он “забывается” и не участвует далее в вычислении весов внимания. Количество элементов, участвующих в этом вычислении, называется размером памяти (C).
Окошко памяти сделано по аналогии с работой про adaptive attention span без резкого перепада, а как soft mask, в нём есть участок длины R (гиперпараметр), где вес элемента линейно спадает от единицы до нуля. Это чтобы был нормальный участок, откуда бы поступали градиенты на обучение этого механизма.
Лосс функция содержит дополнительный терм, уменьшающий средний размер используемой памяти C.
Механизм работает независимо для каждого слоя, что потенциально даёт возможность каждому слою специализироваться на своём временном масштабе (а можно, кстати, и заэнфорсить явный bias в этом месте, интересно, что получилось бы).
В работе указано много тонкостей имплементации. Среди них, например, есть абзац про регуляризацию памяти, когда при обучении для каждого батча случайным образом определяется число l<L (максимальный span) и память дальше этого l зануляется. Это чтобы модель не рассчитывала, что память всегда будет содержать специфическую информацию. Есть чувство, что это можно было бы назвать memory dropout и посвятить ему отдельную статью. А тут скромно абзац внутри текста.
Для сравнения бейзлайнами взяты Transformer-XL, Adaptive-Span и Compressive Transformer. Expire-Span демонстрирует очень хорошие результаты.
На задачке с коридором (когда в начале коридора есть цветное пятно, говорящее, в каком из поворотов в конце коридора будет награда, а коридор может быть длинным, так что помнить надо долго) ощутимо бьёт XL даже при малом объёме памяти (что логично, для такой задачи её как раз нужно мало). Есть задача навигации по нескольким комнатам с порталами (надо помнить, какой портал где правильный) и задача про длинные текстовые инструкции для игры, которые надо запоминать и в нужном месте использовать. Тут тоже лучше, чём XL.
Есть задача копирования последовательности, когда идёт сколько-то символов A, потом много B, и потом надо скопировать то самое количество A. Чем больше span, тем лучше работает (до 128К).
GitHub
GitHub - facebookresearch/transformer-sequential: Trains Transformer model variants. Data isn't shuffled between batches.
Trains Transformer model variants. Data isn't shuffled between batches. - facebookresearch/transformer-sequential
Языковое моделирование на Enwik8 дало новую SoTA, побив в том числе и Compressive Transformer.
На PG-19 качество сравнимое с Compressive Transformer, но существенно быстрее.
Есть и другие задачи, где модель ведёт себя достойно.
Что ценно, кроме высокого качества, модель ест меньше памяти, чем Adaptive-Span или Compressive Transformer, а также стабильно быстрее их на обучении, иногда в 2-3 раза.
Есть интересный анализ expiration value в зависимости от конкретных позиций в разных текстах. Показывают, что разные слои модели специализируются на разных объектах: какие-то хранят именованные сущности, другие отмечают пробелы, ещё одни смотрят на заголовки секций в документах.
Мне кажется, и adaptive span раньше был хороший и незаслуженно малоизвестный, а эта штука бьёт его везде. Плюс, она вполне ортогональна многим другим улучшениям трансформеров. Было бы интересно собрать новый GPT с этим механизмом плюс ещё несколькими многообещающими. И также было бы интересно применить его к окологеномным задачам, там длинный контекст часто важен. 128К это прям уже хорошо!
На PG-19 качество сравнимое с Compressive Transformer, но существенно быстрее.
Есть и другие задачи, где модель ведёт себя достойно.
Что ценно, кроме высокого качества, модель ест меньше памяти, чем Adaptive-Span или Compressive Transformer, а также стабильно быстрее их на обучении, иногда в 2-3 раза.
Есть интересный анализ expiration value в зависимости от конкретных позиций в разных текстах. Показывают, что разные слои модели специализируются на разных объектах: какие-то хранят именованные сущности, другие отмечают пробелы, ещё одни смотрят на заголовки секций в документах.
Мне кажется, и adaptive span раньше был хороший и незаслуженно малоизвестный, а эта штука бьёт его везде. Плюс, она вполне ортогональна многим другим улучшениям трансформеров. Было бы интересно собрать новый GPT с этим механизмом плюс ещё несколькими многообещающими. И также было бы интересно применить его к окологеномным задачам, там длинный контекст часто важен. 128К это прям уже хорошо!