표기법에 관한 참고 사항: 이 글은 (양자 역학에서 영감을 받은 제 개인적인 의견으로는) 이 유도 과정의 형태를 매우 명확하게 만들어 주기 때문에 브라-켓 표기법 (bra-ket notation)을 기본으로 사용합니다. 위의 수학 표기법 전환 기능을 사용하면 모든 방정식을 전통적인 굵은 벡터 (bold vectors)와 명시적인 전치 (transposes)를 사용하여 다시 작성합니다. 브라-켓 모드에서 $|v\rangle$는 열 벡터 (column vector), $\langle v|$는 행 벡터 (row vector), $\langle v|u\rangle$는 숫자이며, $|v\rangle\langle u|$는 행렬 (matrix)입니다. 벡터는 기본적으로 오른쪽을 향하며, 키 (keys)는 선형 어텐션 (linear-attention) 상태로 기록될 때 왼쪽을 향합니다. 우리는 하나의 인과적 어텐션 헤드 (causal attention head)와 실수 값 벡터 (real-valued vectors)를 다루며, DeltaNet의 키가 정규화 (normalized)되었다고 가정하고, 상태 (state)가 키 공간 (key space)에서 값 공간 (value space)으로 매핑된다고 가정합니다.

현대의 선형 어텐션 (linear attention) 변형 모델들은 복잡하며, 언뜻 보기에는 무엇을 달성하기 위해 설계되었는지 파악하기가 쉽지 않습니다. 참고를 위해 여기 Kimi Delta Attention (KDA)의 상태 업데이트 방정식 (state update equation)이 있습니다:

[IMG:1]

이들이 이해하기 어려운 이유는 이들이 지난 몇 년 동안 개발되어 온 선형 어텐션 변형 모델 제품군의 최신 버전이며, 그 복잡성이 필연적으로 팽창하여 외부에서 보기에 최신 변형 모델들이 접근 불가능해 보이기 때문입니다.

이 포스트에서는 최신 Qwen 및 Kimi 모델 제품군에서 사용되는 두 가지 모델을 포함한 DeltaNet 선형 어텐션 변형 모델 제품군을 살펴보고, 은닉 상태 (hidden state)에 대해 간단한 사항들을 단언함으로써 어떻게 여러분도 동일한 방정식에 도달할 수 있었을지를 보여줄 것입니다.

우리가 취할 경로는 다음과 같습니다:

softmax 어텐션 (softmax attention) → 선형 어텐션 (linear attention) → DeltaNet → Gated DeltaNet → KDA

KDA를 유도한 후에야 이를 실행하는 순환적 (recurrent) 및 청크 단위 (chunkwise) Triton 프로그램으로 넘어갈 것입니다.

토큰 $x_t$에 대한 쿼리 (query)의 경우, 일반적인 인과적 소프트맥스 어텐션 (causal softmax attention)은 다음과 같습니다:

[IMG:2]

모든 어텐션 가중치 (attention weight)는 스칼라 (scalar)입니다. 이는 하나의 키 (key)와 하나의 쿼리 (query) 사이의 유사성을 측정하며, 그 후 소프트맥스 (softmax)가 해당 쿼리에 대한 모든 점수를 하나의 분포 (distribution)로 변환합니다. 출력값은 값 벡터 (value vectors)의 가중 합 (weighted sum)입니다.

길이가 $L$인 시퀀스에 대해, $L$개의 키-쿼리 (key-query) 쌍이 존재합니다. 자기회귀 추론 (autoregressive inference) 과정에서 우리는 키 (keys)와 값 (values)을 매번 다시 계산하는 대신 캐시 (cache)할 수 있지만, 캐시는 여전히 시퀀스에 따라 커지며 모든 새로운 쿼리 (query)는 여전히 전체 이력을 검사해야 합니다.

이 계산을 재배열하는 데 있어 장애물은 소프트맥스 (softmax)입니다. 소프트맥스의 분모는 현재의 쿼리 (query)와 이전의 모든 키 (key)에 공동으로 의존합니다. 따라서, 일단 이를 제거해 보겠습니다.

명확성을 위해, 상수 스케일 (constant scale)을 쿼리 (query)에 흡수시킵니다. 그러면 의도적으로 단순화된 어텐션 (attention) 형태는 다음과 같습니다.

[IMG:1]

스칼라 내적 (scalar inner product)은 오른쪽으로 이동할 수 있습니다:

[IMG:2]

