На классификации картинок ImageNet и обычный Perceiver работал сносно, но здесь применяется более общий декодер и Perceiver IO стабильно лучше.

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

optical flow estimation

Pretrained Transformers as Universal Computation Engines
Kevin Lu, Aditya Grover, Pieter Abbeel, Igor Mordatch
Статья: https://arxiv.org/abs/2103.05247
Код: https://github.com/kzl/universal-computation

Прикольная работа весны этого года из Беркли/Google/FB, про которую нелишне ещё раз поговорить.

Общая идея в том, что обычные предобученные трансформеры для NLP (GPT-2) на самом деле довольно легко генерализуют на другие модальности (комп.зрение, вычислительные операции или предсказание структуры белка) без файнтюнинга self-attention или полносвязных слоёв. Получается, что предобучение на языковых задачах выучивает какие-то более глубокие и универсальные структуры вычислительной реальности, помогающие в довольно далёких задачах иной природы.

Для работы берут предобученный трансформер (в данном случае GPT-2), замораживают все его веса кроме линейных входных и выходных слоёв, а также позиционных эмбеддингов и параметров LayerNorm. Эти незамороженные веса потом будут файнтюниться, а такой замороженный трансформер называется Frozen Pretrained Transformer (FPT).

И этот FPT обучается решать задачи с качеством, аналогичным обучению полного трансформера или LSTM, несмотря на то, что файнтюнилась только 0.1% всех параметров и вообще не трогались параметры self-attention. Это интересно.

Задачи взяты следующие:

1. Bit memory: показывают пять битовых строк длины 1000, затем модели показывают одну из этих пяти с замаскированными с вероятностью 0.5 битами, и модель должна восстановить оригинальную строку.

2. Bit XOR: модели дают две битовых строки длины пять (по одному биту за раз), и она должна выдать поэлементный XOR.

3. ListOps: на вход дают последовательность операций со списками (похоже на LISP) и на выходе нужно получить результат этих операций.

4. MNIST: стандартный MNIST, где надо классифицировать картинку с цифрой, но на вход поступают 64 токена с патчами 4x4.

5. CIFAR-10: аналогично

6. CIFAR-10 LRA: модификация с Long-Range Arena, всё переведено в grayscale и flattened и подаётся по токену размера 1 за раз. По сути sequential CIFAR.

7. Remote Homology Detection: по последовательности аминокислот белка предсказать метку (всего 1195 классов). При этом никакого предобучения на базах последовательностей типа Pfam.

В случае GPT-2 Base n_dim=768, n_layers=12, входные и выходные размерности d_in, d_out зависят от задачи, максимальная длина последовательности l.

Для кейса с CIFAR-10 (d_out=10, d_in=16, l=64) получаются обучаемые параметры:
- выходной слой (768*10 = 7680 параметров)
- входной слой (16*768 = 13056)
- параметры LayerNorm (4*768*12 = 36684)
- позиционные эмбеддинги (64*768 = 49512)

Всё скейлится линейно по параметрам датасета. Для базовой модели GPT-2 со 124M параметров это всего лишь 0.086%, а для GPT-2 XL вообще 0.029%.

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

Сравниваются с 1) трансформером, который обучается с нуля на задачу, с 2) LSTM также обучающемся сразу на задачу.

На перечисленном пуле задач FPT выглядит не хуже, а местами лучше альтернатив. То есть перенос на другие модальности работает.

В некоторых случаях с маленькими датасетами обучить 12-слойный трансформер с нуля не получается, поэтому берут трансформер поменьше, например трёхслойный для CIFAR. С другой стороны для FPT увеличение модели только улучшает результат.

Сравниваются с разным другим предобучением (случайная инициализация без предобучения, предобучение на bit memory задаче, предобученный ViT -- с ним есть тонкость, потому что он обучался как энкодер, а GPT-2 -- декодер).

Здесь, во-первых, интересно, что случайная инициализация даёт весьма хороший результат (частый кейс и сильный бэйзлайн, рандомные фичи рулят). Во-вторых, предобучение на языке явно рулит.

Из более глубокого копания в случайных фичах показывают, что случайно инициализированный трансформер работает явно лучше, чем случайно инициализированная LSTM. То есть что-то есть в архитектуре. Есть в self-attention какой-то более подходящий inductive bias для универсальных вычислений.

Дообучение FPT сходится быстрее, чем дообучение случайных фич (обычно в несколько раз, но на кейсе XOR это 40x).

Пытались обнаружить какие-то интересные паттерны внимания на входе, но не получилось найти что-то интерпретируемое.

Также показывают (через разницу между test/train accuracy), что FPT не оверфитится, в отличие от обучаемых с нуля, в том числе более мелких трансформеров. И как я писал выше, более тяжёлые модели FPT только улучшают результат, в отличие от обучаемых с нуля.

Более хитрая инициализация весов, так чтобы их статистики совпадали со статистиками GPT-2, в большинстве кейсов до некоторой степени помогает поднять качество трансформера, обучаемого с нуля, но всё равно не до уровня FPT.

