Kimi Delta Attention, 수식만 보면 무섭지만 결국 ‘덮어쓰기 가능한 메모리’ 이야기

왜 linear attention부터 봐야 하나

• 이 글의 출발점은 Kimi Delta Attention(KDA) 수식이 갑자기 튀어나온 괴물이 아니라는 것임

  • 글은 softmax attention -> linear attention -> DeltaNet -> Gated DeltaNet -> KDA 순서로 따라감
  • 최신 Qwen, Kimi 모델 계열에 쓰이는 DeltaNet 변종이 왜 이런 모양이 됐는지 역추적하는 방식임

• 일반적인 causal softmax attention은 새 query가 과거 key 전체를 훑어야 함

  • 길이 T 시퀀스에서는 key-query 쌍이 T^2개 생김
  • autoregressive inference에서 key/value cache를 써도 캐시는 시퀀스와 함께 커지고, 새 토큰은 여전히 전체 히스토리를 봐야 함

• linear attention의 트릭은 softmax를 잠깐 빼고 계산을 재배열하는 데 있음

  • value x key outer product를 누적하면 과거 전체를 고정 크기 행렬 S_t로 압축할 수 있음
  • 읽을 때는 이 상태 행렬에 query를 곱하면 됨
  • 복잡도는 시퀀스 길이에 대해 quadratic이 아니라 linear로 내려감

DeltaNet은 값을 더하지 않고 오차를 쓴다

• plain linear attention의 문제는 write가 assignment가 아니라 addition이라는 점임

  • 어떤 key에 value를 저장하고 바로 같은 key로 읽으면, 기존 상태가 내놓던 값에 새 value가 더해짐
  • old state가 이미 맞는 값을 내놓고 있었다면 새 write 때문에 값이 두 배로 튈 수도 있음
  • key들이 서로 완전히 직교하지 않으니, 한 write가 다른 기억에도 간섭함

• DeltaNet은 그래서 ‘value 전체’가 아니라 ‘예측 오차’를 기록함

  • 먼저 현재 상태가 새 key에서 어떤 value를 예측하는지 S_{t-1} k_t로 읽음
  • 실제 value와 예측 value의 차이만 beta_t만큼 써넣음
  • beta_t = 1이면 해당 key 방향에서는 즉시 실제 value로 맞춰지고, 더 작으면 기존 기억에서 목표 쪽으로 일부만 이동함

• 같은 update는 온라인 학습 관점에서도 나온다는 게 포인트임

  • 현재 key-value 쌍 하나에 대해 reconstruction loss를 잡고 gradient descent를 한 스텝 하면 DeltaNet update가 그대로 나옴
  • 메모리 관점에서는 ‘기존 association을 얼마나 교체할지’이고, 학습 관점에서는 beta_t가 step size임

Gated DeltaNet과 KDA의 차이

• DeltaNet은 특정 key 방향의 write 문제는 고치지만, 오래된 정보의 수명 문제는 아직 남음

  • linear state는 과거 전체를 한 행렬에 압축하므로, 어떤 오래된 토큰 하나만 골라 건너뛰기 어렵다
  • query와 겹치는 저장 방향은 계속 출력에 영향을 줄 수 있음

• Gated DeltaNet은 상태를 쓰기 전에 먼저 전역적으로 줄이는 retention gate를 넣음

  • alpha_t라는 스칼라 gate로 이전 상태 전체를 같은 비율로 보존하거나 잊음
  • 순서가 중요함: 먼저 잊고, retained state에서 예측하고, 그 예측 오차를 다시 write함

• Kimi Delta Attention은 여기서 한 단계 더 가서 key 채널마다 retention을 따로 둠

  • 스칼라 alpha_t를 벡터로 바꾸고, 이를 Diag(alpha_t) 대각 행렬로 만들어 상태 오른쪽에 곱함
  • 결과적으로 어떤 key 채널은 지우고, 다른 key 채널은 남기는 식의 더 세밀한 forget이 가능해짐
  • 개념적 변화는 작아 보이지만 효과는 큼: whole-state forgetting에서 per-key-channel forgetting으로 바뀌는 셈임