과거에 의존하는 모든 것은 이제 고정된 크기의 행렬 $S$로 모을 수 있습니다:

[IMG:3]

그러면 어텐션은 쓰기 (write) 후 읽기 (read)가 이어지는 순환 (recurrent) 구조가 됩니다:

[IMG:4]

다음 항등식 (identity)이

[IMG:5]

전체 트릭의 핵심입니다. 외적 (outer product)은 행렬이고, 내적 (inner product)은 숫자입니다. 우리는 더 이상 모든 과거의 키 (key)와 값 (value)을 저장하지 않습니다. 대신 그들의 합산된 외적 (outer products)을 고정된 크기의 상태 (state) $S$에 저장합니다.

이는 시퀀스 길이에 대해 이차적 (quadratic)이 아닌 선형적 (linear)입니다. 즉, 토큰을 한 번만 스캔하며 매 단계마다 동일한 상태 (state)를 업데이트합니다. 우리는 소프트맥스 (softmax)의 정규화 (normalization)와 선택성 (selectivity)을 포기함으로써 그 효율성을 얻었습니다. 더 정교한 선형 어텐션 (linear-attention) 방법들은 특징 맵 (feature maps)과 정규화 도구 (normalizers)를 사용하지만, 이 꾸밈없는 형태는 DeltaNet을 고안하게 만든 메모리 문제를 여실히 보여줍니다.

키-값 (key-value) 쌍 $(k_t, v_t)$를 작성하고, 즉시 동일한 키 (key)로 새로운 상태를 쿼리 (query)한다고 가정해 봅시다:

[IMG:6]

쓰기 (write) 작업이 메모리를 $v_t$로 되돌리지는 않습니다. 대신 메모리가 이미 반환하고 있던 값에 $v_t k_t^ op$를 더하게 됩니다.

만약 이전 상태가 이미 올바른 값을 생성했다면, 가산적 쓰기 (additive write)는 새로운 상태가 그 값의 두 배를 생성하게 만듭니다. 더 일반적으로, 키 (keys)들은 서로 직교 (orthogonal)하지 않으므로, 모든 쓰기 (write)는 이전의 쓰기 (writes)와 간섭할 수 있습니다. 선형 어텐션 (linear attention)은 우리에게 조밀한 연상 메모리 (associative memory)를 제공했지만, 그 업데이트 방식은 우리가 원하는 = 에 가깝기보다는 += 처럼 동작합니다.

DeltaNet은 무조건적인 선형 어텐션 (linear-attention) 쓰기 (write)를 델타 규칙 (delta-rule) 보정으로 대체합니다. 이를 도출하는 데는 두 가지 유용한 방법이 있습니다.

토큰 (token)을 쓰기 (write) 전에, 메모리 (memory)에게 새로운 키 (key)와 현재 연관된 것이 무엇인지 묻습니다:

만약 메모리가 $v$를 반환하기를 원한다면, 전체 값 (value)을 더해서는 안 됩니다. 오직 차이 (difference)만을 더해야 합니다:

학습 가능한 쓰기 강도 (write strength) $eta$를 도입하고 다음과 같이 정의합니다:

그다음 현재 키 (key)에 이 오차 (error)를 씁니다:

이제 즉시 동일한 키를 읽습니다:

$eta = 1$일 때, 결과는 정확히 $v$가 됩니다. $eta$가 작을수록 이전 예측 (prediction)을 목표값 방향으로 부분적으로 이동시킵니다.

이 보정 (correction)은 키 공간 (key space)에서도 국소적 (local)입니다. 현재 키와 직교 (orthogonal)하는 임의의 쿼리 (query)에 대해,

따라서 랭크-1 (rank-one) 쓰기는 모든 직교 방향은 그대로 둔 채 선택된 키 방향의 응답 (response)만을 변경합니다.

동일한 업데이트가 온라인 학습 (online learning) 목적 함수 (objective)로부터 도출됩니다. 현재의 키-값 (key-value) 쌍을 선형 사상 (linear map) $W$에 대한 하나의 학습 예시 (training example)로 취급합니다:

상태 (state)에 대한 그 그레이디언트 (gradient)는 다음과 같습니다:

이는 명백히 외적 (outer product)입니다: 값 공간 (value-space)의 예측 오차와 해당 오차가 관찰된 키 브라 (key bra)의 곱입니다. $eta$ 크기의 경사 하강법 (gradient-descent) 단계를 수행하면:

