Какая вычислительная сложность у self-attention и почему

Что проверяют

Понимаете ли вы, откуда берётся квадрат и во что он упирается на практике: во время, в память или в то и другое. Следом обычно спрашивают, как с этим живут модели с длинным контекстом и что меняет FlashAttention.

Как отвечать

Разложите слой для последовательности длины n и размерности d:

• проекции Q, K, V и выходная: умножения матрицы n × d на d × d, это O(n·d²);
• матрица внимания: Q·Kᵀ даёт n × n оценок, каждая это скалярное произведение длины d, итого O(n²·d). Столько же стоит умножение весов на V;
• память: наивная реализация хранит матрицу n × n на каждую голову и слой, а при обучении ещё и для обратного прохода.

Итого O(n²·d + n·d²). Многоголовость порядок не меняет: h голов по d/h. Отсюда вывод: пока n заметно меньше d, основную стоимость дают проекции и FFN, а квадратичный член доминирует на длинных контекстах.

Что с этим делают:

• FlashAttention: считает внимание блоками в быстрой памяти GPU и не материализует матрицу целиком. Операций всё те же O(n²·d), но память линейна, а скорость растёт за счёт меньшего обмена с HBM;
• разреженное внимание: локальное окно плюс глобальные токены (Longformer, BigBird, sliding window), O(n·w·d) при окне w;
• аппроксимации: Linformer сжимает K и V по длине, Performer приближает softmax ядрами, обе линейны по n;
• KV-кеш на инференсе: новый токен стоит O(n·d), но кеш растёт линейно, его урезают через MQA и GQA.

Где спотыкаются

Отвечают «O(n²)» и забывают про d и проекции. Тогда не получается объяснить, почему на коротких последовательностях внимание вовсе не узкое место.

Считают, что FlashAttention делает внимание линейным по времени. Он точный, а выигрывает на памяти и работе с ней.

Путают обучение и генерацию: при авторегрессивном выводе с кешем квадрат размазан по шагам, а главной проблемой становится объём KV-кеша.