Post #30

Продолжаем серию постов про архитектуры и методы, позволяющие бороться с квадратичной сложностью attention в трансформерах. В прошлый раз мы говорили про Linformer, Reformer, Performer и AFT, сегодня — про Sparse Transformers, Longformer, MEGA. Не стоит забывать про всякие hardware оптимизации (например, FlashAttention), но это находится за пределами этого поста. Sparse Transformers: Sparse Transformers предлагает метод подсчета attention между подмножеством токенов: Factorized Self-Attention. Основной идеей является использование p голов с разными паттернами, за счет которых каждый i-ый токен может “посмотреть” на j-ый не более чем за p+1 шагов. Для p=2 мы можем использовать следующую комбинацию (fixed pattern): первая голова смотрит на все предыдущие k токенов, вторая — на токены с фиксированных позиций, обеспечивая протекание информации из прошлого. Такой паттерн позволяет снизить сложность до n корень из n. Метод используется в GPT-3, и возможно какая-то его модификация — GPT-3.5/4. SWA из Mistral напоминает этот же принцип. Longformer: Идея похожа на Sparse Transformers. Разница в следующем: Longformer создавался для задач с длинными последовательностями, поэтому авторы добавили global attention pattern, когда мы смотрим на токены важные для конкретной задачи: в классификации — CLS токен, в задачах QA — токены из вопроса. В итоге имеем комбинацию локального и глобального вниманий. Ну и дополнительно используются всякие трюки, которые подобрали эмпирически: разный размер окна в зависимости от слоя, dilated окна для глубоких слоев и тд. Тут получаем даже линейную сложность, потому что глобальное внимание привносит только константное число доп. токенов. MEGA: Довольно тяжело объяснить идею в нескольких предложениях, поэтому посмотрите сначала на картинку. Основных идей здесь 3: 1. Используем экспоненциальное сглаживание эмбеддингов для дальнейшего получения матриц Q и K. EMA позволяет агрегировать локальный контекст с настраиваемыми весами. 2. Добавляем два gating mechanisms: первый для протекания изначального эмбеддинга X, второй — для протекания EMA эмбеддинга X’. 3. Gated Attention имеет всего одну голову, но показывается, что за счет гейтингов и EMA он эквивалентен Multi-Head Attention. Данные модификации не меняют сложность, они больше про inductive bias. Но есть вариация MEGA chunked, которая снижает сложность до линейной.