Глава 3.1 — Стоимость внимания

Содержание

  1. Продолжаем с того места, где закончилась часть II: архитектура решена, стоимость остаётся
  2. Почему сложность self-attention квадратична, в деталях
  3. Разреженное внимание: Longformer и BigBird
  4. Линейное внимание: переформулировка softmax
  5. FlashAttention: тот же ответ, вычисленный честно
  6. FlashAttention-2 и тенденция, которую она отражает
  7. Взгляд с точки зрения собеседования
  8. Вопросы для самопроверки
  9. Источники

1. Продолжаем с того места, где закончилась часть II: архитектура решена, стоимость остаётся

Часть II оставила нам конкретный, устоявшийся ответ на вопрос "что такое Transformer": стек блоков, каждый из которых сочетает многоголовое self-attention с позиционно-независимой полносвязной сетью, обёрнутый в некоторую схему нормализации, с позиционной информацией, внесённой либо аддитивно во вход, либо непосредственно в вычисление внимания через что-то вроде RoPE или ALiBi. Это полный архитектурный рецепт, и это тот самый рецепт, что лежит в основе практически каждой крупной языковой модели, находящейся сегодня в промышленной эксплуатации. Но рецепт, работающий на бумаге, и рецепт, работающий в масштабе окна контекста в сто тысяч токенов, — это две разные вещи, и разрыв между ними — тема этой главы.

Этот разрыв возникает из единственного вычислительного факта, о котором часть II упомянула лишь мимоходом: self-attention, в своей исходной формулировке, требует, чтобы каждый токен последовательности вычислял оценку совместимости с каждым другим токеном. Это нормально, когда ваши последовательности длиной в несколько сотен токенов. Это становится доминирующей стоимостью во всей модели — затмевая полносвязные слои, затмевая всё остальное — как только последовательности доходят до десятков или сотен тысяч токенов, а именно в этом направлении и движется область, гоняясь за более длинными документами, более длинными диалогами и более длинными цепочками рассуждений. Эта глава посвящена тому, почему так происходит и что с этим было сделано; следующая глава берётся за смежную, но отдельную проблему — то, что даже если бы вы могли вычислять внимание по сколь угодно длинной последовательности бесплатно, позиционные кодировки, с которыми вы обучались, могут попросту не знать, что делать с позициями, которых они никогда не видели.

2. Почему сложность self-attention квадратична, в деталях

Вспомним механику масштабированного скалярного внимания: имея запросы $Q$, ключи $K$ и значения $V$, каждый формы $(n, d)$ для последовательности длины $n$ и размерности головы $d$, внимание вычисляет $\text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$. Именно в слагаемом $QK^T$ и кроется проблема. $Q$ имеет форму $(n, d)$, а $K^T$ — форму $(d, n)$, так что их произведение — это матрица $(n, n)$ — по одной записи для каждой пары токенов в последовательности. Вычисление этой матрицы занимает $O(n^2 d)$ времени, а простое её хранение занимает $O(n^2)$ памяти. Каждая другая операция в блоке Transformer — линейные проекции, полносвязные слои, нормализация — масштабируется линейно по $n$. Одно только внимание масштабируется квадратично, и как только $n$ становится достаточно большим, квадратичный рост побеждает линейный вне зависимости от того, насколько мал постоянный множитель.