이는 즉각적인 재구성 (reconstruction)을 요구함으로써 얻은 업데이트와 정확히 일치합니다. 두 해석은 동일합니다:

• 메모리 연산 (memory operation)으로서, $eta$는 이전 연관 (association)을 얼마나 강하게 대체할지를 제어합니다.
• 온라인 학습 (online learning)으로서, $eta$는 스텝 사이즈 (step size)입니다.
• 선형 대수 (linear algebra)로서, 변화는 랭크-1 외적 (rank-one outer product)입니다.

오차를 확장하면 DeltaNet이 구조화된 상태 전이 (state transition)와 새로운 입력 (new input)의 결합임을 알 수 있습니다:

단위 키 (unit key)에 대해, $W$는 현재 키 방향으로 고윳값 (eigenvalue) $eta$를 가지며, 모든 직교 방향에 대해서는 고윳값 $1$을 가집니다. 이는 새로운 연관을 추가하기 전에 현재 키를 따라 기존의 연관을 제거합니다.

DeltaNet은 쓰기 (write)를 수정합니다. 하지만 아직 상태 (state)의 수명 (lifetime)을 수정하지는 못했습니다.

선형 상태 (linear state)는 전체 이력 (history)을 하나의 행렬 (matrix)로 압축합니다. 읽기 (read) 연산은 해당 토큰이 $S$로 접혀 들어간 (folded into) 이후에는 개별적인 오래된 토큰을 건너뛰도록 선택할 수 없습니다. 쿼리와 겹치는 모든 저장된 방향 (stored direction)이 기여하게 됩니다. 델타 규칙 (delta rule)은 현재 키 주변의 상태를 보정할 수 있지만, 다른 방향에 있는 오래된 정보 (stale information)는 여전히 남아 있어 미래의 읽기 (reads)를 왜곡할 수 있습니다.

따라서 우리는 이전 상태를 사용하기 전에 이를 잊을 수 있는 방법이 필요합니다. 학습 가능한 스칼라 유지 게이트 (learned scalar retention gate)를 $g_t$라고 합시다:

$S_t = S_{t-1} imes g_t$

이 게이트가 적용된 상태에 대해 동일한 델타 규칙 (delta rule)을 실행합니다:

$ ext{prediction}_t = k_t^ op S_t$
$ ext{residual}_t = v_t - ext{prediction}_t$
$S_t = S_t + eta_t ( ext{residual}_t imes k_t^ op)$

이것이 Gated DeltaNet입니다. 순서가 중요합니다. 먼저 잊고, 유지된 상태로부터 예측한 다음, 그 예측을 수정합니다. 만약 잊기 전에 예측했다면, 오차 (error)는 우리가 업데이트하는 메모리와는 다른 메모리를 설명하게 될 것입니다.

재귀 (recurrence)를 확장하면 다음과 같습니다:

$S_t = S_{t-1} imes g_t + eta_t (v_t - k_t^ op (S_{t-1} imes g_t)) k_t^ op$

델타 규칙은 타겟팅된 교체 (targeted replacement)를 제공하고, 스칼라 게이트는 전역적 삭제 (global erasure)를 제공합니다. 이들은 서로 다른 문제를 해결하며 상호 보완적입니다.

하지만 여전히 전체 행렬에 대해 단 하나의 결정만을 내립니다. 모델은 모든 키 채널 (key channel)을 동일한 비율로 유지하거나 잊어야만 합니다.

Kimi Delta Attention은 Gated DeltaNet의 스칼라 유지를 벡터 $g_t$로 대체합니다. 이 벡터를 대각선에 배치합니다:

$G_t = ext{diag}(g_t)$

우리의 상태 (state)는 키를 값으로 매핑하므로, 키 채널은 $S_t$의 열 (columns)입니다. 우측 곱셈 (right-multiplication)을 통해 모든 채널에 서로 다른 유지 계수를 적용합니다:

$S_t = S_{t-1} imes G_t$

그 외의 모든 것은 우리가 이미 유도한 델타 규칙과 동일합니다:

$S_t = S_{t-1} imes G_t + eta_t (v_t - k_t^ op (S_{t-1} imes G_t)) k_t^ op$

이것이 KDA입니다. Gated DeltaNet과 비교했을 때, 개념적인 변화는 오직 스칼라를 벡터로 승격시킨 것뿐입니다.

