Контекст на целый день лучше десятиминутного? Вспышка внимания и вопрос о длине последовательности

Вот вопрос, на который призвана ответить эта статья: если бы преобразователь мог отслеживать весь торговый день с разрешением в одну секунду вместо десятиминутного окна, стал бы он прогнозировать лучше?

До недавнего времени вы даже не могли спросить. Стандартные потребности во внимании $\mathcal{O}(N^2)$ памяти, поэтому день с 23 400 шагами при 12 головах в float16 требует примерно 12,9 ГБ только для матрицы оценок — больше, чем параметры модели, и больше, чем вам дадут большинство карт. Вопрос был закрыт арифметикой еще до того, как его успели проверить.

Его открывает Flash Attention (Дао и др., 2022). Не путем аппроксимации внимания — оно вычисляет точный тот же результат — а путем реструктуризации вычислений с учетом IO, минимизируя трафик между уровнями памяти графического процессора. Это действительно интересный контент, и большая часть статьи посвящена тому, как он работает: мозаика, повторение онлайн-softmax, $\Theta(N^2 d^2 / M)$ Привязка ввода-вывода и обратный пересчет.

Но механизм – это возможность, а не требование. «Чем дольше контекст, тем лучше» — это эмпирическое утверждение о рынках и постоянной позиции этого блога — из Temporal Fusion Transformers, который обнаружил, что ванильные преобразователи в моделях переобучения финансовых рядов и рекуррентных моделях с коротким периодом ретроспективного анализа остаются конкурентоспособными на высокой частоте - идет в другую сторону. Итак, статья завершается измерением, а не механизмом.

Почему внимание привязано к памяти

Внимание вычисляет $\text{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d_k})\mathbf{V}$ — сам примитив в контексте торговли описан в Трансформаторы временного слияния для прогнозирования многогоризонтного портфеля. Вся проблема заключается в одной строке: матрица промежуточных оценок. $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top/\sqrt{d_k}$ является $N \times N$, он записывается в память, считывается обратно для softmax, записывается снова и снова читается для окончательного matmul — и его необходимо сохранить для обратного распространения ошибки.

Интенсивность арифметических операций внимания равна $\approx \min(d, N)$, то есть около 64 флоп/байт при $d = 64$ — значительно левее точки хребта А100. Он расположен на наклонном потолке пропускной способности, а не на плоском потолке вычислений: графический процессор тратит больше времени на перемещение $\mathbf{S}$ вокруг, чем умножать что-либо. Используемая здесь структура линии крыши — точка гребня, наклонный или плоский потолок и то, почему одни и те же рассуждения решают, стоит ли вообще покупать графический процессор — построена на измеренных цифрах в Когда графический процессор окупается.

Иерархия памяти, которую использует алгоритм

SRAM примерно в 10 раз увеличивает пропускную способность и в 100 раз снижает задержку, при тысячной доли емкости. Все, что делает Flash Attention, следует из этого обмена: отказаться от емкости, купить полосу пропускания и задержку. Тот же шаг «реструктуризация алгоритма вместо покупки оборудования», измеренный в ходе бэктеста ЦП, представляет собой лестницу скорости бэктеста.

Алгоритм Flash Attention

Flash Attention обрабатывает внимание в виде плиток, размер которых соответствует размеру SRAM, и никогда не материализует его полностью. $N \times N$ матрица в HBM вообще.

Раздел $\mathbf{Q}$ в $T_r = \lceil N/B_r \rceil$ рядные блоки и $\mathbf{K}, \mathbf{V}$ в $T_c = \lceil N/B_c \rceil$ колонные блоки, с $B_r, B_c$ выбрал так, чтобы плитка плюс ее аккумуляторы поместились на чипе. Для каждого блока запроса выполните итерацию по всем блокам «ключ-значение»:

For each query block Q_i:
    Initialize: O_i = 0, m_i = -inf, l_i = 0    # output, running max, running sum
    For each KV block (K_j, V_j):
        1. Load Q_i, K_j, V_j from HBM to SRAM
        2. Compute S_ij = Q_i @ K_j^T / sqrt(d)  # in SRAM
        3. Compute local max: m_ij = rowmax(S_ij)
        4. Compute P_ij = exp(S_ij - m_ij)        # in SRAM
        5. Compute local sum: l_ij = rowsum(P_ij)
        6. Update running statistics:
           m_new = max(m_i, m_ij)
           l_new = l_i * exp(m_i - m_new) + l_ij * exp(m_ij - m_new)
           O_i = O_i * (l_i * exp(m_i - m_new) / l_new)
                + P_ij @ V_j * (exp(m_ij - m_new) / l_new)
           m_i = m_new, l_i = l_new
    Write O_i to HBM

