Kimi Delta Attention (KDA)

Kimi Delta Attention - это вариант линейного внимания, используемый в kimi-k3 и Kimi Linear. Если взглянуть на него без контекста, обновление состояния выглядит произвольным:

$$ \widetilde S_t = S_{t-1}\operatorname{Diag}(\alpha_t), \quad \hat v_t = \widetilde S_t k_t, \quad e_t = \beta_t(v_t - \hat v_t), \quad S_t = \widetilde S_t + e_t k_t^{\mathsf T}, \quad o_t = S_t,d_k^{-1/2} q_t $$

Вывод Джейми Дборина (Jamie Dborin) из Doubleword восстанавливает эту формулу шаг за шагом, формулируя по одному требованию за раз вдоль цепочки softmax attention -> linear attention -> DeltaNet -> Gated DeltaNet -> KDA. Каждый шаг - это конкретная претензия к предыдущему. В оригинале используется нотация бра-кет; формулы здесь приведены в стандартном векторном виде, где $S$ отображает пространство ключей в пространство значений, то есть размерность $S$ составляет $d_v \times d_k$, а $k_t$ является единичным вектором.

Отказ от softmax превращает внимание в состояние

Каузальное softmax-внимание вычисляет $T^2$ оценок совместимости ключей и запросов для последовательности длины $T$. Кэширование ключей и значений избавляет от повторных вычислений, но кэш всё равно растёт вместе с последовательностью, и каждый новый запрос вынужден просматривать его целиком.

Препятствием для перегруппировки вычислений служит знаменатель softmax, который одновременно зависит и от текущего запроса, и от всех предыдущих ключей. Если убрать его и включить масштабный коэффициент $d_k^{-1/2}$ в запрос, внимание принимает вид:

$$ o_t = \sum_{i \le t} (k_i^{\mathsf T} q_t), v_i = \Big(\sum_{i \le t} v_i k_i^{\mathsf T}\Big) q_t $$

Весь трюк в том, что $(v k^{\mathsf T})q = (k^{\mathsf T}q),v$: внешнее произведение даёт матрицу, а скалярное - число. Вся информация о прошлом сжимается в одну матрицу фиксированного размера, а внимание превращается в операцию записи с последующим чтением:

$$ S_t = S_{t-1} + v_t k_t^{\mathsf T}, \qquad o_t = S_t q_t $$

Вычислительная сложность теперь линейна по длине последовательности, а память - и именно это делает подход применимым для контекста в 1M токенов - остаётся постоянной, а не растёт с каждым токеном, как в случае с KV-кэшем. Платой за это стала потеря нормализации и избирательности softmax.

Сложение - это не присваивание

Запись в линейном внимании устроена с изъяном. Запишем $v_t k_t^{\mathsf T}$, а затем сразу же запросим новое состояние с тем же ключом:

$$ S_t k_t = S_{t-1}k_t + v_t\underbrace{(k_t^{\mathsf T}k_t)}{1} = S{t-1}k_t + v_t $$

Память возвращает не $v_t$. Она возвращает $v_t$ плюс то, что возвращала ранее. Если старое состояние уже было корректным, новое теперь вернёт удвоенное значение. А поскольку ключи не ортогональны друг другу, каждая запись просачивается во все ранее сохранённые ассоциации. Обновление ведёт себя как += там, где требовалось =.

DeltaNet: запись ошибки вместо значения

DeltaNet исправляет операцию записи. Перед сохранением токена $t$ мы запрашиваем у памяти, что именно она сейчас связывает с новым ключом, $\hat v_t = S_{t-1}k_t$, и добавляем только разницу, умноженную на обучаемую силу записи $\beta_t \in [0,1]$:

$$ e_t = \beta_t(v_t - S_{t-1}k_t), \qquad S_t = S_{t-1} + e_t k_t^{\mathsf T} $$

Считывая тот же ключ обратно, получаем:

$$ S_t k_t = (1-\beta_t)S_{t-1}k_t + \beta_t v_t $$

При $\beta_t = 1$ чтение возвращает в точности $v_t$; меньшие значения интерполируют. Коррекция также локальна: для любого $x$, ортогонального $k_t$, выполняется $(S_t - S_{t-1})x = e_t(k_t^{\mathsf T}x) = 0$, поэтому запись ранга один меняет отклик только вдоль выбранного направления ключа и не затрагивает ортогональные направления.

