Понимание последовательно организованных данных будь то язык, музыка или видео трудная задача, особенно в случаях, когда они сильно зависят от контекста, который их окружает. Например, если человек или какой-либо предмет пропадёт из поля зрения на видеозаписи и появится снова через значительный промежуток времени, многие модели забудут, как он выглядел. В сфере обработки языка нейронные сети с долгой краткосрочной памятью (long short-term memory, LSTM) охватывают достаточный контекст для того, чтобы успешно осуществлять последовательный перевод предложение за предложением. В этом случае контекстное окно (т.е. охват данных, которые модель принимает во внимание при переводе) может содержать от десятка до сотни слов. Более новая модель Трансформера не только улучшила качество последовательного перевода, но может быть использована для генерации целых статей Википедии с помощью суммаризации множества документов. Это возможно благодаря тому, что Трансформер увеличил контекстное окно до тысячи слов. Кроме того, столь обширный рассматриваемый контекст позволяет использовать Трансформер для обработки не только текста, но и пикселей или музыкальных нот, на основе которых можно сгенерировать изображения или музыку.

Однако еще большее расширение контекстного окна Трансформера неизбежно упирается в ограничения. Основу этого фреймворка составляет механизм внимания, который подразумевает оценивание всех возможных пар слов внутри контекстного окна для понимания взаимосвязей между ними. Так, для текста из 100К слов, модель должна будет оценить 100К х 100К пар слов, а это 10 миллиардов пар на каждом этапе, что кажется довольно непрактичным. Еще одну проблему составляет стандартная практика сохранять выход каждого слоя модели. Для приложений, использующих большие контекстные окна, требуется огромный объем памяти для хранения выходов со всех слоев модели (от гигабайтов для моделей с несколькими слоями до терабайтов с тысячами). Это означает, что реальные модели на основе Трансформера, состоящие из множества слоев, могут использоваться только на небольших абзацах текста или для генерации коротких музыкальных фрагментов.

В данной статье представлен Reformer модель на основе Трансформера, которая разработана для работы с большими контекстными окнами до 1 миллиона слов на одном графическом ускорителе и с использованием всего 16 Гб памяти. Reformer объединяет два важнейших метода решения проблем внимания и выделения памяти, не позволявших использовать Трансформеры для больших контекстных окон: хеширование с учетом местоположения (locality-sensitive-hashing, LSH), помогающее снизить сложность обращения к длинным последовательностям, а также обратимые остаточные слои (reversible residual layers) для более эффективного использования доступной памяти.

Проблема внимания

Первый вызов, с которым сталкиваются при применении модели на основе Трансформера к очень большой текстовой последовательности что делать со слоем внимания? LSH предлагает следующий выход: вместо того, чтобы искать среди всех возможных комбинаций пар векторов, можно подсчитать хеш-функцию, которая объединит похожие вектора. Например, в задаче машинного перевода, где каждый вектор из первого слоя нейросети представляет собой слово (и еще больший контекст для последующих слоев), вектора одних и тех слов в разных языках получат один и тот же хеш. На схеме ниже разные цвета соотносятся с разными хешами, а схожие слова имеют один цвет. После того, как слова хешированы, последовательность перемешивают таким образом, что элементы с одним и тем же хешом оказываются вместе, а затем разбиваются на группы чанки (chunks), чтобы распараллелить вычисления. Далее к этим небольшим чанкам (и их ближайшим соседям для фиксации перехода) применяется механизм внимания, что значительно уменьшает вычислительную нагрузку.

Хеширование с учетом местоположения: Reformer получает на вход первого слоя последовательность ключей, где каждый ключ является векторным представлением каждого слова (или пикселя, в случае с изображениями), а затем и более широкого контекста для последующих слоев. LSH применяется к последовательности, после чего ключи сортируются по их хешу и разделяются на чанки. Механизм внимания применяется только к одному чанку и его ближайшим соседям.

Проблема памяти