Повторение онлайн-софтмакса

Уловка, которая делает возможным тайлинг, — это онлайн-softmax. Наивному softmax требуется два прохода по строке: один для нахождения максимума (для численной стабильности), второй для возведения в степень и нормализации. Два прохода по строке, которую вы отказываетесь сохранить, являются противоречием, поэтому Flash Attention продолжает обрабатывать статистику и масштабировать ее по ходу дела.

После блоков $1, \ldots, j$:

и выходной аккумулятор корректируется тем же коэффициентом:

Каждый раз, когда новый блок увеличивает рабочий максимум, ранее накопленный выход задним числом масштабируется на $e^{m^{(j-1)} - m^{(j)}}$ — как будто новый максимум был известен с самого начала. Результат алгебраически идентичен двухпроходному softmax. В точной арифметике это не приближение; это реассоциация. (В конечной точности это другой путь округления, что имеет значение — см. проверку точности ниже.)

сложность ввода-вывода

Это официальное заявление о победе. Flash Attention работает

Доступы HBM, где $M$ размер SRAM по сравнению с $\Theta(Nd + N^2)$ для стандартной реализации. Обратите внимание, что $M$ появляется в знаменателе: чем больше блокнот на кристалле, тем меньше циклов обработки, поэтому алгоритм формулируется с точки зрения иерархии памяти, а не подсчета FLOP. Для типичного $d = 64$ и $M \approx 100$ КБ, соотношение в пользу Flash Attention примерно в 5-10 раз меньше обращений.

Обратный проход: перерасчет вместо сохранения

Обратное распространение ошибки через внимание обычно требует $\mathbf{P}$ матрица, которую прямой проход просто отказался сохранить. Flash Attention пересчитывает плитки из $\mathbf{Q}, \mathbf{K}$ во время обратного прохода, сохраняя только выходные данные $\mathbf{O}$ и статистика softmax $(m, \ell)$ - оба $\mathcal{O}(N)$, нет $\mathcal{O}(N^2)$. Он обменивает скромное количество избыточной арифметики на термин памяти, который и был всей проблемой. Это та же сделка, что и контрольная точка градиента, применяемая при детализации плитки внутри одного оператора.

FA2: параллелизм

Flash Attention 2 (Дао, 2023) сохранил алгоритм и исправил расписание:

1. Меньше FLOP, не относящихся к Matmul. FA1 тратил реальное время на изменение масштаба, поиск максимального значения и возведение в степень — операции, которые выполняются на ядрах CUDA, а не на тензорных ядрах. FA2 откладывает изменение масштаба до конца внутреннего цикла.
2. Параллелизм по длине последовательности. FA1 распараллеливает только пакеты и головы. FA2 также распараллеливает блоки запросов. Это особенно важно для торгового случая, где у вас часто есть одна очень длинная последовательность на актив и размер пакета 1-4 — именно тот режим, при котором пакетный параллелизм истощает графический процессор.
3. Разделение работы по деформации. Каждая деформация использует отдельное подмножество блоков запроса вместо разделения вычисления оценки и сокращения по деформациям, удаляя перекрестное сокращение.

Заявленный результат: ~70% теоретических пиковых значений FLOP на A100 против ~35% на FA1.

FA3: Механика бункера

Flash Attention 3 (Дао, Шах, 2024 г.) зависит от архитектуры H100:

