gonzo-обзоры ML статей
16.7K subscribers
1.85K photos
3 files
775 links
Авторы:
Гриша Сапунов, ранее руководитель разработки Яндекс-Новостей, ныне CTO Intento. Области интересов: AI/ML/DL, биоинформатика.
Лёша Тихонов, ранее аналитик в Яндексе, автор Автопоэта, Нейронной Обороны... Области интересов: discrete domain, NLP, RL.
Download Telegram
Обзор интересной статьи про применение contrastive learning для биологии
Forwarded from DL in biology (Зоя Червонцева)
Evolution Is All You Need: Phylogenetic Augmentation for Contrastive Learning

Текст:
https://arxiv.org/abs/2012.13475
Код: нет, но есть формулы

Статья-видение. Авторы предлагают новую разновидность contrastive learning для обучения именно биологических моделей.

Contrastive learning -- одна из популярных концепций самообучения (self-supervised learning). Самообучение часто бывает полезно, когда данных много, а меток для них мало. Можно предобучиться на неразмеченных данных и что-то понять про их структуру -- а потом уже приступать к классификации. Для стадии предобучения данные часто размножают (аугментируют), и в этой работе авторы заметили, что биологические данные и так уже аугментированы эволюцией, причем очень подходящим для contrastive learning образом.

Собственно идея контрастирования заключается в следующем. Представим, что у нас есть картинки, и для каждой картинки определены два разных аугментирующих преобразования -- например, поворот на 90 градусов и размытие. Цель -- научиться понимать, какие картинки являются вариантами друг друга. Для этого нужно найти такую функцию (энкодер), которая отобразит два варианта одной исходной картинки в максимально близкие точки, а варианты разных картинок -- в далёкие. Будем требовать, чтобы это выполнялось для всех картинок, а степень близости итоговых точек определим через взаимную информацию их распределений. Если хорошо подобрать преобразования, такой энкодер научится отделять важное от неважного.

На практике вместо взаимной информации считают более простую функцию, являющуюся её нижней оценкой. Оптимизируется функция для N пар производных картинок, из которых одна пара “правильная” -- соответствует двум разным преобразованиям одной исходной картинки, а остальные пары “неправильные” -- соответствуют разным картинкам. Преобразования при этом не фиксируют, а сэмплируют каждый раз из некоторого заранее заданного распределения.

Так вот, биология. Авторы предлагают перейти к работе с последовательностями и рассматривать их в контексте ортологических рядов. Исходная “картинка” каждого ряда -- это его предковая последовательность, а все современные последовательности -- преобразования исходной картинки. Теперь можно применить всю ту же процедуру обучения энкодера. Хорошими парами будут считаться пары ортологов, а плохими -- пары из разных семейств. Из некоторого математически мотивированного принципа следует, что хорошие пары лучше составлять из максимально далеких друг от друга последовательностей.

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