Стоит точно понимать, откуда на самом деле берётся эта квадратичная стоимость, потому что это не причуда какой-то конкретной реализации — это структурное свойство. Каждому токену нужно знать, для каждого другого токена, насколько тот релевантен, прежде чем он сможет сформировать взвешенную комбинацию значений. Это сравнение всех со всеми по определению, а сравнение всех со всеми среди $n$ элементов стоит $O(n^2)$ вне зависимости от того, насколько изощрённо реализована арифметика. Удвоение длины контекста не удваивает стоимость внимания, оно её учетверяет — перейдите от 2 тыс. к 128 тыс. токенов, увеличение длины в 64 раза, и вы получаете рост вычислений и памяти для внимания примерно в 4096 раз, при неизменности всего остального — стоимость масштабируется как $(n'/n)^2$, а $64^2 = 4{,}096$. Именно поэтому "просто сделать окно контекста длиннее" никогда не было вопросом простого изменения значения в конфигурации; это вынуждает либо смириться с суровой стоимостью, либо приблизить вычисление, либо найти принципиально более дешёвый способ вычислить то же самое.

Два семейства решений, рассматриваемых в этой главе, атакуют эту проблему с разных философских позиций. Паттерны разреженного внимания (Longformer, BigBird) принимают приближение: они вычисляют внимание по тщательно выбранному подмножеству пар токенов, а не по всем сразу, делая ставку на то, что бо́льшая часть релевантного сигнала улавливается этим подмножеством. Линейное внимание переформулирует всё вычисление через ядерные отображения признаков так, что ему вообще никогда не нужно материализовать матрицу $(n,n)$, ценой отказа от точной нелинейности softmax. FlashAttention занимает совершенно иную позицию: она ничего не меняет математически, вычисляет в точности тот же выход внимания, но перестраивает то, как вычисление обращается к памяти GPU, так что на практике оно работает значительно быстрее. Понимание того, что это три подлинно разные стратегии — приблизить паттерн, переформулировать алгебру или оптимизировать отображение на оборудование, — составляет суть этой главы.

3. Разреженное внимание: Longformer и BigBird

Самый прямой способ избежать вычисления матрицы внимания $(n,n)$ — просто заранее решить, что большинство из этих $n^2$ пар вообще не нужно вычислять. Это идея, лежащая в основе паттернов разреженного внимания, и Longformer — наглядная иллюстрация этой идеи. Longformer ограничивает большинство токенов локальным скользящим окном: каждый токен обращает внимание только на фиксированное число соседей с обеих сторон, что снижает стоимость этой части внимания до $O(n \cdot w)$ для размера окна $w$ — линейно по длине последовательности, а не квадратично. Однако одного скользящего окна было бы слишком ограничивающе, поскольку это означает, что информация может распространяться лишь на несколько токенов за слой, а некоторым задачам действительно нужно, чтобы токен обращал внимание на что-то далёкое (токену классификации нужно видеть весь документ, токену вопроса нужно видеть отрезок ответа там, где бы он ни находился). Решение Longformer — назначить небольшое число "глобальных" токенов (задаче-специфичных токенов вроде токена [CLS], или в некоторых постановках — конкретных важных токенов), которые обращают внимание на все остальные токены последовательности и на которые обращают внимание все остальные токены, без ограничения окном. Этот гибрид дешёвого локального внимания плюс небольшая доля дорогого, но редкого глобального внимания даёт бо́льшую часть модельной выгоды от полного внимания за малую долю его стоимости.

BigBird обобщает эту идею, комбинируя три паттерна внимания, а не два: локальное окно (как у Longformer), набор глобальных токенов (как у Longformer), и вдобавок случайный паттерн внимания, где каждый токен также обращает внимание на небольшое число случайно выбранных других токенов. Значимость случайного компонента легко недооценить: он улучшает теоретико-графовую связность паттерна внимания. Если представить паттерн внимания как граф, где ребро означает "этот токен может передать информацию тому токену", то одно лишь локально-глобальное внимание всё ещё может оставлять некоторые пары токенов на расстоянии многих "прыжков" друг от друга, а значит, информации приходится проходить через несколько слоёв, чтобы их соединить. Добавление разреженных случайных рёбер сокращает эти пути, и авторы BigBird показывают, что именно это свойство связности позволяет разреженному паттерну сохранять теоретическую выразительность полного внимания — конкретно, они показывают, что их механизм разреженного внимания является универсальным аппроксиматором функций последовательность-в-последовательность и тьюринг-полным, свойства, на которые наивный паттерн только-локального-окна претендовать не может.

Оба подхода принимают приближение в обмен на переход от $O(n^2)$ к $O(n)$ или $O(n \log n)$ по длине последовательности, и оба были проверены на действительно длинных документных задачах, где полная стоимость $O(n^2)$ была ранее недопустимой. Тем не менее компромисс здесь реален: выбор того, какой паттерн использовать (размер окна, сколько глобальных токенов, сколько случайности), сам по себе является проектным решением, которое должно соответствовать задаче, и ошибка в этом выборе означает молчаливый отказ от связей внимания, которые имели значение.

4. Линейное внимание: переформулировка softmax

Разреженное внимание сохраняет softmax в точности как было и просто вычисляет его по меньшему числу пар. Линейное внимание идёт другим путём: оно меняет саму алгебру внимания так, что матрица $(n,n)$ никогда не формируется, ни точно, ни приближённо, ни по какому-либо подмножеству. Ключевой ход, предложенный Katharopoulos et al., состоит в том, чтобы заметить, что роль softmax во внимании — превратить оценку сходства $Q_iK_j^T$ во что-то вроде неотрицательного веса, и что это можно заменить более общей функцией сходства $\text{sim}(q, k)$, выраженной как ядро: $\text{sim}(q,k) = \phi(q)^T\phi(k)$ для некоторого отображения признаков $\phi$. Как только сходство выражается таким образом, выход внимания для данного запроса можно переписать, используя ассоциативность матричного умножения: вместо вычисления $(\phi(Q)\phi(K)^T)V$ — которое всё ещё формирует промежуточную матрицу $(n,n)$ — вы вычисляете $\phi(Q)(\phi(K)^T V)$, где $\phi(K)^T V$ — это матрица $(d, d)$, которая вообще не растёт с длиной последовательности. Эта перестановка порядка и есть весь трюк: те же три матрицы, перемноженные в другом порядке, никогда не требуют существования в памяти объекта $(n,n)$, а суммарная стоимость падает до $O(n d^2)$ — линейно по $n$.

Эта переформулировка имеет второе, пожалуй, ещё более важное следствие: она показывает, что авторегрессивное линейное внимание можно вычислять в точности как рекуррентную нейронную сеть. Поскольку матрица $(d,d)$ $\phi(K)^T V$ может накапливаться инкрементально по мере поступления новых токенов — добавляя внешнее произведение нового токена $\phi(k_t)^T v_t$ к бегущей сумме, — трансформеры с линейным вниманием можно запускать во время инференса с постоянной стоимостью на токен и постоянной памятью, независимо от того, сколько токенов уже сгенерировано, в точности как скрытое состояние RNN. Katharopoulos et al. явно формулируют это как "трансформеры — это RNN", и это по-настоящему полезная мысленная модель: квадратичный по стоимости, полностью параллельный взгляд на внимание во время обучения и постоянный по стоимости, последовательный рекуррентный взгляд на внимание во время инференса — это две стороны одного и того же вычисления, связанные тем, используете ли вы ассоциативность матричного умножения.

Цена этой эффективности — потеря точного softmax. Softmax-внимание обладает особым, полезным свойством: это нормализованное, резко заострённое распределение, которое при достаточной ёмкости может эффективно вести себя как жёсткий поиск, обращая внимание почти полностью на один или два по-настоящему релевантных токена. Ядерные отображения признаков, приближающие сходство softmax, как правило, дают более гладкие, менее заострённые взвешивания, и выбор отображения признаков $\phi$ становится реальным проектным решением с реальными последствиями для качества — неудачное отображение признаков может измеримо навредить способности модели к тому виду резкого, избирательного поиска, в котором внимание хорошо. Варианты линейного внимания получили применение в отдельных, критичных по эффективности сценариях, но полное softmax-внимание осталось выбором по умолчанию в большинстве передовых моделей, именно потому что эта заострённость, как правило, важна для качества так, как это трудно полностью восстановить с помощью ядерного приближения.

Этот вердикт был верен на 2020 год и заслуживает пересмотра, потому что открытый в этом разделе рекуррентный взгляд оказался более значимой половиной вклада. Если авторегрессивное линейное внимание — это RNN с матричным состоянием фиксированного размера, то интересен не выбор ядерного отображения признаков, а правило, по которому это состояние обновляется, и «прибавить новое внешнее произведение к бегущей сумме» у Katharopoulos et al. — лишь простейший из множества вариантов. Именно на этот вопрос отвечают Mamba, RWKV и модели на дельта-правиле, и глава 3.6 разбирает его целиком.

5. FlashAttention: тот же ответ, вычисленный честно

И разреженное, и линейное внимание меняют то, что вычисляется. FlashAttention не меняет ничего в том, что вычисляется — она вычисляет в точности тот же выход softmax-внимания, что и исходная формулировка, бит в бит в пределах точности плавающей запятой, — а вместо этого меняет то, как вычисление отображается на оборудование GPU. Чтобы понять, почему одно это способно дать большой прирост скорости, нужно разделить две вещи, которые легко спутать: сколько операций с плавающей запятой выполняет алгоритм и сколько времени ему реально требуется на выполнение. У современных GPU есть небольшой объём чрезвычайно быстрой памяти на кристалле (SRAM) и гораздо больший объём более медленной памяти вне кристалла (память с высокой пропускной способностью, HBM). Наивная реализация внимания вычисляет полную матрицу оценок $(n,n)$ в HBM, применяет к ней softmax в HBM, а затем умножает на $V$, читая и записывая эту целую матрицу $(n,n)$ в медленную память многократно. Для длинных последовательностей именно это перемещение данных между HBM и вычислительными блоками GPU — а не сама арифметика — является реальным узким местом. Внимание на современном оборудовании ограничено пропускной способностью памяти, а не вычислениями, а это значит, что алгоритм, выполняющий ровно то же число операций с плавающей запятой, но перемещающий меньше данных, может быть существенно быстрее по реальному времени выполнения.

Решение FlashAttention состоит в том, чтобы вообще никогда не материализовать полную матрицу $(n,n)$ в HBM. Она работает через тайлинг: разбивает запросы, ключи и значения на блоки, достаточно маленькие, чтобы поместиться в быструю SRAM, и обрабатывает внимание блок за блоком, вычисляя частичные статистики softmax для каждого блока и объединяя их инкрементально по мере прохождения по последовательности. Это требует определённой численной аккуратности — нельзя вычислить нормализацию softmax, пока не увидели все оценки, но и не хочется держать все оценки в памяти, — и FlashAttention решает это с помощью техники "онлайн softmax", поддерживающей бегущий максимум и бегущую сумму, перемасштабируя ранее накопленные частичные результаты по мере поступления новых блоков, так что итоговый результат математически идентичен стандартному softmax, вычисленному сразу целиком. Итоговый эффект в том, что FlashAttention выполняет то же суммарное число FLOPs, что и стандартное внимание (на самом деле, несколько больше, из-за пересчёта, который она выполняет во время обратного прохода, чтобы избежать хранения всей матрицы оценок), но значительно меньше чтений и записей в HBM, а поскольку пропускная способность HBM и есть настоящее узкое место, это выливается в большой прирост реальной скорости и в объём памяти, масштабирующийся линейно по длине последовательности, а не квадратично.

Это подлинно иная категория оптимизации по сравнению с разреженным или линейным вниманием, и стоит явно зафиксировать это различие в своей мысленной модели, потому что интервьюеры прощупывают именно этот момент: FlashAttention — не приближение и вообще не меняет выходы модели; это системная оптимизация неизменного математического объекта. Это значит, что она сочетается со всем остальным в этой главе — можно применять IO-осведомлённые техники тайлинга и к паттернам разреженного внимания тоже, и люди так и делают, — и это значит, что принятие FlashAttention не стоит ничего с точки зрения качества модели, что во многом объясняет, почему её приняли по сути повсеместно и практически мгновенно во всей области.

6. FlashAttention-2 и тенденция, которую она отражает

FlashAttention-2 совершенствует ту же базовую идею, а не заменяет её, и эти усовершенствования стоит знать хотя бы на общем уровне, поскольку они иллюстрируют, сколько производительности всё ещё оставалось "на столе" из-за инженерных особенностей оригинала, отдельно от каких-либо изменений в базовом алгоритме. FlashAttention-2 улучшает распределение работы между параллельными блоками потоков и варпами GPU, сокращает число операций, не являющихся матричным умножением (которые сравнительно дороги на тензорных ядрах GPU, оптимизированных под матричное умножение), и перестраивает порядок циклов так, чтобы больше потоковых мультипроцессоров GPU оставались занятыми на протяжении всего вычисления, особенно на тех нагрузках — таких как вывод с длинной последовательностью при умеренном размере батча, — где оригинальная FlashAttention оставляла заметно недоиспользованными вычислительные ресурсы GPU. Ничто из этого не меняет то, что вычисляется; это второй проход по выжиманию из того же тайлированного, IO-осведомлённого алгоритма большей доли теоретической пропускной способности оборудования.

FlashAttention-3 продолжает ту же линию, нацеливаясь конкретно на архитектуру Hopper от NVIDIA (GPU уровня H100), используя аппаратную асинхронность, которую конструкция FlashAttention-2 не могла задействовать: перекрытие тензорных ядер, выполняющих матричное умножение, с отдельными блоками, перемещающими данные по памяти, вместо строгого чередования вычислений и перемещения данных, плюс нативную поддержку низкоточного исполнения FP8 в тех частях вычисления, где цена по точности пренебрежимо мала. Результат приближается к пиковой пропускной способности, которую тензорные ядра Hopper реально способны выдать на внимании, и это тот же паттерн, что уже установил FlashAttention-2: каждое новое поколение оборудования открывает новые возможности асинхронности и точности, под которые можно переинженерить тот же самый лежащий в основе тайлированный, IO-осведомлённый алгоритм, не меняя того, что именно вычисляется.

Более широкий урок, который стоит извлечь из FlashAttention и FlashAttention-2 вместе взятых, состоит в том, что "стоимость внимания" — это не единое фиксированное число, определяемое исключительно сложностью $O(n^2)$, — это движущаяся цель, форма которой определяется тем, насколько хорошо алгоритм согласован с иерархией памяти и параллелизмом оборудования, на котором он выполняется. Это тема, которая будет повторяться на протяжении оставшейся части этой книги: архитектурные решения и системные решения на практике неразделимы, и глава о "стоимости внимания", обсуждающая только асимптотическую сложность, упустила бы половину того, что на самом деле определяет, практично ли обслуживать заданную длину контекста. Имея на столе как маршрут приближения через разреженное/линейное внимание, так и точный IO-осведомлённый маршрут, следующая глава переходит к проблеме, которую эти техники сами по себе не решают: даже при достаточно дешёвом внимании позиционные кодировки модели должны вести себя разумно на длинах, на которых она никогда не обучалась, а это вопрос обобщения, а не вычислительной стоимости.

7. Взгляд с точки зрения собеседования

"Почему self-attention квадратично и что именно растёт по мере роста квадратичного слагаемого?" Сильный ответ определяет конкретно произведение $QK^T$: вычисление попарных оценок совместимости между каждой парой из $n$ токенов обязательно порождает матрицу $(n,n)$, и как её вычисление, так и её хранение стоят $O(n^2)$. Также стоит отметить, что это структурное свойство — оно следует из самого определения "каждый токен сравнивает себя с каждым другим токеном" — и не является артефактом конкретной реализации, что и объясняет, почему FlashAttention способна исправить постоянный множитель и паттерн доступа к памяти, но не меняет базовую асимптотику $O(n^2)$ по FLOPs.

"FlashAttention часто описывают как делающую внимание "бесплатным". Это точно?" Нет, и сильный кандидат оспаривает эту формулировку: FlashAttention сокращает время выполнения и объём памяти за счёт IO-осведомлённости, а не за счёт сокращения числа операций с плавающей запятой (на самом деле она выполняет несколько больше FLOPs из-за пересчёта в обратном проходе). Экономия возникает из осознания того, что внимание на современных GPU ограничено пропускной способностью памяти, поэтому минимизация чтений/записей HBM через тайлинг и слитые (fused) ядра даёт большой реальный прирост скорости, хотя асимптотическая вычислительная сложность не меняется.

"Когда вы бы обратились к паттерну разреженного внимания вроде Longformer, а когда к FlashAttention?" Это решения разных проблем, и сильный ответ прямо это утверждает: FlashAttention — это подключаемая точная оптимизация с нулевой ценой для качества, поэтому по сути никогда нет причины её не использовать. Паттерны разреженного внимания — это подлинное приближение, обменивающее часть моделирующей способности на изменение асимптотической сложности, и оно оправдано именно тогда, когда длины последовательностей достаточно велики, а стоимость памяти/вычислений полного внимания достаточно высока, чтобы даже точная, IO-осведомлённая реализация всё ещё была слишком дорогой, — и когда структура задачи (например, длинные документы, где локальность плюс несколько глобальных якорных токенов правдоподобно улавливают бо́льшую часть релевантных зависимостей) делает конкретный паттерн разреженности разумным индуктивным смещением.

"От чего вы отказываетесь, переходя на линейное внимание, и почему оно не заменило softmax-внимание в передовых моделях?" Сильный ответ называет механизм: линейное внимание заменяет сходство softmax ядерным отображением признаков так, что $\phi(Q)(\phi(K)^TV)$ можно вычислить, никогда не формируя матрицу $(n,n)$, что даёт стоимость $O(n)$ и эквивалентную рекуррентную формулировку для инференса. Цена в том, что большинство практических ядерных отображений признаков дают более гладкие, менее заострённые распределения внимания по сравнению с точным softmax, что, как правило, вредит способности модели к резкому, избирательному поиску по длинным контекстам — возможности, которая на практике оказалась достаточно важной, чтобы большинство передовых моделей придерживались точного (ускоренного FlashAttention) softmax-внимания, а не линейного, оставляя линейное/рекуррентно-подобное внимание для сценариев, где экстремальные длины последовательностей делают этот компромисс оправданным.

"BigBird заявляет тьюринг-полноту для своего паттерна разреженного внимания. Почему это важно и почему одно лишь локальное скользящее окно не обладает этим свойством?" Сильный ответ связывает это со связностью графа: рассматривая паттерн внимания как граф над токенами, чисто локальное окно сохраняет короткие информационные пути только между близкими токенами, так что двум далёким друг от друга токенам может понадобиться много слоёв, чтобы обменяться информацией, что фактически ограничивает выразительность модели для задач, требующих дальних зависимостей. Комбинация локальных, глобальных и случайных рёбер в BigBird поддерживает малый диаметр графа (любые два токена могут достичь друг друга за немного "прыжков"), что авторы используют, чтобы доказать, что их разреженный паттерн сохраняет ту же теоретическую выразительную силу (универсальную аппроксимацию, тьюринг-полноту), что и полное внимание, в отличие от наивного оконного паттерна.

8. Вопросы для самопроверки

  1. Отталкиваясь от определения масштабированного скалярного внимания, объясните точно, какая операция даёт стоимость $O(n^2)$ и почему её нельзя избежать более умной реализацией того же вычисления.
  2. Если удвоить длину контекста модели, использующей стандартное полное внимание, примерно во сколько раз вырастут вычисления и память для внимания и почему это делает "просто увеличить окно контекста" нетривиальным инженерным решением?
  3. Объясните, чем отличаются конструкция Longformer "локальное окно плюс глобальные токены" и конструкция BigBird "локальное плюс глобальное плюс случайное", и почему именно случайный компонент BigBird улучшает поток информации по последовательности.
  4. Пройдите через алгебраическую перестановку, которая делает линейное внимание линейным по длине последовательности, и объясните, почему эта перестановка также даёт эквивалентную рекуррентную (RNN-подобную) формулировку, полезную во время инференса.
  5. Что значит, что внимание "ограничено пропускной способностью памяти", а не "ограничено вычислениями" на GPU, и как этот факт мотивирует весь замысел FlashAttention?
  6. Почему FlashAttention описывается как вычисляющая "тот же результат быстрее", а не как приближение, и что это различие подразумевает относительно того, меняет ли она выходы модели или динамику обучения?
  7. Назовите один аспект, в котором FlashAttention-2 улучшает FlashAttention, и объясните, почему это улучшение касается использования оборудования, а не базового тайлированного/IO-осведомлённого алгоритма.

9. Источники

  • Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150. https://arxiv.org/abs/2004.05150
  • Zaheer, M., et al. (2020). Big Bird: Transformers for Longer Sequences. NeurIPS 2020. arXiv:2007.14062. https://arxiv.org/abs/2007.14062
  • Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020. arXiv:2006.16236. https://arxiv.org/abs/2006.16236
  • Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135. https://arxiv.org/abs/2205.14135
  • Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. arXiv:2307.08691. https://arxiv.org/abs/2307.08691
  • Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608. https://arxiv.org/abs/2407.08608

Тест для самопроверки

Проверьте понимание главы с помощью короткого теста — вопросы по теории и небольшие расчёты.

Какая именно операция в масштабированном скалярном внимании порождает стоимость $O(n^2)$?

Объяснение: Q имеет форму (n,d), а Kᵀ — форму (d,n), так что QKᵀ — это матрица (n,n), по одной записи на пару токенов. Это сравнение всех со всеми структурно, а не причуда реализации.

В чём ключевое различие между тем, как разреженное внимание (Longformer/BigBird) и линейное внимание снижают стоимость внимания?

Объяснение: Разреженное внимание точно вычисляет выбранное подмножество пар. Линейное внимание переставляет порядок в φ(Q)(φ(K)ᵀV), так что объект (n,n) никогда не материализуется, снижая стоимость до O(nd²).

Почему FlashAttention ускоряет внимание, вообще не меняя его математический результат?

Объяснение: FlashAttention фактически выполняет несколько больше FLOPs (из-за пересчёта на обратном проходе), но значительно меньше обращений к HBM — поскольку узкое место именно в пропускной способности, а не в арифметике, это даёт большое реальное ускорение.

Почему важен случайный компонент внимания BigBird, помимо того, что уже даёт паттерн «локальное окно плюс глобальные токены»?

Объяснение: Случайные рёбра сокращают диаметр графа, что авторы BigBird используют, чтобы показать, что их разреженный паттерн сохраняет теоретическую выразительность полного внимания (универсальная аппроксимация, тьюринг-полнота).

Self-attention вычисляет матрицу $(n,n)$ попарных оценок. Для последовательности из $n = 1000$ токенов, сколько записей содержит эта матрица?

Объяснение: 1000 × 1000 = 1 000 000 попарных записей.

Если удвоить длину контекста модели, во сколько раз вырастут вычисления и память внимания, учитывая, что стоимость масштабируется как $O(n^2)$?

Объяснение: Удвоение n даёт (2n)² = 4n² — рост в 4 раза, а не в 2.

Переход от контекста в 2 тыс. токенов к 128 тыс. токенов — это увеличение длины последовательности в 64 раза. Примерно во сколько раз вырастут вычисления внимания, учитывая, что стоимость масштабируется как $O(n^2)$?

Объяснение: 64² = 4096 — последовательность длиннее в 64 раза стоит примерно в 4096 раз дороже по вычислениям и памяти внимания.

Longformer использует локальное окно размера $w = 512$ для последовательности длины $n = 100{,}000$. Примерно во сколько раз полное внимание $O(n^2)$ дороже локального внимания Longformer $O(n \cdot w)$ — то есть чему равно $n / w$? (Округлите до целого числа.)

Объяснение: 100 000 / 512 ≈ 195.3, то есть полное внимание стоит примерно в 195 раз дороже, чем локальная часть внимания Longformer.