То же самое обновление получается и из совершенно другой постановки. Представим текущую пару как один обучающий пример для линейного отображения $S$ с функцией потерь реконструкции $\mathcal L_t(S) = \tfrac12\lVert Sk_t - v_t\rVert_2^2$. Её градиент равен $(Sk_t - v_t)k_t^{\mathsf T}$ - наглядно внешнее произведение ошибки в пространстве значений на ключ, для которого эта ошибка возникла. Один шаг градиентного спуска с шагом $\beta_t$ из точки $S_{t-1}$ в точности воспроизводит дельта-правило. Таким образом, $\beta_t$ - это сила замещения, если рассматривать состояние как память, или скорость обучения (learning rate), если трактовать процесс как онлайн-обучение; само же изменение в обоих случаях является обновлением ранга один.

Раскрыв скобки, получаем DeltaNet в виде структурированного перехода состояний:

$$ S_t = S_{t-1}\big(I - \beta_t k_t k_t^{\mathsf T}\big) + \beta_t v_t k_t^{\mathsf T} $$

Для единичного ключа матрица $I - \beta_t k_t k_t^{\mathsf T}$ имеет собственное значение $1-\beta_t$ вдоль текущего ключа и $1$ во всех остальных направлениях. Она стирает старую ассоциацию по этому ключу перед записью новой.

Gated DeltaNet: иногда стираться должно всё

DeltaNet исправляет запись, но не время жизни записанного. Как только токен попадает в $S_t$, чтение $S_t q = \sum_{i\le t}(k_i^{\mathsf T}q)v_i$ уже не может его пропустить. Дельта-правило наводит порядок только вокруг текущего ключа; устаревшее содержимое в других направлениях продолжает влиять на каждое последующее чтение.

Gated DeltaNet добавляет обучаемый скалярный вентиль удержания (retention gate) $\alpha_t \in [0,1]$, который применяется до запуска дельта-правила:

$$ \widetilde S_t = \alpha_t S_{t-1} \ (\text{forget}), \quad \hat v_t = \widetilde S_t k_t \ (\text{predict}), \quad e_t = \beta_t(v_t - \hat v_t) \ (\text{correct}), \quad S_t = \widetilde S_t + e_t k_t^{\mathsf T} \ (\text{write}) $$

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

Проблема остаётся в том, что $\alpha_t$ - это одно число для всей матрицы. Каждый канал ключей вынужден сохраняться или затухать с одинаковой скоростью.

KDA: отдельный вентиль на каждый канал ключа

KDA превращает скаляр в вектор $\alpha_t \in [0,1]^{d_k}$ и помещает его на диагональ, $D_t = \operatorname{Diag}(\alpha_t)$. Поскольку состояние отображает ключи в значения, каналы ключей - это столбцы $S$, поэтому умножение справа независимо затухает каждый из них:

$$ \widetilde S_t = S_{t-1}D_t $$

Всё остальное - уже выведенное ранее дельта-правило. Концептуально всё отличие от Gated DeltaNet сводится к замене $\alpha_t \longrightarrow \operatorname{Diag}(\alpha_t)$, что позволяет очищать один канал, сохраняя при этом другой.

Раскрыв рекуррентность, можно увидеть структуру этого перехода:

$$ S_t = S_{t-1}\underbrace{D_t\big(I - \beta_t k_t k_t^{\mathsf T}\big)}_{A_t} + \beta_t v_t k_t^{\mathsf T}, \qquad A_t = D_t - (D_t k_t)(\beta_t k_t)^{\mathsf T} $$

$A_t$ представляет собой диагональную матрицу минус матрицу ранга один - diagonal-plus-low-rank (DPLR). Стоит чётко разделять размерности: DPLR описывает переход $d_k \times d_k$, действующий в пространстве ключей, тогда как сама память по-прежнему остаётся матрицей $S_t$ размерности $d_v \times d_k$.

Сжатая сводка по семейству из четырёх шагов:

В реализациях хранится $g_t = \log\alpha_t \le 0$, а коэффициенты удержания восстанавливаются как $\exp(g_t)$. В транспонированной раскладке $d_k \times d_v$, используемой в эталонном коде, вся рекуррентность укладывается в пять строк:

state = state * g_t.exp().unsqueeze(-1)
prediction = einsum("bhkv,bhk->bhv", state, k_t)
residual = beta_t.unsqueeze(-1) * (v_t - prediction)
state = state + einsum("bhk,bhv->bhkv", k_t, residual)
output = einsum("bhk,bhkv->bhv", q_t * scale, state)