Сегодня 23 января, в 21:00 MSK (через час) планируем обсудить широко известную статью Франсуа Шолле "On the measure of intelligence" (https://arxiv.org/abs/1911.01547)

"Пойдусь по статье, расскажу зачем нам нужно формальное измеримое определение интеллекта, какие к нему есть подходы, какие будут более продуктивными, какие менее. Расскажу, какое определение предлагает сам Франсуа. И наконец, расскажу про предложенный им бенчмарк ARC, в чем его достоинства и почему им стоит заниматься уже сейчас, попробую дать краткий обзор того, что с ним уже произошло за год" @Andreyz4k

Если вы ещё не с нами, но вам интересно участвовать в этой активности в качестве слушателя или докладчика (или ещё как), ответьте, пожалуйста, на вопросы анкеты:

https://forms.gle/6dx54vmruiSVAv9A8

Обещаю сегодня перед семинаром отправить приглашения всем, кто их оставит за это время (до 20:50 примерно).
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
William Fedus, Barret Zoph, Noam Shazeer
Статья: https://arxiv.org/abs/2101.03961

В начале 2021 года Гугл опубликовал статью Switch Transformers, которая в медиа прошла под вывеской “Гугл обучил трансформер с 1.6T параметров и побил GPT-3 со своими 175B параметров!”.

Это в целом не совсем корректный и полный месседж хотя бы, например, потому, что ещё полгода назад у Гугла уже была модель на 600B параметров в работе про GShard с примерно такой же архитектурой, а также потому, что сравнивать архитектуру стандартного трансформера (как в GPT-3) с архитектурой MoE (лежащей в основе Switch Transformer или GShard) по общему числу параметров не очень корректно.

Работа и направление на самом деле интересные. В чём суть?

Начать стоит издалека.

Архитектура под названием Mixture of Experts (MoE) в нейросетях является одним из направлений в рамках подхода conditional computation, суть которого в том, что граф вычислений определяется входными данными. В целом сам подход MoE достаточно старый, прослеживается как минимум с начала 1990-х годов (http://www.cs.toronto.edu/~fritz/absps/jjnh91.pdf) и применялся с разного рода моделями, далеко не только нейросетевыми.

В нейросетевых MoE обычно есть множество “экспертов” (отдельных слоёв или даже целых сетей), в которые направляет данные на обработку специальный механизм роутинга (gating network). Этот механизм по входным данным генерирует распределение вероятностей “подходящих” экспертов (и по сути является обучаемым полносвязным слоем с софтмаксом), из которых обычно выбирается сколько-то (немного) топовых экспертов, и в них отправляются данные. Выходы экспертов суммируются и направляются затем в последующий слой. Остальные невыбранные эксперты ничего не получают и в этот момент не работают. Такой вариант, реализованный в виде слоя, называется Sparsely-Gated Mixture-of-Experts Layer, здесь разреженность относится именно к выбору экспертов.

В начале 2017 года весьма представительный коллектив авторов (включая Хинтона, Джеффа Дина, а также Noam Shazeer, одного из авторов работ про более поздние Switch Transformer и GShard) опубликовал знаковую работу про этот метод (“Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”, https://arxiv.org/abs/1701.06538), где авторы применили слой MoE между двумя слоями рекурретной сети для задачи языкового моделирования (в целом та же задача, которую решает GPT), а также для машинного перевода.

Надо сказать, уже в той работе января 2017 года появилась модель на 137B параметров и это было почти за полтора года до первой GPT (хронологию GPT можно посмотреть здесь: https://blog.inten.to/gpt-3-language-models-are-few-shot-learners-a13d1ae8b1f9).

Данная модель содержала 131072 эксперта (эксперт представлял собой простую feed-forward сеть с одним скрытым слоем), из которых активными в каждый момент времени были только 4. То есть из этих 137B параметров (относящихся, правда, не только к экспертам, но зато не относящихся к большим матрицам эмбеддингов, которые из этого подсчёта исключены) в реальности в каждый момент времени работает очень небольшое (на порядки меньшее) число. За счёт огромной разреженности по экспертам, такие модели требуют заметно меньше вычислений, чем dense модели с равным числом параметров.

Так что для более адекватного сравнения с “обычными” трансформерами было бы правильно ввести какую-то другую метрику, что-то типа “количество активных весов” или как-то ещё и сравниваться по ней. У обычных трансформеров, кстати, это число тоже может быть сильно меньше полного числа параметров (но всё же не настолько сильно, как у MoE), если они используют функции активации типа ReLU, в таких сетях часть “нейронов” по факту может быть выключена.

Та модель 2017 года показала прекрасные результаты и побила state-of-the-art на языковом моделировании, а также вела себя весьма достойно на машинном переводе.
Другим большим достижением я бы назвал работу про GShard середины 2020 года (“GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding”, https://arxiv.org/abs/2006.16668). Среди авторов Дмитрий Лепихин (пользуясь случаем, хочу передать привет :) ) и другие, включая всё того же Noam Shazeer.

В этой работе MoE применили к трансформерам, в которых каждый второй feed-forward слой в энкодере и декодере заменили на MoE с top-2 активными в каждый момент времени экспертами. Модель отскейлили до 600B параметров (2048 экспертов, 36 слоёв трансформера) и обучали на 2048 TPU v3 под задачу мультиязычного машинного перевода, и получили заметное улучшение качества.

Чтобы трансформер с MoE можно было так отскейлить, разработали модуль GShard, предоставляющий API для шардирования моделей, по сути представляющий собой аннотации для тензоров.

В работе даже попробовали обучить модель на 1T параметров (2048 экспертов, 60 слоёв), но с ней были проблемы по части стабильности при обучении на bfloat16 (остальное обучалось на float32; напоминаю, что про форматы чисел с плавающей точкой есть подробная статья https://moocaholic.medium.com/fp64-fp32-fp16-bfloat16-tf32-and-other-members-of-the-zoo-a1ca7897d407).

Ну и собственно Switch Transformer во многом продолжает эту линию. Алгоритм роутинга в MoE упростили, теперь активен только один эксперт всё в том же feed-forward слое (был отдельный эксперимент по добавлению MoE в слой внимания, и это даёт лучший результат, но работает на float32, а на bfloat16 расходится). Каждый токен отправляется в соответствующего ему эксперта.

Делали это поверх Mesh-Tensorflow, в работе есть сколько-то оптимизаций для улучшения обучения (например, везде используют bfloat16, что ощутимо облегчает коммуникацию, и только внутри функции роутера всё приводят к float32; более хитрые инициализации и регуляризации).

Сравниваются с моделями семейста T5: T5-base (223M), T5-large (739M), для сравнения с ними делают Switch-base (7.4B) и Switch-large (26.3B) так, чтобы по количеству вычислений (FLOPS) они совпадали.

Делают предобучение на Masked Language Model (MLM, по типу BERT’а) и файн-тюнинг на задачах из GLUE, SQuAD, SuperGLUE и т.д.

Был ещё дополнительный лосс, помогающий балансировать нагрузку по различным экспертам, чтобы никакой из них не перетрудился и не стал узким местом, к кому стоит очередь. Исключать, кстати, такого всё равно нельзя и на случай, когда кто-то окажется перегружен и не успеет обработать какой-то из входящих токенов, этот токен (вернее его эмбеддинг) пробрасывается на выход, благо там везде residual connections. Было предположение, что качество от этого страдает, поэтому сделали No-Token-Left-Behind роутинг, который, если обнаруживает, что токен отправлен на перегруженного эксперта, рероутит его на другого. Но это на практике ничего не дало.

Новых state-of-the-art здесь не получают, но зато показывают, что, во-первых, Switch Transformer превосходит обычный трансформер по качеству; во-вторых, он хорошо скейлится по числу экспертов (чем больше, тем выше качество; при этом не забывайте, что активен в каждый момент только один эксперт, так что это не ансамбль в традиционном виде), а также, в-третьих, при фиксированном объёме вычислений достигает лучшего результата, чем традиционный трансформер (в 7 раз быстрее для base варианта).

Большой Switch Transformer можно отдистиллировать (про дистилляцию у нас в канале вообще было много всего) в обычный, например, Switch-base в T5-base, и это даст качество выше, чем если с нуля обучать аналогичный обычный (T5-base), сохраняется примерно 30% улучшения. Сравнились также с мультиязычным mT5, получили прирост на всех языках (про T5 и mT5 мы писали https://t.me/gonzo_ML/442).
По сути Switch Transformer добавляет новое измерение для распараллеливания. К традиционным data- (когда разные части датасета обрабатываются копиями модели на разных узлах) и model-parallelism (когда части одной модели вычисляются на разных узлах) добавляется expert-parallelism (когда разные эксперты вычисляются на разных узлах). И, конечно, всё это можно комбинировать. И даже нужно, если модель очень большая.

И вот, собственно, очень большая модель возникает. Авторы обучают модели Switch-XXL (395B параметров, 64 эксперта), аналогичную по количеству вычислений T5-XXL (13B, 8.7T FLOPS/sequence), а также Switch-C (1571B, 2048 экспертов, более лёгкая для вычислений, “всего” 890B FLOPS/sequence) — вот она, пресловутая модель на 1.6T параметров. Switch-C, кстати, обучали только с использованием expert-parallelism.

Из ожидаемого, Switch-XXL обошёл по качеству T5-XXL. Из менее ожидаемого, Switch-C даёт качество ниже, чем Switch-XXL за то же число шагов, но поскольку по вычислениям он в разы легче (примерно 10x), то при одинаковом вычислительном бюджете он добирается до фиксированного качества раньше, чем Switch-XXL. Также из неожиданного, что Switch-C нестабильности при обучении не демонстрирует, а Switch-XXL иногда нестабилен. Ещё из неожиданного есть то, что Switch-XXL получает state-of-the-art на предобучении, но не получает его на задачах после файн-тюнинга на SuperGLUE (но в целом положительная корреляция есть) — эта динамика требует отдельного изучения.

Практический вопрос: нужно ли всё это, если у вас нет суперкомпьютера с кучей TPU/GPU? Авторы показали, что улучшение есть даже на двух экспертах, а это можно использовать и с маленькими моделями, влезающими на скромные машины.

В общем интересное направление, можно ожидать перевода в практическую плоскость даже для “простых смертных”.
MoE из статьи 2017 года с LSTM
MoE Transformer из статьи про GShard
Switch Transformer
Скейлинг по количеству экспертов. Даже 2 эксперта это хорошо.
Сравнение Switch Transformer и MoE Transformer
Switch быстрее