И хотя LSH решает проблему использования механизма внимания, вопрос с памятью остается открытым. Один слой нейронной сети часто требует до нескольких гигабайт памяти и обычно помещается на одном GPU, так что, если бы использовалась модель с одним слоем, то можно было бы обрабатывать даже длинные последовательности. Но в случае обучения многослойной модели с помощью градиентного спуска, активации с каждого слоя должны быть сохранены для использования их во время обратного распространения ошибки. Типовая модель на основе Трасформера имеет десятки и более слоев, поэтому при кэшировании значений со всех слоев память очень быстро заканчивается.

Второе нововведение, предложенное в Reformer'е, заключается в следующем: чтобы не хранить вход для каждого слоя в памяти, его просто переподсчитывают каждый раз, когда происходит обратное распространение ошибки. Это достигается с помощью использования обратимых слоев, в которых для восстановления активаций любого промежуточного слоя используются активации последнего слоя, что равносильно запуску сети в обратном направлении. В типичной остаточной нейронной сети каждый слой в стеке пополняет вектора, проходящие через сеть. Обратимые слои, напротив, имеют два набора активаций для каждого слоя. Один из них соответствует стандартному процессу, который был описан выше, и постепенно обновляется по мере перехода от одного слоя к другому; второй же фиксирует только изменения первого. Таким образом, для того чтобы запустить нейросеть в обратном порядке, достаточно вычесть активации, применяемые к каждому слою.

Обратимые слои: (a) В стандартной остаточной нейросети активации с каждого слоя используются для обновления входных данных в следующем слое. (b) В обратимой нейронной сети существует два набора активации, из которых только один обновляется после каждого слоя. (с) Этот подход позволяет запускать сеть в обратном направлении для того, чтобы восстановить все промежуточные значения.

Применения Reformer'а

Применение описанных выше новых подходов в Reformer'е позволяет повысить его эффективность, а также применять его для обработки текстовых последовательностей до 1 миллиона слов длиной на одном графическом ускорителе с использованием всего 16 Гб памяти. Благодаря высокой эффективности, Reformer может напрямую обрабатывать текстовые данные с контекстным окном намного более большим, чем практически все современные наборы данных. Возможно, такая способность Reformer'а вдохновит сообщество к созданию подобных крупных наборов данных.

Однако существует сфера, где нет недостатка в данных с обширным контекстом это генерация изображений. В этом Colab-ноутбуке представлены примеры применения Reformer'а для графических данных. Начиная с фрагментов изображения, показанных в верхнем ряду на схеме ниже, Reformer способен пиксель за пикселем сгенерировать целые изображения (см. нижний ряд).

Сверху: фрагменты изображений, использованные в качестве входной последовательности Reformer'а. Снизу: Дополненные полнокадровые изображения. Оригинальные изображения взяты из набора данных Imagenet64.

В то время как применение Reformer'а к задачам обработки изображений и видео демонстрирует большой потенциал, его использование для обработки текста кажется еще более многообещающим. Reformer может обрабатывать сразу и на одном устройстве целые романы. Так, в этом Colab-ноутбуке демонстрируется обработка Преступления и наказания. В будущем, когда будет создано достаточное количество текстовых наборов данных с большим контекстным окном, такие разработки, как Reformer, позволят генерировать длинные и связные произведения.

Вывод

Разработчики полагают, что Reformer закладывает основу для будущего использования моделей на основе Трансформера как для обработки длинных текстовых данных, так и для задач за рамками обработки естественного языка. Следуя традиции открытых исследований, авторы уже начали изучать возможность применения таких моделей для еще более длинных последовательностей и улучшения сценария использования позиционного кодирования, а также приглашают всех ознакомиться со статьей о Reformer'е, изучить их код и предложить свои идеи. В сфере глубокого обучения на данный момент существует мало текстовых наборов данных с большим контекстом, однако в реальном мире такие данные повсюду. Возможно, читателям этой статьи удастся найти новое применение Reformer'а. А пока можно начать с этого Colabноутбука и связаться с разработчиками, если возникнут проблемы или вопросы.

Авторы

• Авторы оригинала Nikita Kitaev, ukasz Kaiser
• Перевод Смирнова Екатерина
• Редактирование и вёрстка Шкарин Сергей