Два расписания вычислений для одной и той же рекуррентности

Вторая половина статьи посвящена исполнению, и суть в том, что рекуррентный и блочный (chunkwise) kernel'ы KDA - это не два разных механизма, а два расписания вычислений, дающие идентичные состояния и результаты.

Fused recurrent - это буквальная транскрипция: одна программа Triton на последовательность, value-голову и тайл значений шириной 32 (BV = 32, BK = next_power_of_2(K)), где каждая владеет тайлом состояния [BK, BV] и последовательно перебирает токены в цикле. Предсказание и чтение - это редукции, запись - внешнее произведение. Это хорошо подходит для фазы декодирования (decode), когда за раз обрабатывается только один токен, но плохо подходит для обучения, поскольку векторные операции не превращаются в крупные матричные умножения, необходимые тензорным ядрам.

Chunkwise обрабатывает $C$ токенов одновременно и должен воспроизводить тот же результат, несмотря на то что ошибка каждого токена зависит от записей, сделанных ранее внутри того же блока. Решение состоит в том, чтобы сначала вычислить предварительные ошибки в предположении, что внутриблочных записей не было:

$$ \bar e_i = \beta_i\big(v_i - S_c D_{0:i}k_i\big) $$

а затем скорректировать их. Каждый член коррекции представляет собой скалярный коэффициент, умноженный на более раннюю ошибку, поэтому определим:

$$ \rho_{ij} = \beta_i, k_j^{\mathsf T} D_{j+1:i}, k_i, \qquad j < i $$

где $D_{j+1:i}$ - накопительное затухание, переносящее запись в позиции $j$ вперёд к позиции $i$. Объединение $\rho_{ij}$ в строго нижнетреугольную матрицу $R_c$ превращает всю цепочку зависимостей в одну линейную систему: $E_c = \bar E_c\big((I + R_c)^{-1}\big)^{\mathsf T}$. Плотное обращение матриц здесь не требуется: $I + R_c$ треугольная с единичной диагональю, так что задача сводится к каузальному решению треугольной системы (triangular solve) для каждого канала значений.

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

$$ S_{c+1} = S_c D_{0:C-1} + E_c K_c^{\mathrm{end}} $$

где $K_c^{\mathrm{end}}$ содержит ключи, затухающие к границе блока. Выходные значения формируются второй нижнетреугольной матрицей коэффициентов $A^{qk}c$, построенной из $\chi{ij} = s,k_j^{\mathsf T}D_{j+1:i}q_i$ для $j \le i$ (диагональ включена, так как KDA читает после записи, и токен $i$ видит собственную запись):

$$ O_c = s,S_c Q_c^{\mathrm{boundary}} + E_c\big(A^{qk}_c\big)^{\mathsf T} $$

Первое слагаемое считывает затухшее входящее состояние, второе - добавляет каузальный вклад записей внутри блока.

Реальная реализация на Triton представляет собой конвейер запусков, а не один kernel. Сначала вычисляются локальные для блока накопительные логарифмические затухания, благодаря чему $D_{j+1:i}$ получается из разности двух префиксных сумм вместо перемножения цепочки векторов удержания; затем строятся матрицы взаимодействия $A^{qk}$ и $A^{kk}$, где $A^{kk}$ формирует WY-подобное представление скорректированных записей блока; запускается один state kernel, выполняющий единственное сканирование между блоками; после чего output kernel вычисляет все блоки и тайлы параллельно. В готовом коде выполнено дополнительное слияние (fusion): диагональные блоки взаимодействия по 16 токенов вычисляются до объединённого kernel'а внедиагональных элементов и треугольного решения.

Почему эта архитектура полезна за пределами дата-центров

Рекуррентное состояние KDA компактно и фиксировано по размеру, что делает обработку длинного контекста на скромном оборудовании в принципе реализуемой. deltafin запускает KDA-слои K3 на CPU во время декодирования именно потому, что состоянию там комфортнее, чем в цепочке вызовов на GPU, а его чистый PyTorch-shim для доступных только под CUDA kernel'ов fla сходится с блочным путём с точностью порядка 1e-9. Сама Moonshot объясняет возможность запуска на ноутбуке именно заслугой KDA, а не какими-либо наработками по квантованию.

Qwen в этом семействе остановилась на шаге Gated DeltaNet: qwen3-8-flash-next строит каждый из своих 12 блоков из трёх слоёв Gated DeltaNet и одного слоя полного внимания (full-attention).