Longformer

Longformer — это трансформер для работы с длинными документами. Он заменяет полное самовнимание скользящим окном внимания вокруг каждого токена и глобальным вниманием к нескольким токенам, выбранным с учётом задачи. Благодаря этому вычислительные затраты растут линейно, а не квадратично относительно длины входных данных. Из Белтагии, Мэтью Питерс и Арман Кохан из Allen Institute for AI представили Longformer в статье 2020 года «Longformer: трансформер для длинных документов». Кодировщики, такие как BERT и RoBERTa, обрабатывают не более 512 токенов, поэтому более длинные тексты приходилось разбивать на фрагменты, теряя связи между ними. Выпущенные модели Longformer обрабатывали за один проход 4096 токенов — это контекстное окно в восемь раз длиннее.

Как работает внимание в Longformer#

Каждый слой сочетает локальные и глобальные шаблоны в рамках одного механизма внимания:

• Внимание со скользящим окном: Каждый токен учитывает фиксированное окно соседних токенов: в выпущенных моделях оно охватывает 256 токенов с каждой стороны.
• Внимание со скользящим окном и разрежением: Пропуски между позициями, к которым направлено внимание, позволяют головке охватить больше позиций при тех же затратах. Авторы использовали разрежение в нескольких головках верхних слоёв своих языковых моделей на уровне символов.
• Глобальное внимание: Выбранные токены учитывают все позиции, а все позиции учитывают их. Для этого используются отдельные проекции запросов, ключей и значений. Выбор зависит от задачи: токен [CLS] используется для классификации, а токены вопроса — для ответов на вопросы.

Наложение слоёв расширяет охват окна так же, как наложение свёрточных слоёв увеличивает рецептивное поле свёрточной нейронной сети (CNN): l слоёв с окном ширины w охватывают l × w токенов. Каждый слой вычисляет примерно n × w оценок внимания вместо n × n. Разница наглядно видна при подсчёте пар в маске Longformer:

import torch

window = 512  # Longformer-base: 256 tokens on each side of every token

for n in (1024, 2048, 4096):
    pos = torch.arange(n)
    # Sliding window: each token attends to neighbors within window / 2 positions
    mask = (pos[:, None] - pos[None, :]).abs() <= window // 2
    # Global attention on token 0 ([CLS]): it sees every token and every token sees it
    mask[0, :] = True
    mask[:, 0] = True
    print(f"{n} tokens: full attention {n * n / 1e6:.1f}M scores, Longformer {mask.sum().item() / 1e6:.1f}M")

Он выводит 0,5 млн, 1,0 млн и 2,0 млн оценок Longformer против 1,0 млн, 4,2 млн и 16,8 млн для полного внимания: при каждом удвоении входных данных вычислительные затраты на внимание Longformer удваиваются, а затраты на полное внимание увеличиваются в четыре раза. В реальных реализациях вычисляется только полоса оценок, что и обеспечивает экономию памяти.

Обучение и результаты#

Longformer начинается с контрольной точки RoBERTa. Авторы восемь раз скопировали 512 обученных позиционных эмбеддингов, чтобы охватить 4096 позиций, а затем продолжили предварительное обучение языковой модели с маскированием на длинных документах. После дообучения модель стабильно превосходила RoBERTa на задачах с длинными документами и на тот момент достигла лучших результатов на тестах WikiHop и TriviaQA для ответов на вопросы. Языковая модель на уровне символов языковая модель также установила новые результаты на текстах text8 и enwik8.

Longformer-Encoder-Decoder (LED) добавляет декодер для суммаризации текста и других задач генерации текста. Её возможности показали на примере суммаризации полных научных статей с arXiv. В Transformers от Hugging Face реализованы обе модели, а глобальные токены передаются как global_attention_mask. Позже появились специализированные версии, например Clinical-Longformer, предварительно обученная на больших корпусах клинических текстов и протестированная на 10 задачах, включая распознавание именованных сущностей, ответы на вопросы и классификацию документов.

Longformer и родственные модели#

• Transformer-XL: Читает длинный текст по сегментам и позволяет каждому сегменту учитывать кэшированные скрытые состояния предыдущего сегмента. За счёт рекуррентной обработки модель увеличивает контекст при генерации слева направо.
• BigBird: Модель Google того же года, которая добавляет случайные связи внимания к оконному и глобальному вниманию и тоже обрабатывает 4096 токенов.
• Разреженное внимание: Общее семейство фиксированных шаблонов внимания, к которому относится Longformer.
• Flash attention: Вычисляет точное полное внимание в тайлах быстрой памяти на кристалле, поэтому объём памяти растёт линейно с длиной последовательности, хотя количество оценок остаётся квадратичным.

Локальное и глобальное внимание в компьютерном зрении#

Изображения высокого разрешения создают ту же проблему, поскольку визуальный трансформер (ViT) преобразует изображение в тысячи токенов фрагментов. Swin Transformer решает эту задачу с помощью сдвигаемых локальных окон, а Segment Anything Model (SAM) от Meta, запускаемая в пакете Ultralytics, чередует блочные оконное и глобальное внимание для сегментации изображений. Кодировщик изображений ViT-B использует внимание внутри неперекрывающихся окон 14 × 14 на сетке фрагментов 64 × 64 в 8 из 12 блоков, а в блоках 2, 5, 8 и 11 применяет глобальное внимание ко всем 4096 фрагментам.