• 이 전이는 diagonal-plus-low-rank(DPLR) 형태로 정리됨

  • KDA의 key-space transition은 대각 행렬에서 rank-one 행렬을 뺀 형태로 쓸 수 있음
  • 메모리 상태 자체는 여전히 d_v x d_k 행렬이고, DPLR은 그 상태에 작용하는 key-space 전이를 설명하는 말임

구현은 recurrent와 chunkwise 두 갈래

• decode에는 fused recurrent Triton 커널이 자연스럽게 맞음

  • 한 번에 새 토큰 하나를 처리하므로 state를 token 순서대로 업데이트하면 됨
  • 커널은 state *= exp(g_t), prediction reduction, residual 계산, outer product write, output reduction 흐름을 거의 그대로 옮긴 형태임
  • launch는 sequence, value head, 32-wide value tile 단위로 나뉨

• 학습이나 긴 prefill에서는 recurrent 방식이 GPU를 잘 못 먹임

  • 토큰별 vector operation은 tensor core가 좋아하는 큰 matrix multiplication으로 잘 바뀌지 않음
  • 그래서 같은 recurrence를 chunk 단위로 재배열해 matrix product와 triangular solve 중심으로 바꿈

• chunkwise KDA는 여러 토큰을 한 번에 처리하되 causal dependency를 복원해야 함

  • 먼저 각 토큰이 incoming state만 본다고 가정한 provisional error를 병렬로 계산함
  • 이후 chunk 안에서 앞선 write들이 뒤 토큰 prediction에 미치는 영향을 lower-triangular system으로 풀어 실제 error를 구함
  • I + R_c가 대각이 1인 삼각 행렬이라 일반 dense inverse가 아니라 causal triangular solve로 처리 가능함

• chunk 끝 state와 chunk 내부 output도 matrix multiplication으로 정리됨

  • chunk 끝 state는 incoming state에 누적 decay를 적용한 항과, chunk 내부 write들을 end boundary까지 decay시킨 항으로 나뉨
  • token output은 incoming chunk state를 읽는 항과, 같은 chunk 안에서 이미 발생한 write의 causal contribution을 더하는 항으로 구성됨

---

기술 맥락

• 이 글에서 중요한 기술 선택은 softmax attention의 T^2 비용을 고정 크기 recurrent state로 바꾸는 거예요. 긴 컨텍스트 모델에서는 과거 key/value를 계속 들고 읽는 비용이 커지니까, 과거를 행렬 하나에 접어 넣는 방식이 매력적이거든요.

• 그런데 그냥 접어 넣으면 문제가 생겨요. plain linear attention은 메모리에 값을 ‘대입’하는 게 아니라 계속 ‘더하기’를 하니까, 같은 key 방향의 기존 기억과 새 기억이 서로 섞여요. DeltaNet이 오차만 쓰는 이유가 바로 이 지점이에요.

• Gated DeltaNet과 KDA의 차이는 잊는 단위예요. Gated DeltaNet은 상태 전체를 같은 비율로 줄이지만, KDA는 key 채널별로 retention을 다르게 줘요. 그래서 오래된 정보와 필요한 정보를 더 세밀하게 다룰 여지가 생겨요.

• 구현에서도 선택지가 갈려요. decode는 토큰이 하나씩 들어오니 fused recurrent Triton 커널이 낮은 지연시간에 맞고, 학습이나 긴 prefill은 chunkwise로 바꿔 tensor core가 잘 먹는 행렬 연산을 늘리는 쪽이 유리해요.

• 그래서 이 글의 결론은 수식 하나를 외우라는 게 아니에요. KDA는 ‘고정 크기 메모리’, ‘오차 기반 write’, ‘채널별 forget’, ‘실행 스케줄 분리’가 합쳐진 설계라고 보면 훨씬 덜 무섭게 읽혀요.