- 소프트맥스 어텐션에서 출발해 고정 크기 상태를 쓰는 선형 어텐션, 오류만 기록하는 DeltaNet, 전체 상태를 감쇠하는 Gated DeltaNet, 채널별로 감쇠하는 Kimi Delta Attention(KDA)까지 단계적으로 유도함
- 기본 선형 어텐션은 과거 key-value 외적의 합을 상태 (S_t)에 저장해 시퀀스 길이에 선형으로 동작하지만, 새 값을 대입하지 않고 기존 연관에 더하는 가산 쓰기 간섭이 발생함
- DeltaNet은 현재 key에서 예측한 값과 목표 value의 차이에 (\beta_t)를 곱해 기록하며, 즉시 재구성 조건·온라인 경사하강·랭크 1 상태 갱신이라는 세 해석이 같은 수식으로 이어짐
- Gated DeltaNet은 스칼라 (\alpha_t)로 상태 전체를 먼저 감쇠하고, KDA는 이를 대각 행렬 (D_t=\operatorname{Diag}(\alpha_t))로 확장해 key 채널마다 서로 다른 비율로 정보를 유지하거나 삭제함
- 같은 KDA 점화식을 디코드용 융합 순환 Triton 커널과 학습·긴 프리필용 청크 방식으로 실행하며, 청크 방식은 토큰 내부 의존성을 삼각 풀이로 복원해 행렬 곱으로 재구성함
표기법과 전개 순서
- bra-ket 표기에서 (\lvert q\rangle)는 열벡터, (\langle k\rvert)는 행벡터, (\langle k\vert q\rangle)는 스칼라, (\lvert v\rangle\langle k\rvert)는 행렬임
- 하나의 인과적 어텐션 헤드와 실수 벡터를 사용하고, DeltaNet key는 정규화되어 있으며 상태는 key 공간에서 value 공간으로 사상한다고 가정함
- 전개 순서는 소프트맥스 어텐션 → 선형 어텐션 → DeltaNet → Gated DeltaNet → KDA이며, 마지막에 순환 및 청크형 Triton 구현으로 연결됨
- DeltaNet 계열 가운데 두 변형은 최신 Qwen과 Kimi 모델 계열에서 사용됨
이차 복잡도 어텐션에서 선형 상태로
- 일반적인 인과적 소프트맥스 어텐션은 key와 query의 유사도를 계산하고, 모든 과거 key에 대한 점수를 분포로 정규화한 뒤 value 벡터의 가중합을 출력함
- 길이 (T)인 시퀀스에는 (T^2)개의 key-query 쌍이 있음
- 자동회귀 추론에서는 key와 value를 캐시할 수 있지만 캐시 크기가 시퀀스와 함께 늘어남
- 새로운 query도 전체 과거를 확인해야 함
- 소프트맥스 분모는 현재 query와 모든 이전 key에 공동으로 의존하므로 계산 순서를 단순히 재배치하기 어려움
- 소프트맥스를 제거하면 출력을 과거 key-value 외적의 합으로 묶을 수 있음
- (S_t=\sum_{i\le t}\lvert v_i\rangle\langle k_i\rvert)
- (S_t=S_{t-1}+\lvert v_t\rangle\langle k_t\rvert)
- (\lvert o_t\rangle=S_t\lvert q_t\rangle)
- 핵심 항등식은 ((\lvert v\rangle\langle k\rvert)\lvert q\rangle=\langle k\vert q\rangle\lvert v\rangle)이며, 모든 과거 key와 value 대신 합산된 외적을 고정 크기 (d_v\times d_k) 상태에 저장함
- 토큰을 한 번 순회하므로 시퀀스 길이에 선형으로 동작하지만, 그 대가로 소프트맥스의 정규화와 선택성을 잃음
- 더 정교한 선형 어텐션은 특징 맵과 정규화 항을 사용함
선형 어텐션의 가산 쓰기 문제
- 정규화된 현재 key에 (\lvert v_t\rangle\langle k_t\rvert)를 기록한 직후 같은 key로 읽으면 (S_t\lvert k_t\rangle=S_{t-1}\lvert k_t\rangle+\lvert v_t\rangle)가 됨
- 새 쓰기는 메모리가 (v_t)를 반환하도록 대입하지 않고 기존 반환값에 (v_t)를
+=방식으로 더함 - 이전 상태가 이미 올바른 값을 반환한다면 같은 value가 두 배가 되며, key들이 서로 직교하지 않으므로 각 쓰기가 기존 쓰기와 간섭할 수 있음
- 선형 어텐션은 압축된 연관 메모리를 제공하지만, 필요한
=에 가까운 갱신 대신 가산 갱신을 수행함
DeltaNet: 값 대신 예측 오류 쓰기
- DeltaNet은 새 key에 대한 기존 예측 (\widehat v_t=S_{t-1}k_t)를 먼저 읽고 전체 value 대신 차이만 기록함
- (e_t=\beta_t(v_t-S_{t-1}k_t))
- (S_t=S_{t-1}+e_tk_t^\mathsf T)
- 학습된 쓰기 강도 (\beta_t)는 ([0,1]) 범위임
- 같은 key로 즉시 다시 읽으면 ((1-\beta_t)S_{t-1}k_t+\beta_tv_t)가 됨
- (\beta_t=1)이면 정확히 (v_t)를 반환함
- 더 작은 값은 기존 예측을 목표 방향으로 일부만 이동시킴
- 갱신은 key 공간에 국소적임
- 현재 key와 직교하는 query 방향에서는 외적 갱신이 0이므로 응답이 바뀌지 않음
- 현재 key 방향의 연관만 선택적으로 교체함
-
재구성 손실로 유도하기
- 상태 (S)를 선형 사상으로 보고 현재 key-value 쌍의 손실을 (\frac12\lVert Sk_t-v_t\rVert_2^2)로 두면 기울기는 ((Sk_t-v_t)k_t^\mathsf T)임
- (S_{t-1})에서 크기 (\beta_t)로 경사하강 한 단계를 수행하면 DeltaNet의 갱신식과 정확히 같아짐
- 같은 갱신을 세 방식으로 해석할 수 있음
- 메모리 연산에서 (\beta_t)는 기존 연관 교체 강도임
- 온라인 학습에서 (\beta_t)는 학습률임
- 선형대수에서는 예측 오류와 key의 랭크 1 외적임
-
구조화된 상태 전이
- 갱신을 전개하면 (S_t=S_{t-1}(I-\beta_tk_tk_t^\mathsf T)+\beta_tv_tk_t^\mathsf T)가 됨
- 단위 key에 대해 (I-\beta_tk_tk_t^\mathsf T)는 현재 key 방향에서 고윳값 (1-\beta_t), 모든 직교 방향에서 고윳값 1을 가짐
- 기존 key 방향의 연관을 먼저 제거하고 새 연관을 더하지만, 상태 전체의 수명 관리는 아직 해결하지 못함
Gated DeltaNet: 상태 전체를 먼저 잊기
- 하나의 행렬에 전체 과거를 압축하면 이미 상태에 합쳐진 개별 토큰만 선택해 건너뛸 수 없음
- DeltaNet은 현재 key 주변을 교정하지만, 다른 방향의 오래된 정보는 남아 향후 읽기에 계속 기여할 수 있음
- Gated DeltaNet은 학습된 스칼라 유지 게이트 (\alpha_t\in[0,1])를 적용함
- (\widetilde S_t=\alpha_tS_{t-1})로 잊음
- (\widehat v_t=\widetilde S_tk_t)로 예측함
- (e_t=\beta_t(v_t-\widehat v_t))로 교정함
- (S_t=\widetilde S_t+e_tk_t^\mathsf T)로 기록함
- 잊기 → 예측 → 교정 → 쓰기 순서가 중요함
- 감쇠 전에 예측하면 오류를 계산한 메모리와 실제로 갱신하는 메모리가 달라짐
- 델타 규칙은 목표 key에 대한 교체를, 스칼라 게이트는 전역 삭제를 담당해 서로 다른 문제를 해결함
- 다만 하나의 (\alpha_t)가 행렬 전체에 적용되므로 모든 key 채널을 같은 비율로 유지하거나 잊어야 함
Kimi Delta Attention: 채널별 감쇠
- Kimi Delta Attention은 스칼라 (\alpha_t)를 (d_k)차원 벡터로 바꾸고 (D_t=\operatorname{Diag}(\alpha_t))를 구성함
- 상태가 key 공간에서 value 공간으로 사상되므로 key 채널은 (S)의 열에 해당하며, 오른쪽 곱 (S_{t-1}D_t)가 각 열에 서로 다른 유지율을 적용함
- KDA는 다음 순서로 동작함
- (\widetilde S_t=S_{t-1}D_t)로 key 채널별 감쇠
- (\widehat v_t=\widetilde S_tk_t)로 예측
- (e_t=\beta_t(v_t-\widehat v_t))로 교정
- (S_t=\widetilde S_t+e_tk_t^\mathsf T)로 기록
- (o_t=S_t(d_k^{-1/2}q_t))로 읽기
- Gated DeltaNet에서 KDA로의 개념적 변화는 (\alpha_t)를 (D_t)로 승격한 것뿐이지만, 한 채널을 지우면서 다른 채널은 유지할 수 있음
-
대각-저랭크 전이
- KDA를 전개하면 (S_t=S_{t-1}A_t+\beta_tv_tk_t^\mathsf T)이며, (A_t=D_t(I-\beta_tk_tk_t^\mathsf T))임
- (A_t=D_t-b_ta_t^\mathsf T), (b_t=D_tk_t), (a_t^\mathsf T=\beta_tk_t^\mathsf T)로 쓸 수 있어 대각-저랭크(DPLR) 전이가 됨
- DPLR은 key 공간에서 작용하는 (d_k\times d_k) 전이를 뜻하며, 메모리 상태 자체는 여전히 (d_v\times d_k) 행렬임
- 계열마다 다음 기능이 추가됨
- 선형 어텐션: 고정 크기 순환 메모리
- DeltaNet: 목표 방향의 선택적 교체
- Gated DeltaNet: 상태 전체 감쇠
- KDA: key 채널별 감쇠
- 구현은 보통 (g_t=\log\alpha_t\le0)를 저장한 뒤 (\exp(g_t))로 유지율을 구함
- 전치된 (d_k\times d_v) 레이아웃의 5단계 참조 구현은
naive_recurrent_kda에서 확인할 수 있음
디코드용 융합 순환 Triton 커널
- KDA에는 두 가지 주요 실행 방식이 있음
- 융합 순환 방식: 디코드, 짧은 시퀀스, 상태 유지형 서빙에 적합함
- 청크 방식: 학습과 긴 프리필에 적합함
fused_recurrent_kda_fwd는 시퀀스·value head·32폭 value 타일마다 하나의 Triton 프로그램을 실행함BK는 일반적인 지원 구성에서 key 차원을 덮음- 각 프로그램은 전치 상태의
[BK, BV]타일을 소유하고 토큰을 순서대로 순회함 - 서로 다른 value 타일, head, 시퀀스는 독립적으로 실행됨
- 커널은 상태 감쇠, key에 대한 예측 축약, residual 계산, 외적 쓰기, query 읽기 축약을 점화식 그대로 수행함
- 한 번에 새 토큰 하나만 들어오는 디코드에는 적합하지만, 벡터 연산을 Tensor Core에 효율적인 큰 행렬 곱으로 바꾸지 못해 학습과 긴 프리필에는 불리함
Chunkwise KDA: 점화식을 행렬 곱으로 재배치
- Chunkwise KDA는 (C)개 토큰을 함께 처리하면서 토큰별 순환 방식과 정확히 같은 상태와 출력을 만들어야 함
- 각 청크는 두 결과를 계산함
- 들어오는 상태 (S_c)에서 청크 전체를 처리한 뒤의 (S_{c+1})
- 청크 내부 모든 토큰의 인과적 출력
- 각 토큰의 델타 오류가 같은 청크 안의 이전 쓰기에 의존한다는 점이 핵심 난점임
-
누적 감쇠와 임시 오류
- 토큰 (i)의 대각 감쇠를 (D_i), 청크 경계부터 토큰 (i)까지의 누적 감쇠를 (D_{0:i}=D_0D_1\cdots D_i)로 둠
- 토큰 (j)의 쓰기가 토큰 (i)까지 전달될 때는 (D_{j+1:i})가 적용되며, 대각 행렬이므로 감쇠 행렬끼리는 교환 가능함
- 먼저 청크 내부의 다른 쓰기를 무시한 임시 오류를 병렬 계산함
- (\bar e_i=\beta_i(v_i-S_cD_{0:i}k_i))
- 첫 토큰을 제외한 임시 오류는 이전 청크 내부 쓰기의 영향을 빠뜨렸으므로 그대로 사용할 수 없음
-
인과 의존성 복원
- 이전 토큰 (j)가 현재 토큰 (i)의 오류에 미치는 계수를 (\rho_{ij}=\beta_i k_j^\mathsf TD_{j+1:i}k_i)로 정의함
- 실제 오류는 (e_i=\bar e_i-\sum_{j<i}\rho_{ij}e_j) 형태로 순차 의존성을 가짐
- (\rho_{ij})를 엄격한 하삼각 행렬 (R_c)에 넣으면, 쌓아 놓은 오류 행렬은 (E_c=\bar E_c(A_c^{kk})^\mathsf T), (A_c^{kk}=(I+R_c)^{-1})로 계산됨
- 일반적인 밀집 역행렬은 필요하지 않음
- (I+R_c)는 대각 원소가 1인 삼각 행렬임
- 각 value 채널에 대해 인과적 삼각 풀이를 수행하면 됨
-
청크 종료 상태 계산
- 들어오는 상태는 청크의 모든 감쇠를 통과하고, 각 청크 내부 쓰기는 자신보다 뒤에 있는 감쇠만 통과함
- 청크 끝까지 감쇠된 key들을 (K_c^{\mathrm{end}})에 행으로 쌓으면 상태를 다음 행렬 곱으로 정리할 수 있음
- (S_{c+1}=S_cD_{0:C-1}+E_cK_c^{\mathrm{end}})
- 여러 랭크 1 외적 쓰기를 하나의 행렬 곱으로 합쳐 청크 전체 상태를 한 번에 전진시킴
-
청크 내부의 모든 출력 계산
- KDA는 현재 토큰을 쓴 뒤 읽으므로 토큰 (i)의 출력에는 자신의 쓰기도 포함됨
- 이전 쓰기 (j)가 query (i)에 미치는 계수를 (\chi_{ij}=s,k_j^\mathsf TD_{j+1:i}q_i), (j\le i)로 정의함
- 계수를 하삼각 읽기 행렬 (A_c^{qk})에 배치함
- 상삼각의 0은 미래 토큰의 기여를 차단함
- 대각 원소는 현재 토큰이 자신의 쓰기 이후 읽는 동작을 반영함
- 경계에서 각 query까지 감쇠된 벡터를 (Q_c^{\mathrm{boundary}})에 쌓으면 전체 출력은 다음과 같음
- (O_c=sS_cQ_c^{\mathrm{boundary}}+E_c(A_c^{qk})^\mathsf T)
- 첫 행렬 곱은 감쇠된 청크 진입 상태를 읽고, 두 번째 곱은 청크 내부의 인과적 쓰기 기여를 더함
Chunkwise Triton 파이프라인
- 청크 구현은 하나의 거대한 커널이 아니라 여러 커널 호출로 이어지는 파이프라인임
- 먼저 청크 내부 누적 로그 감쇠를 계산함
- 두 prefix sum의 차이로 유지 벡터를 길게 곱하지 않고 (D_{j+1:i})를 표현함
- 이어서 인과적 (A^{qk})와 (A^{kk}) 상호작용 행렬을 만들고, (A^{kk})로 청크의 교정된 쓰기에 대한 WY 형태를 구성함
- 상태 커널은 유일한 청크 간 순회를 수행함
- 각 청크로 들어오는 상태를 생성함
- 청크의 델타 오류를 해소함
- 진입 상태가 계산된 뒤에는 출력 커널이 서로 다른 청크와 타일의 토큰을 병렬 처리할 수 있음
- 실제 구현은 먼저 16토큰 대각 상호작용 블록을 계산한 뒤 융합된 비대각 및 삼각 풀이 커널을 실행함
chunk_kda_fwd가 단계를 조정하며 주요 진입점은chunk_kda_fwd_intra,chunk_gated_delta_rule_fwd_h,chunk_gla_fwd_o_gk임- 코드의
v_new는 해소된 오류임 h는 청크 진입 상태임kg는 청크 끝까지 감쇠된 key임
- 코드의
- 순환 방식과 청크 방식은 서로 다른 어텐션이 아니라 동일한 KDA 점화식의 두 실행 스케줄임
- 순환 방식은 저지연 디코드를 위한 직렬 벡터 연산임
- 청크 방식은 Tensor Core 중심의 학습과 프리필을 위한 행렬 연산임