Генеративное моделирование с разреженными трансформерами (sparse transformers)23.04.2019 07:00

Мы разработали Sparse Transformer — глубокую нейронную сеть, которая устанавливает новые рекорды в прогнозировании того, что произойдет дальше в последовательности: будь то текст, изображения или звук. В ней используется алгоритмическое усовершенствование механизма внимания (attention), позволяющее извлекать закономерности из последовательностей, которые в 30 раз длиннее тех, что были доступны ранее.

Одна из существующих проблем в исследованиях ИИ — моделирование долгосрочных, тонких взаимосвязей в сложных данных, таких как изображения, видео или звуки. Sparse Transformer включает в себя реформулирование механизма самовнимания O (N2) O (N^2) Transformer⁠ в виде O (NN) O (N \sqrt{N}) , а также ряд других улучшений, что позволяет применять его непосредственно к этим богатым типам данных. Раньше модели, используемые для таких данных, создавались специально под один домен или с трудом масштабировались на последовательности длиной более нескольких тысяч элементов. Наша же модель способна обрабатывать последовательности с десятками тысяч элементов, используя сотни слоев, и достигает лучших в отрасли показателей в самых разных областях. В OpenAI мы используем её для создания систем ИИ, обладающих большими возможностями понимания мира.

Глубокое внимание

В архитектуре Transformer каждый выходной элемент связан с каждым входным, а весовые коэффициенты между ними рассчитываются динамически в зависимости от контекста — этот процесс называется вниманием (attention). Считается, что это делает модели Transformer более гибкими по сравнению с моделями с фиксированными схемами связей, однако на практике для каждого слоя и каждого головного устройства внимания требуется создание матрицы внимания размером N×N N\times N , что может потреблять огромные объемы памяти при работе с данными, содержащими множество элементов, такими как изображения или необработанный аудиопоток.

Один из способов сократить это потребление — пересчитывать матрицу внимания из чекпоинтов во время обратного распространения ошибки. Это хорошо зарекомендовавший себя в глубоком обучении метод снижения потребления памяти ценой увеличения объема вычислений. Применительно к матрице внимания в Transformer это означает, что пиковые затраты памяти перестают зависеть от количества слоев, что позволяет обучать сети существенно большей глубины, чем раньше. На практике мы выяснили, что модели Transformer глубиной до 128 слоев превосходят более мелкие сети на тестовых задачах, таких как CIFAR-10.

Чтобы обучать эти модели увеличенной глубины, мы внесли ряд корректировок в порядок операций в трансформаторе и изменили схему инициализации. Все подробности можно найти в нашей работе.

Разреженное внимание

Однако даже вычисление одной-единственной матрицы внимания может стать непрактичным при очень больших входных данных. Вместо этого мы используем разреженные паттерны внимания, в которых каждая выходная позиция вычисляет весовые коэффициенты только для подмножества входных позиций. Когда это подмножество мало по сравнению со всем набором входных данных (например,  N \sqrt{N}  элементов вместо N N ), итоговое вычисление внимания становится выполнимым даже для очень длинных последовательностей, обладая алгоритмической сложностью O (NN) O (N \sqrt{N})  вместо O (N2) O (N^2) .

Чтобы оценить жизнеспособность такого подхода, мы сначала визуализировали обученные паттерны внимания для глубоких моделей Transformer на изображениях и обнаружили, что многие из них демонстрируют интерпретируемые и структурированные паттерны разреженности. На каждом из изображений ниже показано, на какие входные пиксели (подсвечены белым) опирается данная головка внимания для прогнозирования следующего значения на изображении. Когда области входа сосредоточены на небольших подмножествах и демонстрируют высокую степень регулярности, этот слой поддается разрежению. Некоторые из них представлены здесь для 128-слойной модели на изображениях CIFAR-10:

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

Первая версия,  строчное (strided) внимание, грубо говоря, эквивалентна тому, что каждая позиция обращается к своей строке и своему столбцу, и похожа на паттерн внимания, изученный сетью выше. (Заметим, что внимание по столбцам можно эквивалентно сформулировать как внимание к строке транспонированной матрицы). Вторая версия,  фиксированное (fixed) внимание, учитывает фиксированный столбец и элементы после последнего элемента столбца — этот паттерн показался нам полезным, когда данные не укладывались в двумерную структуру (например, текст). За более подробными сведениями мы отсылаем читателей к нашей работе.

Экспериментальные результаты

Sparse Transformers установили новые рекорды по оценке плотности распределения (density estimation) для наборов данных CIFAR-10, Enwik8 и Imagenet 64.

Мы также обнаружили, что разреженное внимание достигает меньшей функции потерь (loss), чем полное внимание, помимо того, что оно значительно быстрее (сравнения см. в нашей работе). Это может указывать на полезную индуктивную предвзятость (inductive bias), обеспечиваемую нашими паттернами разреженности, либо на скрытые проблемы оптимизации при плотном внимании.

Генерация изображений

Модели Transformer, использующие разреженное внимание, судя по всему, обладают представлением о глобальной структуре, что можно качественно оценить по дополнению изображений (image completions). Ниже мы визуализируем модель, обученную на ImageNet с разрешением 64×64 64\times 64 :

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

Генерация сырых аудиосигналов

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

Эта модель была обучена на сырых фрагментах классической музыки и использует разреженное внимание для генерации последовательностей длиной 65 000. Это соответствует примерно 5 секундам необработанного аудио, и в каждом из приведенных ниже фрагментов мы объединили несколько выборок.

Релиз кода

Обычно реализация разреженного внимания связана с нарезкой матриц запросов и ключей на блоки, поэтому для облегчения экспериментов мы реализовали набор блочно-разреженных ядер⁠, которые эффективно выполняют эти операции на GPU. Мы открываем исходный код этих ядер и предоставляем примеры функций разреженного внимания в этом репозитории⁠.

Будущая работа и ограничения

• Представленные нами паттерны разреженного внимания — это лишь предварительные шаги в направлении эффективного моделирования длинных последовательностей. Мы считаем, что исследование различных паттернов и комбинаций разреженности полезно, и что обучение разреженным паттернам является особенно перспективным направлением исследований для нового поколения архитектур нейронных сетей.
• Даже с учетом описанных выше улучшений авторегрессионная генерация последовательностей все еще кажется непрактичной для изображений или видео очень высокого разрешения. Однако созданные нами оптимизированные операции внимания могут оказаться полезными примитивами для объединения с другими подходами к моделированию данных высокой размерности, такими как многомасштабные подходы.

Если вы заинтересованы в развитии возможностей ИИ и помощи в продвижении нашей миссии по обеспечению его пользы для человечества, мы нанимаем сотрудников⁠!

Спасибо Ашишу Васвани (Ashish Vaswani) за полезные обсуждения, а также Иоханнесу Оттербаху (Johannes Otterbach), Марку Чену (Mark Chen), Прафулле Дхаривалу (Prafulla Dhariwal), Дэвиду Луану (David Luan) и Лукашу Кайзеру (Lukasz Kaiser) за комментарии к рукописи.