그 효과는 상당합니다. 한 채널은 비워질 수 있는 반면 다른 채널은 유지될 수 있습니다.

KDA의 수정 사항을 확장하면 다음과 같습니다:

$S_t = S_{t-1} imes G_t + eta_t v_t k_t^ op - eta_t (k_t^ op S_{t-1} G_t) k_t^ op$

키 공간 (key-space)의 전이는 다음과 같습니다:

$S_t = S_{t-1} (G_t - eta_t k_t k_t^ op)$

여기서

$M_t = G_t - eta_t k_t k_t^ op$

따라서 $M_t$는 대각 행렬에서 계수 1인 행렬 (rank-one matrix)을 뺀 것입니다. 즉, 대각선-플러스-저계수 (diagonal-plus-low-rank), 또는 DPLR 전이입니다. "DPLR"은 키 공간에 작용하는 전이를 설명합니다. 메모리 상태 자체는 여전히 행렬 $S_t$입니다.

전체 과정을 이제 간결하게 요약할 수 있습니다:

구현 시에는 보통 $S_t$를 $k_t, v_t$와 함께 저장한 다음, $g_t = ext{exp}( ext{retention factors})$를 통해 유지 계수를 얻습니다. 참조 코드에서 사용된 전치 레이아웃 (transposed layout)에서 재귀는 단 다섯 줄로 이루어집니다:

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

공식 naive_recurrent_kda 참조를 확인하세요.

위의 재귀(recurrence)는 자기회귀 디코딩 (autoregressive decode)을 위한 자연스러운 구현입니다. KDA에는 두 가지 주요 실행 체제 (execution regimes)가 있습니다:

재귀적 Triton 실행은 시퀀스, 가치 헤드, 그리고 32-폭 가치 타일당 하나의 프로그램을 사용합니다:

BK = triton.next_power_of_2(K)
BV = 32
grid = (triton.cdiv(V, BV) * N * HV,)

fused_recurrent_kda_fwd 실행 코드를 참조하세요.

BK는 일반적인 지원 구성에서 키 차원 (key dimension)을 커버합니다. 각 프로그램은 구현된 전치된 상태 (transposed state)의 [BK, BV] 타일을 소유하며 토큰을 순서대로 루프(loop)합니다. 서로 다른 가치 타일, 헤드 및 시퀀스는 독립적으로 실행됩니다.

커널은 재귀 식을 거의 그대로 옮겨 놓은 것과 같습니다:

state *= tl.exp(g_t[:, None])
prediction = tl.sum(state * k_t[:, None], axis=0)
residual = beta_t * (v_t - prediction)
...

예측 (prediction)과 읽기 (read)는 리덕션 (reductions)이며, 쓰기 (write)는 외적 (outer product)입니다. 이는 한 번에 하나의 새로운 토큰만 사용 가능한 디코딩에 매우 탁월합니다. 하지만 이러한 벡터 연산은 텐서 코어 (tensor cores)가 가장 효율적으로 작동하는 대규모 행렬 곱셈 (matrix multiplications)으로 전환되지 않기 때문에, 학습 및 긴 프리필에는 덜 매력적입니다.

이 점이 정확히 동일한 재귀에 대한 두 번째 관점을 모색하게 하는 동기가 됩니다.

청크 단위 (Chunkwise) KDA는 토큰들을 함께 처리합니다. 이는 토큰별 재귀와 정확히 동일한 상태와 출력을 생성해야 하지만, 작업을 행렬 곱셈으로 재구성합니다.

각 청크에 대해 우리는 두 가지 결과가 필요합니다:

• 유입된 상태 (incoming state)가 주어졌을 때, 전체 청크 이후의 상태;
• 청크 내부의 모든 인과적 토큰 출력 (causal token output).

유일한 어려움은 토큰 $t$의 델타 오차 (delta error)가 동일한 청크 내의 이전 토큰들에 의해 수행된 쓰기 작업에 의존한다는 점입니다. 4개의 토큰을 사용한 예시는 이러한 의존성을 명확하게 보여줍니다.

토큰 $t, t+1, t+2, t+3$을 가져와 다음과 같이 정의합니다.

청크 경계와 토큰 사이의 누적 감쇠 (cumulative decay)는

토큰 $t$에서의 쓰기 (write)가 토큰 $t+1$로 전달되는 감쇠는

중간에 개입하는 감쇠가 없을 때의 값입니다. 이 모든 행렬은 대각 행렬 (diagonal matrix)이므로, 서로 교환 법칙 (commute)이 성립합니다.