1. Специализация асинхронной деформации. Тензорный ускоритель памяти Хоппера (TMA) перемещает HBM→SRAM асинхронно. FA3 разделяет варпы на производителей, выдающих загрузку TMA для следующего блока KV, и потребителей, вычисляющих текущий блок, поэтому перемещение данных скрывается за арифметикой.
2. Чередование matmul и softmax. $\mathbf{Q}_i\mathbf{K}_j^\top$ поскольку один блок работает на тензорных ядрах, в то время как softmax для предыдущего блока работает на ядрах CUDA — двух разных аппаратных модулях, действительно параллельных, а не разделенных по времени.
3. FP8 с некогерентной обработкой. H100 выполняет FP8 с пропускной способностью, вдвое превышающей FP16. Наивное внимание FP8 разрушается выбросами; FA3 случайным образом вращает векторы перед блочным квантованием, чтобы распределить величину выбросов по координатам, что приводит к снижению числовой ошибки в 2,6 раза по сравнению с простым FP8.

Причинно-следственная маскировка — это то, где торговля получает скидку

Причинно-следственная маскировка обязательна для временных рядов — модель не должна учитывать будущее — и при мозаике это не дополнительные затраты, а экономия. Любой тайл, ключи которого полностью находятся в будущем относительно его запросов, начисто пропускается, никогда не загружается и никогда не вычисляется, что сокращает примерно половину работы. В PyTorch это is_causal=True; больше ничего не требуется.

Интеграция состоит из восьми строк

Почти ни один код, который вам нужен, не касается Flash Attention. Замените явный путь матрицы оценок для объединенного ядра:

scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_head)
scores = scores.masked_fill(causal_mask, float("-inf"))
out = torch.matmul(torch.softmax(scores, dim=-1), v)

out = torch.nn.functional.scaled_dot_product_attention(
    q, k, v,
    dropout_p=self.dropout if self.training else 0.0,
    is_causal=True,
)

Вот и все изменение. q, k, v имеют форму (batch, heads, seq, head_dim); причинная маска исчезла, потому что ее создает ядро. Чтобы получить полную аннотированную торговую модель PyTorch, в которую можно поместить это — входную проекцию, блоки и 3-классную голову вверх/плоскую/вниз – используйте DeepLOB, а полный процесс обучения см. в разделе Трансформаторы временного синтеза. Строительство здесь четвертой копии этих лесов ничему не научит.

Требования: вычислительные возможности >= 8,0 (A100, H100, RTX 3090+), входные данные половинной точности, PyTorch >= 2,0. Убедитесь, что быстрый путь действительно связан с torch.backends.cuda.sdp_kernel диагностика и torch.cuda.max_memory_allocated() — SDPA автоматически возвращается к математическому ядру, если какое-либо предварительное условие не выполняется, и автоматический возврат выглядит точно так же, как работающая модель, которая просто работает медленно.

Измерение: окупается ли более длинный контекст?

Все вышесказанное говорит о том, что контекст 32K или 128K теперь доступен. Здесь ничего не говорится о том, полезно ли это. Честный эксперимент:

Обучите ту же архитектуру на $N \in \{512, 4096, 32768\}$ с SDPA-Flash в серии BTC, используемой в других частях этой серии, с фиксированными параметрами, оптимизатором и целью, поэтому длина последовательности является единственной переменной. Сообщите о двух вещах:

1. Стоимость. Измеряется количество настенных часов за эпоху и torch.cuda.max_memory_allocated() в каждом $N$.
2. Преимущество. Прогнозируемая эффективность за пределами выборки по сравнению с $N$, на расколе вперед.

В более раннем проекте этой статьи содержалась таблица показателей памяти на длину последовательности, полученная аналитически из $\mathcal{O}(N)$ формула активации. Эти строки удалены: они никогда не измерялись и не соответствовали арифметике бюджета памяти, приведенной в статье. Производное число, представленное в таблице результатов, является сфабрикованным результатом, и в этом блоге он не публикуется.

Интересное свойство этого эксперимента состоит в том, что его можно опубликовать в любом направлении. Если производительность за пределами выборки монотонно возрастает с $N$, что оправдывает всю программу с длинным контекстом. Если он стабилизируется на нескольких тысячах шагов или ухудшается, это более сильная часть — компаньон честному негативу — и это означало бы, что стена памяти никогда не была обязательным ограничением для торговых преобразователей.

Больше контекста — больше емкости, следовательно, больше поверхности для переобучения

Существует конкретная причина ожидать неизменного или отрицательного результата. Трансформаторы временного синтеза уже документально подтверждает, что ванильные преобразователи наивно применялись к переобучению финансовых рядов — у них отсутствуют временные индуктивные смещения, а рекуррентные модели с коротким обзором назад остаются конкурентоспособными на высокой частоте. Расширение контекста с 512 до 32 768 шагов не добавляет информации, пропорциональной длине; предельное 32-тысячное отставание почти эффективного ценового ряда несет в себе очень мало. Что он надежно добавляет, так это количество подходящих параметров.