Из параметров, которые обучают, наиболее сильный эффект от LayerNorm. Более хитрые размораживания могут улучшить результат, в частности включение файнтюнинга полносвязных слоёв помогает, но для слоёв внимания это не работает, только хуже.

Есть эксперимент, где обучают только выход (что по сути получается похожим на reservoir computing/ESN) и оно в целом работает, но хуже FPT (зато быстрее в сотни и тысячи раз).

Пробуют вместо GPT-2 также BERT, энкодер T5 и Longformer. В целом поведение похоже, только у T5 просадка (но интересно было бы посмотреть на T5 с декодером).

В общем это интересный набор исследований. В контексте движения про Foundation Models мы и правда можем оказаться через неокторое время в ситуации, когда у нас есть огромные предобученные модели, демонстрирующие как минимум good enough результаты на куче разнородных задач, и даже если мы не понимаем, почему именно они работают, их можно хоть зашивать в железо (после различной оптимизации, и возможно даже в аналоговом виде) и предоставлять как универсальную микросхему для решения большого класса задач, по сути Микропроцессор 2.0. Если это будет достаточно массово и дёшево, может родиться новый класс доступных и умных продуктов с интеллектом “на месте”, а не в облаке.

Decision Transformer: Reinforcement Learning via Sequence Modeling
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch
Статья: https://arxiv.org/abs/2106.01345
Code: https://github.com/kzl/decision-transformer
Разбор от Янника: https://www.youtube.com/watch?v=-buULmf7dec

Работа от примерно той же группы авторов, что и в предыдущем посте про “Pretrained Transformers as Universal Computation Engines” (https://news.1rj.ru/str/gonzo_ML/711). Тоже по сути про универсальность трансформеров для решения разных задач, теперь для RL.

Здесь авторы подошли к задаче offline reinforcement learning и моделируют её как задачу conditional sequence modeling, используя авторегрессионный трансформер, аналогичный GPT.

Offline RL по сути несколько специфическая область, когда агент не имеет возможности итеративно взаимодействовать со средой, и вынужден довольствоваться как-то полученным (возможно каким-то совершенно другим неизвестным агентом) накопленным и сохранённым опытом в виде массива наблюдений.

Трансформеры уже применялись в RL (и мы даже писали про них, например здесь: https://news.1rj.ru/str/gonzo_ML/148), но обычно они встраиваются в традиционные RL алгоритмы, например, обучая policy через TD learning (в упомянутой выше работе использовалась on-policy адаптация метода Maximum a Posteriori Policy Optimization, называющаяся V-MPO).

В текущей работе авторы пробуют изменить парадигму обучения и подходят к задаче как к моделированию последовательности практически как при работе с языком.

На вход GPT-like модели подаётся траектория, закодированная последовательностью троек (R, s, a), где a — это действие, s — состояние среды, а R в данном случае не полученный при этой комбинации (s,a) reward, а return-to-go (сумма будущих вознаграждений), и это сделано для того, чтобы моделировать будущие действия не на основании прошлых вознаграждений, а на основании будущего желаемого возврата.

Для получения следующего действия на вход прилетает вся прошлая траектория + текущие значения желаемого возврата R_t и текущее состояние среды s_t, и для позиции этого s_t предсказывается действие a_t. То есть для K шагов получается 3K токенов (я бы даже сказал 3K-1).

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

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

Ну и собственно обученный трансформер работает так: на вход подают желаемый перформанс (возврат, return-to-go) и текущее состояние среды, модель генерирует действие, оно выполняется в среде, уменьшаем возврат на полученное значение вознаграждения и повторяем пока не достигли окончания эпизода.

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

Оцениваются на задачах дискретного и непрерывного управления (соответственно несколько игр из бенчмарка Атари и бенчмарк D4RL в OpenAI Gym), сравнивают с model-free методами TD learning (Conservative Q-Learning, CQL, который на тот момент вроде как был SoTA, и с некоторыми другими), а также с imitation learning.

Для Атари это трансформер размерности 128 с 6 слоями и 8 головами, для Gym всего трёхслойный с одной головой.

На Атари метод сравним с CQL и бьёт более старые алгоритмы. На задачах непрерывного управления из D4RL почти на всех задачах Decision Transformer побил традиционные RL методы.

Авторы пытаются анализировать, не занимается ли Decision Transformer behaviour cloning’ом и показывают, что он более эффективен, чем просто imitation learning на подмножестве датасета.

Отдельно изучают, важна ли глубина контекста, и да, на контексте длины 1 трансформер работает существенно хуже.

И в целом метод неплохо выучивает полиси, где надо делать long-term credit assignment, даже если обучается на рандомных полиси. В ситуации с разреженными и отложенными наградами также работает достойно.

Такие вот дела. Были когда-то прикольные работы, где GPT играла в Го (https://arxiv.org/abs/2007.03500) или шахматы (https://arxiv.org/abs/2008.04057), а здесь тот же универсальный трансформер по типу GPT решает более-менее любые RL задачи. Интересно, что будет, когда обучат какую-нибудь гигантскую модель на пиллиарды параметров.