[CMU] Unlimiformer: Long-Range Transformers with Unlimited Length Input
Amanda Bertsch, Uri Alon, Graham Neubig, Matthew R. Gormley
Статья: https://arxiv.org/abs/2305.01625
Код: https://github.com/abertsch72/unlimiformer
Делу борьбы с limited attention span трансформеров посвящается!
Как известно, ограничение стандартного трансформера на размер контекста связано с квадратичной сложностью механизма внимания. И типичные трансформеры имеют размер контекста порядка 512-2048 токенов, в редких случаях побольше. У GPT-4 он 8192 токена, и есть пока не доступная массово версия на 32k (но вряд ли это классический механизм внимания). Для многих реальных датасетов этого недостаточно.
Было много подходов увеличить размер этого контекста и понизить вычислительную сложность трансформеров, про многие такие решения мы в своё время писали: Performer (https://hottg.com/gonzo_ML/404), Longformer (https://hottg.com/gonzo_ML/292), Big Bird (https://hottg.com/gonzo_ML/381) и многое другое. В массы, правда, почему-то эти решения не пошли.
Очередное решение радикально решает вопрос увеличения attention span, теперь он “бесконечный”. Это достигается за счёт использования retrieval-based механизма, который можно добавить к уже обученному трансформеру.
Работает это так для архитектуры encoder-decoder.
Входная последовательность разбивается на чанки с перекрытием, они кодируются энкодером трансформера, оставляется при этом только средняя половина выходной последовательности каждого чанка. Это нужно для сохранения контекста с обеих сторон.
Полученные эмбеддинги отправляются во внешнее векторное хранилище, например, типа FAISS (https://engineering.fb.com/2017/03/29/data-infrastructure/faiss-a-library-for-efficient-similarity-search/, https://github.com/facebookresearch/faiss).
В декодере перед cross-attention делается k-nearest neighbor (kNN) search во внешнем хранилище и для каждой головы каждого слоя внимания выбирается набор токенов, которые надо учесть. Таким образом вместо того, чтобы в декодере смотреть на все N (=размер контекста) токенов из энкодера, мы смотрим на N “лучших” токенов входной последовательности, поднятой из базы. А декодер всё так же имеет техническое ограничение на размер контекста в N токенов.
Никаких обучаемых параметров в kNN-search нету, и этот механизм можно встроить в любой уже обученный трансформер.
Это всё в каком-то смысле развитие идеи Memorizing Transformers (https://arxiv.org/abs/2203.08913), но там хранилища эмбеддингов подразумевались для каждой головы каждого слоя, и это было слишком много, поэтому в той работе ограничились одним слоем энкодера. Работа про Memorizing Transformers отличалась также и наличием обучаемых весов для гейта, который комбинировал локальное и kNN внимание, поэтому его нельзя было легко применить к уже обученной LM.
В текущей работе переформулировали механизм внимания так, что по новой формуле с другим порядком вычисления достаточно одного хранилища для всех голов и слоёв. В новой формулировке QK^T = (h_d * W_q * W^T_k) * h^T_e, имея скрытое состояние декодера h_d и сделав его проекцию в новое пространство под каждый слой и голову внимания (умножив на W_q и W^T_k), можно довольно просто поднять из векторного хранилища соответствующие скрытые состояние энкодера h_e. Вроде как такая формулировка встречается впервые в retrieval-based моделях.
Перечисленные изменения можно добавить в уже обученный трансформер и ничего в нём больше не менять. Но можно ещё и пофайнтюнить. Пробовали разные режимы, с малой и большой вычислительной стоимостью. При малой стоимости Unlimiformer добавляется только на тесте, а при большой он используется и в обучении.
Потестили на датасетах для суммаризации длинных документов. Добавляли Unlimiformer к BART_base (https://hottg.com/gonzo_ML/153) с контекстом 1024 токена и PRIMERA (https://aclanthology.org/2022.acl-long.360/) с контекстом 4096 токенов.
Подход работает и улучшает результат. Файнтюнинг добавляет качества и бьёт альтернативные улучшалки базовых моделей на Memorizing Transformers и SLED (https://arxiv.org/abs/2208.00748).
>>Click here to continue<<