Итак, зачистка $N$ следует рассматривать как то, что есть: поиск по выбору модели, с тем же механизмом, который этот блог применяет к любому другому поиску. Три длины последовательности, умноженные на все остальное, являются подсчетом попыток, и победитель должен очистить Сниженный коэффициент Шарпа, рассчитанный на основе этого количества попыток и ворот PBO, а не просто победить своих соседей. В противном случае «победа в длинном контексте» неотличима от выбора лучшего из трех шумных серий.

Проверка точности, потому что «точно» выполняет большую работу

Flash Attention является точным в точной арифметике. Приложенная к нему рекомендация — запускать в fp16 или bf16, а на H100 рассматривать FP8 — нет. Это отдельные утверждения, и на практике доминирует второй: повторное связывание суммы и понижение точности до половинной — оба возмущения, и статья, в которой введена гарантия упорядочения, не должна затем отмахиваться от точности.

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

• Вычисление внимания с помощью SDPA-Flash в bf16 и эталонной реализации fp64 на идентичных входах; сообщите максимальную относительную ошибку выходного тензора.
• Доведите дело до принятия решения: для модели, выдающей метку вверх/вниз/вниз, сообщите количество меток, переключающихся между двумя путями, как долю от общего числа решений.

Небольшое, ограниченное и объяснимое разногласие является признаком правильного быстрого пути. Неограниченный означает, что рекомендация FP8 никогда не была безопасной для этой модели. Ни одно из чисел неизвестно, пока оно не будет запущено.

Когда к этому обратиться

Сжато до решения, которое имеет ту же форму, что и руководство по принятию решений графического процессора:

• Свыше ~2000 временных шагов на графическом процессоре CUDA: да, безоговорочно. Это однострочное изменение, дающее точный результат, и выигрыш растет с увеличением $N$. Не существует сценария, в котором вы бы хотели, чтобы материализовалось $N \times N$ вместо этого путь.
• **Менее ~512 временных шагов, на ЦП или с архитектурами, не требующими внимания (CNN, SSM, такие как Mamba): не имеет значения. ** Слева от гребня фиксированные накладные расходы — это все затраты, и внимание никогда не было вашим узким местом.
• Вышеуказанные пороги — это фольклор, а не измерения — они взяты из общей литературы, а пересечение вашей собственной модели и карты — это десятистрочный ориентир. Запустите его, а не доверяйте круглым цифрам.

Заключение

Flash Attention — это чистый и действительно важный результат: соблюдая иерархию памяти и повторно связывая softmax, он вычисляет точное внимание с $\mathcal{O}(N)$ память вместо $\mathcal{O}(N^2)$и $\Theta(N^2 d^2 / M)$ IO-bound объясняет, почему. Внедрение его в торговый преобразователь представляет собой однострочное изменение без затрат на точность точных арифметических действий и с большим выигрышем в памяти.

Чего он не делает, так это отвечает на вопрос вверху. Он преобразует «контекст на целый день невозможен» в «контекст на целый день стоит дешево», что является изменением стоимости эксперимента, а не его результата. Падение стены памяти — это приглашение к измерению, и именно измерение превращает это из бумажного резюме в открытие.

Ссылки

1. Дао Т., Фу Д.Ю., Эрмон С., Рудра А., Ре, К. «FlashAttention: быстрое и эффективное в использовании памяти точное внимание с распознаванием ввода-вывода». НейрИПС (2022). arXiv:2205.14135
2. Дао, Т. «FlashAttention-2: более быстрое внимание с лучшим параллелизмом и разделением работы». ICLR (2024). arXiv:2307.08691
3. Дао Т., Шах Дж. «FlashAttention-3: быстрое и точное внимание с асинхронностью и низкой точностью». НейрИПС (2024). arXiv:2407.08608
4. Васвани А. и др. «Внимание — это все, что вам нужно». НейрИПС (2017).
5. Милаков М., Гимельшейн Н. «Онлайн-расчет нормализатора для softmax». arXiv:1805.02867 (2018).