먼저 각 토큰이 적절히 감쇠된 유입 상태 (incoming state)는 볼 수 있지만, 해당 청크 내부의 다른 쓰기들은 볼 수 없다고 가정해 봅시다:

4개의 토큰에 대해, 이는 4개의 임시 값 공간 오차 켓 (value-space error kets)을 생성합니다:

이들은 병렬로 계산하기 쉽지만, 첫 번째를 제외한 나머지는 모두 틀렸습니다. 동일한 청크 내의 이전 쓰기들도 예측에 기여하기 때문입니다.

토큰 $t$는 청크 내에 이전 쓰기가 없으므로,

토큰 $t+1$은 $ ext{decay}_{t, t+1}$을 통과한 후 토큰 $t$의 쓰기를 봅니다:

토큰 $t+2$는 앞선 두 개의 쓰기를 모두 봅니다:

토큰 $t+3$은 세 개 모두를 봅니다:

모든 괄호는 스칼라 (scalar)입니다. 인과적 키-키 계수 (causal key-key coefficient)를 다음과 같이 정의합니다.

그러면 네 개의 방정식 모두 다음과 같은 간결한 형태를 가집니다.

계수들을 엄격한 하삼각 행렬 (strictly lower-triangular matrix)로 모읍니다:

오차 켓들을 열로 쌓습니다:

$ ext{error}_{t+1}$에 대해서도 마찬가지입니다. 그러면 인과적 치환 (causal substitutions)은 다음과 같습니다.

구현 시 일반적인 밀집 역행렬 (dense inverse)을 형성할 필요는 없습니다. $M$은 대각 성분이 1인 삼각 행렬이므로, 이 연산은 모든 값 채널 (value channel)에 독립적으로 적용되는 인과적 삼각 풀기 (causal triangular solve)입니다.

청크의 끝에서, 유입 상태는 4개의 감쇠를 모두 통과했습니다. 청크 내의 각 쓰기는 자신 이후의 감쇠만을 통과했습니다:

행이 경계 끝에 도착하는 키 (keys)인 행렬을 다음과 같이 정의합니다:

$E$가 오차 켓들을 열로 쌓기 때문에, 네 개의 외적 (outer-product) 쓰기는 하나의 행렬 곱셈이 됩니다:

이것이 첫 번째로 요구되는 청크 결과입니다. 순환 상태 (recurrent state)를 한 번에 4개의 토큰만큼 전진시킵니다.

KDA는 쓰기 후에 읽습니다 (reads after writing). $s_{t}$가 토큰 $t$ 이후의 로컬 상태라면,

네 개의 출력을 전개합니다:

인과적 쿼리-키 계수 (causal query-key coefficient)를 다음과 같이 정의합니다.

그리고 계수들을 하삼각 읽기 행렬 (lower-triangular read matrix)에 배치합니다:

0은 인과성 (causality)을 강제합니다. 대각 성분이 포함된 이유는 토큰이 자신의 쓰기를 수행한 후에 읽기를 하기 때문입니다.

이제 경계 감쇠가 적용된 쿼리 켓 (query kets)들을 열로 쌓습니다,

그리고 출력 켓 (output kets)들도 동일한 방식으로 쌓습니다. 네 가지 출력은 모두 다음과 같습니다.

첫 번째 행렬 곱 (matrix product)은 적절하게 감쇠된(decayed) 유입 상태 (incoming state)를 읽습니다. 두 번째는 청크 (chunk) 내부에서 수행된 쓰기 (writes)의 인과적 기여 (causal contribution)를 더합니다. 이것이 두 번째로 요구되는 청크 결과입니다.

청크 단위 구현 (chunkwise implementation)은 위의 방정식들을 하나의 거대한 커널 (monolithic kernel) 대신 커널 실행 (kernel launches)의 파이프라인으로 변환합니다.

먼저 청크 로컬 누적 로그 감쇠 (chunk-local cumulative log decays)를 계산합니다. 두 접두사 합 (prefix sums) 사이의 차이는 긴 유지 벡터 (retention vectors) 체인을 명시적으로 곱하지 않고도 인코딩합니다. 그런 다음 인과 (causal) 및 상호작용 (interaction) 행렬을 구성합니다. 이는 청크의 교정된 쓰기 (corrected writes)에 대한 WY 스타일 표현 (WY-style representation)을 형성하는 데 사용됩니다.