掘金2026년 11월 2일 21:30중국어 → 한국어

Kimi K3 경량 복제 분석: 채널 단위로 망각 제어하는 KDA

Kimi K3를 경량으로 재현하는 연재에서 KDA를 채널 단위 망각 제어 구조로 설명하고 등장 배경과 위치를 정리했다.

중국어 원문을 AI로 번역했습니다. 고유명사와 수치는 원문 표기를 우선하며, 중요한 판단에는 아래 출처 원문을 함께 확인하세요.

0x00 개요

0x01 진화 맥락과 위치

- 1.1 전신

- 1.2 동시기 경쟁 기법

- 1.3 KDA 논문

- 1.4 위치

- 1.5 KimiDeltaAttention

0x02 원리 설명

- 2.1 직관적 비유 관점

- 2.2 전문가 관점

- 2.3 KDA 원리도

- 2.4 심층 논의

- 2.5 범례

- 2.6 각종 메커니즘 비교

0x03 하위 지식 포인트 분해

0x04 하위 지식 포인트 심층 분석

- 4.1 상태 재귀와 delta rule (Eq.1)

- 4.2 유계 감쇠 log-decay (Eq.5)

- 4.3 channel-wise decay와 저랭크 매개변수화 (Eq.2)

- 4.4 Head-wise RMSNorm(head별로 정규화한 RMSNorm)

- 4.5 출력 게이트 (Eq.6)

- 4.6 Gated MLA와 KDA/MLA의 상보성

0xFF 참고

이 문서는 KIMI K3 KDA(Kimi Delta Attention)를 재현하는 첫 번째 부분이며, 글이 너무 길어서 두 회로 나누어 발행합니다.

GDN의 head-wise 거친 망각과 바닐라 선형 어텐션의 비선택적 누적과 달리, KDA는 Diag(α_t) channel-wise 유계 감쇠 + delta-rule 기록을 통해 고정 크기 상태 행렬 S上에서 제어 가능한 기억, 조밀한 chunkwise 병렬 처리, 그리고 위치 인코딩이 없어도 자체 일관적인 선형 어텐션을 구현합니다.

참고:

- 프로젝트 github.com/RossiLH/Mod…와 시리즈 블로그는 필자의 하나의 시도이자 학습 과정입니다. 독자의 요구를 충족하는지 확신할 수 없으며, 스스로 성찰하고 개선할 것이고, 여러분의 소중한 의견도 바랍니다.

- KIMI K3, DeepSeek V4, PI0.5, PI0.7 등 모델을 다룹니다. 이후 새로운 모델을 추적할 수 있습니다.

- 전 CPU 재현이며 GPU를 포함하지 않고, 대규모 학습(이는 또 다른 방대하고 심오한 영역)과 하드웨어도 다루지 않습니다.

- 모델 구성 요소는 가능한 한 원본 대규모 모델의 축소판이며, 구성 요소에 오픈소스 버전이 없으면 근사 구성 요소로 대체합니다.

- 각 구성 요소가 어떤 문제를 해결하는지, 그 문제가 어디에 어려움이 있는지, 왜 이렇게만 해결할 수 있는지, 해법의 원리는 무엇인지, 그리고 구성 요소들이 어떻게 맞물려 하나의 전체를 이루는지 설명합니다. 중점은 용어를 외우는 것이 아니라 하나의 주선을 이해하는 것입니다.

- 각 주요 구성 요소에는 바로 실행 가능한 Gradio demo가 있습니다. GPU가 필요 없습니다. 또한 정적 그림을 생성할 수 있는 여러 시각화 스크립트(plot_*.py)가 있습니다.

- 모든 파일의 줄 번호는 로컬 소스 스냅샷 기준이며, 이후 소스가 변경되면 줄 번호는 이동할 수 있지만 함수명과 호출 관계는 여전히 위치를 찾는 진입점으로 사용할 수 있습니다.

- 각 글이 가능한 한 독립적으로 완성될 수 있도록 일부 내용이 서로 다른 글에 모두 존재하여 시리즈 차원에서 중복이 있을 수 있으니 이해해 주시기 바랍니다.

- 바닐라 선형 어텐션(Linear Transformer): S = Σ k v^T, 행렬 덧셈은 교환 가능합니다. 결함: 고정 상태 용량 제한, 표현력 약함, 선택적 망각 없음.

- Gated Linear Attention(GLA): channel-wise gating의 대표이지만, 선형 어텐션(외적 누적)을 그대로 사용하며 delta-rule의 "기록 시 이전 기여를 먼저 빼는" 메커니즘이 없습니다.

- Gated DeltaNet(GDN): delta rule(DeltaNet, Mamba-2와 동일 기원)에 head-wise forget gate를 더했습니다. 결함: gate가 각 head마다 하나의 스칼라라서 망각 입도가 너무 거칠어 서로 다른 feature 차원에 대해 차별화된 기억 제어를 할 수 없습니다.

KDA의 목표는 GDN의 delta rule을 기반으로 망각 입도를 channel 수준으로 세분화하는 것입니다. 이를 위해서는 상태 업데이트의 전이 행렬에 대각 gate Diag(α_t)를 추가하면서 동시에 chunkwise 분할 병렬(DPLR / Diagonal-Plus-Low-Rank, 대각+저랭크 구조)을 유지해야 합니다.

| 방안명 | 핵심 메커니즘 | 장점 | 단점 | | --- | --- | --- | --- | | KDA | delta rule + channel-wise 대각 decay gate + 유계 sigmoid 감쇠 | 기억 제어 가능, 조밀한 chunkwise 가능, NoPE 위치 정보 자체 일관성 | 상태 용량 고정, 장거리 정확 검색은 full attention보다 약함 | | Gated DeltaNet | delta rule + head-wise forget gate | 단순하고 안정적 | gate 입도가 거칠어 채널 수준 기억 제어가 나쁨 | | Mamba-2 | SSD(구조화 상태 공간 쌍대), 선택성 강함 gate | 하드웨어 고효율 | SSM 상태 공간 이론 필요, decay 의미가 delta rule과 다름 | | Softmax attention(MLA) | token-to-token 전체 상호작용 | 정확 검색, 전역 정보 흐름 | KV cache가 시퀀스에 따라 선형 증가 | | 희소 어텐션 | 선정된 token에만 softmax | 긴 문맥 비용 제어 가능 | 선택을 위해 여전히 전체 KV 저장 필요, 표현 상한은 full attention을 넘지 않음 |

KDA와 직접 관련된 논문은 두 편이며, 이 문서는 기술 보고서(K3 보고서)를 주선으로 하고 Kimi Linear 보고서를 알고리즘 세부 사항의 보충으로 삼습니다.

| 논문 | 출처 | | --- | --- | | Kimi Linear: An Expressive, Efficient Attention Architecture | (Kimi Linear 기술 보고서, KDA의 알고리즘 정의 출처) arXiv:2510.26692 (2025-10-30); 저자: Kimi Team 등; 기관: Moonshot AI | | Kimi K3: Open Frontier Intelligence | (K3 기술 보고서, KDA가 K3에서 최종 형태 + 시스템 공동 설계) arXiv:2607.24653 (2026-07-27); 저자: Kimi Team; 기관: Moonshot AI |

- KDA는 우선 Kimi Linear 기술 보고서에서 완전하게 정의되었다. 이는 Gated DeltaNet(GDN)을 확장하여, 거친(head-wise) forget gate를 channel-wise 대각 gate Diag(α_t) 로 교체한다.

- K3 보고서 §2.1.1은 K3에서의 KDA 구현 세부 사항(유계 감쇠 Eq.5, 전랭크 출력 게이트 Eq.6)을 추가로 제시하며, §5.1에서는 KDA의 알고리즘-시스템 협조 설계(FlashKDA, KCP, KDA-aware 프리픽스 캐시)를 체계적으로 서술한다.

paper-eq-5

paper-eq-6

Attention의 추상화는 "반드시 Softmax가 나타나야 한다"가 아니라,

Key에 따라 Value를 어떻게 읽을지 결정하는 것이다.

KDA는 상태 행렬을 통해 이를 완성한다: y_t = q_tᵀ S_t . 따라서 그것은 여전히 "Query가 Key-Value 관계에 따라 내용을 읽는" 것이며, 단지 다음과 같을 뿐이다:

- Full Attention: token 수준 KV를 저장하고, 읽을 때 Softmax를 계산한다;

- Linear/KDA: 압축 상태를 저장하고, 읽을 때 커널/행렬 곱셈을 사용한다;

- KDA는 또한 망각이 있는 온라인 업데이트를 추가한다.

따라서 이렇게 말할 수 있다: KDA는 어텐션의 재귀적, 커널화, 게이팅이 있는 구현 중 하나이다. 어텐션의 본질은 내용 주소 지정(content addressing)이며, Softmax라는 구체적 함수가 아니다.

KimiDeltaAttention 코드는 다음과 같다:

class KimiDeltaAttention(nn.Module): """KDA 전체 레이어: 투영(q/k/v/beta/decay) + 핵심 재귀 + 출력 게이트. 파라미터 구조(보고서 Eq. 2 / Eq. 5 / Eq. 6): - ``q_proj/k_proj/v_proj``: 각 head의 key/value 차원으로의 선형 투영. - ``q_conv/k_conv/v_conv``: 인과적 depthwise ShortConv(폭 4)로, 위치 인코딩이 없는 모델에 토큰 단위의 \"변위\" 정보를 주입한다(보고서 §2.1.1; Kimi Linear 논문은 이를 위치 인코딩 대체물이라고 부른다). - ``beta_proj``: delta-rule 기록 강도 ``beta = sigmoid(x)``로, \"이번 기록이 얼마나 많은 옛 기억을 덮어쓰는지\"를 제어한다(``beta``가 1에 가까우면 해당 방향의 기억을 완전히 재작성함을 의미). - ``alpha_down/alpha_up/alpha_bias``: **저랭크** decay logits ``z = W_up W_down x + b``(Eq. 2)로, channel별 log 감쇠의 원천. - ``a_log``: Eq. 5의 head별 로그 스케일 ``A_h``로, 논문에서는 0으로 초기화한다. - ``o_norm``: head-wise RMSNorm(Eq. 6의 첫 번째 단계). - ``g_proj``: 전랭크, 입력 의존적 출력 게이트 ``sigmoid(W_g x)``(Eq. 6의 두 번째 단계). """ def __init__(self, dim: int, cfg: KDAConfig, *, layer_idx: int = 0 , eps: float = 1 e- 6 ): super().__init__() self.dim = dim self.cfg = cfg self.layer_idx = layer_idx h, dk, dv = cfg.num_heads, cfg.head_dim_k, cfg.head_dim_v self.num_heads, self.head_dim_k, self.head_dim_v = h, dk, dv k_dim, v_dim = cfg.k_dim, cfg.v_dim # k_dim = H*Dk, v_dim = H*Dv # --- head로 분할되지 않은 넓은 벡터로 선형 투영한 뒤 ShortConv에 넘겨 국소 혼합을 수행 ---------- self.q_proj = nn.Linear(dim, k_dim, bias= False ) self.k_proj = nn.Linear(dim, k_dim, bias= False ) self.v_proj = nn.Linear(dim, v_dim, bias= False ) self.q_conv = ShortConv(k_dim, cfg.conv_kernel) # 폭 4의 인과적 depthwise 컨볼루션 self.k_conv = ShortConv(k_dim, cfg.conv_kernel) self.v_conv = ShortConv(v_dim, cfg.conv_kernel) # --- delta-rule 기록 강도 beta와 저랭크 decay logits ---------------------- self.beta_proj = nn.Linear(dim, h, bias= False ) # -> (B, T, H) head마다 스칼라 하나 # 저랭크 decay logits(Eq. 2): z = W_up W_down x + b_alpha. self.alpha_down = nn.Linear(dim, cfg.alpha_rank, bias= False ) # rank로 축소 self.alpha_up = nn.Linear(cfg.alpha_rank, k_dim, bias= False ) # k_dim으로 다시 확장 self.alpha_bias = nn.Parameter(torch.empty(k_dim)) # channel별 편향 # Eq. 5의 head별 log-scale A_h; 논문에서는 0으로 초기화. self.a_log = nn.Parameter(torch.zeros(h)) # --- 출력 게이트(Eq. 6): head-wise RMSNorm + 전랭크 입력 의존 게이트 ------------ self.o_norm = HeadRMSNorm(dv, eps) self.g_proj = nn.Linear(dim, v_dim, bias= False ) # 전랭크 출력 게이트 self.o_proj = nn.Linear(v_dim, dim, bias= False ) self.reset_parameters() def forward( self, x: Tensor, *, reset_mask: Tensor | None = None, cache: KDACache | None = None, initial_state: Tensor | None = None, ) -> tuple [Tensor, Tensor] : """``x``는 ``(batch, seq, dim)``이다. 반환은 ``(output, final_state)``. 형태 선택: 학습/prefill은 chunkwise(블록 내 병렬, 블록 간 상태 직렬 전달)로 가고, 자기회귀 decode(`` t == 1 ``이고 이미 상태가 존재)는 단일 스텝 ``kda_step``으로 간다. 두 형태는 수치적으로 완전히 동등하다(``kda_ref``에는 ground truth 검증용 ``kda_recurrent``가 있다).

``` b, t, _ = x.shape q, k, v, beta, g, conv_states = self._project(x, cache) # (B, T, H, D) -> (B, H, T, D): 재귀는 T 차원에서 진행되고, head를 중간으로 올려 브로드캐스트하기 쉽게 한다. qh, kh, vh = (z.transpose(1, 2) for z in (q, k, v)) gh = g.transpose(1, 2) bh = beta.transpose(1, 2) state_in = initial_state if initial_state is not None else (cache.state if cache else None) # 중간 계산은 fp32 사용(activation이 bf16/fp16일 때), 재귀 상태 누적에서 정밀도를 잃지 않는다. work_dtype = torch.float32 if x.dtype != torch.float64 else torch.float64 chunk = self.effective_chunk_size(work_dtype) if t == 1 and state_in is not None: # decode 형태: 단일 token 단일 스텝 업데이트. 이 브랜치는 "이전 상태가 있을 때"만 활성화됨에 유의, # decode의 첫 token도 chunkwise를 탈 수 있도록 보장한다(초기 상태가 0). out, state = kda_step( qh[:, :, 0].to(work_dtype), kh[:, :, 0].to(work_dtype), vh[:, :, 0].to(work_dtype), gh[:, :, 0].to(work_dtype), bh[:, :, 0].to(work_dtype), state_in.to(work_dtype), ) out = out.unsqueeze(2) # (B,H,1,Dv) else: # prefill / 무상태 decode: chunkwise 병렬 형태(블록 내 병렬, 블록 간 직렬). fn = self._kernel() out, state = fn( qh, kh, vh, gh, bh, chunk_size=chunk, initial_state=state_in, reset_mask=reset_mask, dtype=work_dtype, ) # Eq. 6 출력 게이트: head-wise RMSNorm -> 전체 랭크로 입력에 의존하는 게이트 -> 출력 프로젝션. out = out.transpose(1, 2).reshape(b, t, self.num_heads, self.head_dim_v) out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1) out = out * torch.sigmoid(self.g_proj(x)) out = self.o_proj(out) if cache is not None: # decode 시 이번 라운드의 상태를 cache에 다시 써서 다음 token이 사용하게 한다. cache.state = state cache.conv_q, cache.conv_k, cache.conv_v = conv_states return out, state ```

당신이 하나의 고정 크기 "기억 화이트보드"(상태 행렬 S)만 앞에 두고 서가(KV cache)는 없는 사서라고 상상해 보자.

핵심 동작은 "먼저 channel별로 색을 바래게 한 뒤, delta로 써 넣는 것"이다. 새 책(token)이 한 권 올 때마다 두 가지 일을 해야 한다.

- 먼저 망각: 각 책의 각 주제가 가진 자기만의 "바랜 속도"(channel-wise decay α)로 화이트보드 위의 옛 내용 전체를 조금씩 옅게 만든다.

- 다음으로 갱신: 새 책의 내용을 delta rule에 따라 써 넣는다. 다만 쓰기 전에, 이미 알고 있던 새 책 내용에 해당하는 기존 정보를 "지워 버리고"(v − S^T k), 새 정보 자체만 쓴다. 이렇게 하면 중복 누적이 생기지 않는다.

핵심 차이: softmax 서가는 무한히 커서 모든 책을 정확히 기억할 수 있고, KDA 화이트보드는 영원히 이만한 크기지만 매번 얼마나 잊고 무엇을 갱신할지 정확히 결정할 수 있다. 컨텍스트가 기억 용량을 넘어설 때 KDA는 MLA 전역 레이어와 혼합하여 보완한다—"작은 관리인 + 정기적인 대정리".

입력과 출력:

- 입력: x ∈ R^{B×T×d}, 여기서 B=batch, T=시퀀스 길이, d=은닉 차원.

- 프로젝션 후: query/key q_t, k_t ∈ R^{d_k}(L2 정규화 후), value v_t ∈ R^{d_v}, head마다 하나씩.

- 상태: S_t ∈ R^{d_k × d_v}, 고정 크기이며, T에 따라 커지지 않는다.

핵심 재귀(Eq.1, 단일 head 관점):

S_t = (I − β_t k_t k_t ᵀ) Diag(α_t) S_{t-1} + β_t k_t v_t ᵀ (Eq.1)

- Diag(α_t): channel-wise 감쇠 인자, α_t = exp(g_t), g_t ≤ 0은 log-decay. 각 행(channel)이 독립적으로 망각한다.

- Diag(α_t) S_{t-1}: channel별로 옛 상태를 감쇠한다.

- β_t: delta-rule 쓰기 강도 (0,1), "옛 기억의 방향을 덮어쓰는" 정도를 제어한다.

- (I − β_t k_t k_tᵀ): 랭크-1 프로젝션으로, 상태에서 k_t와 정렬된 성분을 β_t만큼 약화한다. 즉, k 방향을 따라 옛 상태에서 "새 key와 겹치는 성분"을 빼는 것—이것이 delta rule의 "새로 쓰기 전에 옛것 지우기"다.

- + β_t k_t v_tᵀ: 새 기억을 쓴다.

출력: o_t = S_tᵀ q_t (각 head에 대해).

등가 재작성(수치적으로 더 안정, 실제 구현):

S_t = S_{t -1 } + β _t k_t ( v_t − S_{t -1 }ᵀ k_t )ᵀ

즉 먼저 "이번에 쓸 새 정보" = v_t 에서 "구 상태가 k_t 방향으로 이미 기억하고 있는 내용"을 뺀 것을 계산하여, 증분만 쓴다.

파고들기: "선택적 망각 + 온라인 학습률"의 delta rule은 선형 회귀의 재귀 최소제곱(RMS)과 동형이며, k_t를 입력, v_t를 감독(supervision), S를 선형 사상 가중치로 보면, S_tᵀ k_t는 예측이고, v_t − S_tᵀ k_t는 잔차이며, β_t는 학습률이다. 따라서 KDA는 어텐션을 "가변 학습률의 온라인 최소제곱"으로 모델링하여, 이론상 표현력이 순수 누적의 선형 어텐션보다 강하며, 이는 자연스럽게 "압축이 곧 지능"이라는 이론적 뒷받침을 갖게 한다.

먼저 공식을 보자 (Eq.1)

S_t = (I - beta_t * k_t k_t ^T) Diag ( alpha_t ) S_{t -1 } + beta_t * k_t v_t ^T └─────────── 토큰별 전이 행렬 ───────────────┘ └── 새 정보 쓰기 ──┘ alpha_t = exp ( g_t ) 채널별 보존 인자 ( 0 < alpha <= 1 )

K3 아키텍처 개요, 좌하단이 KDA 모듈 다이어그램

KIMI-K3-Arch

- 이론적으로 왜 work하는가: KDA는 상태 전이 행렬을 DPLR 특례——Diag(α_t)(대각) + β_t k_t k_tᵀ(랭크-1)——로 매개변수화하는데, 이것이 바로 "chunkwise 병렬화 가능"의 구조적 전제다. 일반 DPLR에 비해 KDA의 전이 행렬은 고전적 delta rule과 일치하며(S의 열이 key 방향을 따라 감쇠), 물리적 의미가 명확하다(각 열이 "key 기억의 망각 곡선").

- 하한 감쇠의 필요성: Kimi Linear는 무계 softplus g = −exp(A)·softplus(z)를 사용하여, 누적 감쇠 곱이 언더플로되고 그 역수(K/Γ의 리스케일링)가 오버플로되어, 대각 tile이 명시적 position-pair 경로를 타도록 강제했다. K3는 g = g_min·sigmoid(exp(A)z)로 바꾸어 log-decay를 (g_min, 0)에 고정했으며, g_min=−5일 때 exp(80)=5.5e34로 여전히 BF16 범위 내에 있다—— 각 인과 tile(대각 포함) 모두 밀집 Tensor-Core matmul로 계산할 수 있다 .

- NoPE 자기 일관성: KDA의 감쇠와 delta 업데이트는 본래 위치를 인지하므로(recency는 감쇠 누적으로 인코딩됨), MLA 층은 RoPE 없이도 전역 혼합이 가능하다. 컨텍스트를 1M으로 확장하는 데 어떠한 위치 파라미터 수술도 필요 없다.

- 논문이 해결하지 못한 이론적 문제: 고정 상태 용량이 초장거리 정밀 검색에 갖는 근본적 상한(논문은 선형 어텐션의 검색이 약함을 인정하고 MLA 혼합으로 보완함); 상태 행렬의 용량-정밀도 트레이드오프(state expansion 방향은 여전히 탐구 중).

도해는 다음과 같다:

plot_kda_principle

그림은 2×2 네 칸으로, 모두 실제 재귀(kda_recurrent, 프로젝트 전체의 KDA ground truth)로 돌려 얻은 것이다:

좌상단 —— 상태 행렬 S(기억): KDA의 "기억"은 고정 크기 Dk × Dv 행렬 하나다. 그림의 히트맵은 실제로 32개 토큰을 다 돌린 후의 최종 상태다; T=32개 토큰이 단 하나의 16×16 행렬만 생성한다——이것이 KV cache와 대비되는 핵심 셀링 포인트다: 메모리 O(1), 시퀀스 길이에 따라 증가하지 않음 .

우상단 —— 단일 스텝 세 가지 일 분해: 토큰이 하나 들어올 때마다 상태는 오직 세 가지 일만 한다:

- 감쇠 Diag(alpha) S: 구 기억에 먼저 채널별로 alpha=exp(g)<=1인 보존 인자를 곱한다(구 정보를 "조금 잊는다");

- 검색 S^T k: 현재 key로 기억 속에 이미 저장된 것이 무엇인지 "조회"한다;

- 쓰기 beta * k (v - S^T k)^T: 오차 v - 검색된 기억 만 상태에 쓴다(delta rule의 핵심).

좌하단 —— delta rule의 "교정" 본질: "같은 key가 반복해서 나타나고 v가 점진적으로 변하는" 장면을 인위적으로 구성하여, key로 검색된 값이 어떻게 최신 v로 수렴하는지 보여준다. 일반 선형 어텐션(S += k v^T)은 "뒤섞여 버리지만", delta rule은 v - S^T k를 통해 충돌을 눌러준다——이것이 "기억이 지속적으로 갱신되며 서로 오염되지 않는" 이유다. 그림의 세 선은 실제 kda_step으로 돌려 얻은 검색값이다.

우하단 —— 유계 log-decay(Eq.5): g = g_min * sigmoid(exp(A) z)는 log-decay를 엄격히 (g_min, 0)에 고정하며, 즉 g_min=-5일 때 g ∈ (-5, 0)이다. 이 "하한"은 Q4/Q5/Q6의 수치 안정성 전제다.

한 문장 결론: KDA = 고정 크기 기억 행렬 S(O(1) 메모리) + 토큰별 세 가지 일(감쇠→검색→오차만 쓰기) + 유계 감쇠(chunkwise 오버플로 방지). 표준 MHA의 KV cache가 T에 따라 선형으로 증가 하는 것과 본질적 대비를 이룬다.

메커니즘 | 감쇠 | 쓰기 | 망각 세분도 | 매 Token 게이트의 Shape 선형 어텐션 | \bar S_t=S_{t-1} | S_t=\bar S_t+v_tk_t^\top | 없음 | 없음 | 게이트 없음 DeltaNet | \bar S_t=S_{t-1} | S_t=\bar S_t+\beta_t(v_t-S_{t-1}k_t)k_t^\top | 없음 | β [num_heads] GLA | \bar S_t=S_{t-1}Diag(α_t) | 선형 어텐션과 동일 | 각 Head의 각 Key 채널 | α [num_heads, head_dim] GDN | \bar S_t=\alpha_tS_{t-1} | DeltaNet과 동일 | 각 Head당 하나의 스칼라 | β, α [num_heads] KDA | \bar S_t=S_{t-1}Diag(α_t) | DeltaNet과 동일 | 각 Head의 각 Key 채널, 즉 S의 각 열 | β [num_heads] ; α [num_heads, head_dim]

이 표의 중심축은 다음과 같다. 서로 다른 방법들은 주로 두 지점에서 차이를 만든다.

- 이전 상태를 어떻게 감쇠시키는가;

- 현재 Value를 어떻게 쓰며, 잔차 오차 보정을 사용하는가.

선형 어텐션에는 망각과 오차 보정이 없다. DeltaNet은 잔차 쓰기를 추가한다. GDN은 Head 단위의 통일된 망각을 추가한다. KDA는 망각을 각 Key 채널 단위까지 세분화한다.

KDA(Kimi Delta Attention) ├── 상태 점화식과 delta rule(핵심 연산자) ├── 유계 감쇠 log-decay( Eq . 5 ) ├── channel-wise decay와 저랭크 파라미터화( Eq . 2 ) ├── ShortConv(위치 인코딩 없는 국소 혼합) ├── chunkwise 병렬( Eq . 3 - 4 , DPLR) ├── 단일 스텝 점화 kda_step(decode 형태) ├── 출력 게이트( Eq . 6 ) ├── 디코딩 캐시 KDACache와 clone () 분기 └── KDA 컨텍스트 병렬(KCP, 보고서 § 5.1 . 2 )

코드 대응은 다음과 같다.

use_triton=False 일 때 _kernel() 은 순수 PyTorch kda_chunkwise 를 반환한다.

통일 기호 약속 (이 문서의 모든 절이 따름):

- S_t ∈ R^{d_k × d_v} : 점화 상태 행렬(Head당 하나).

- q, k ∈ R^{d_k} (이미 L2 정규화됨), v ∈ R^{d_v} .

- g_t ∈ R^{d_k} : 채널별 log-decay(≤0); α_t = exp(g_t) 는 retention 계수.

- β_t ∈ (0,1) : delta-rule 쓰기 강도.

- chunk 크기 C (설정 64, tile 16), tile 은 수치 안전을 위한 기본 감쇠 블록.

KDA는 Delta Rule의 잔차 덮어쓰기 메커니즘을 유지하지만, GDN의 스칼라 망각 게이트를 d_k 차원 벡터로 확장했다.

4.1.1 원리

KDA의 기본 가정은 다음과 같다. 각 층, 각 Head에서 모든 요청은 고정 크기의 상태 행렬 S_t 를 유지한다. 이 행렬은 현재 요청이 지금 위치까지 도달한 전역 기억을 나타낸다. 각 역사적 Token마다 별도의 기록을 남기는 것이 아니라, 모든 역사적 Token의 컨텍스트 정보를 동일한 행렬 하나에 압축한다.

KDA가 현재 Token을 처리할 때, 모델은 현재 Token의 Hidden State와 대응하는 투영 파라미터에 따라 다음을 계산한다.

- q_t : Query, 기억으로부터 정보를 읽는 데 사용;

- k_t : Key, 어드레싱 패턴을 정의하는 데 사용;

- v_t : Value, 쓸 내용을 제공하는 데 사용;

- α_t : 망각 게이트, 각 Key 채널이 이전 정보를 얼마나 보존할지 결정;

- β_t : 쓰기 게이트, 현재 잔차를 얼마나 쓸지 결정;

- γ_t : 출력 게이트, 출력 정보의 강도를 제어.

KDA의 상태 업데이트는 다음과 같이 쓸 수 있다.

(1) 먼저 사용

에서 망각 게이트는 이전 상태의 서로 다른 Key 채널이 얼마나 보존될지를 결정한다. α_t 가 벡터이므로 모든 열이 동일한 비율로 감쇠하지는 않는다. j번째 Key 채널은 독립적인 망각 속도 α_{t,j} 를 가지며, 중요한 채널은 더 많이 보존할 수 있고, 사라져야 할 채널은 더 빠르게 감쇠될 수 있다.

만약 상태 행렬의 형상이 다음과 같다면: . 그러면 오른쪽에 Diag(α_t) 를 곱하면 상태 행렬의 각 열을 개별적으로 스케일링한다. 각 열은 하나의 Key 채널에 대응하므로, KDA가 구현하는 것은 "Head별, Key Channel별" 망각이다.

게이트의 규모도 이에 따라 GDN의:

num_heads

에서 다음으로 확장된다.

num_heads × head_dim

遗忘 게이트가 없다면, 순수 선형 어텐션은 그저 다음과 같을 수 있다: . 행렬 덧셈은 교환법칙을 만족하기 때문에, 입력 순서 [A,B]와 [B,A]는 동일한 최종 상태를 만들어낸다. 이는 단순 누적 상태가 본질적으로 위치 정보를 결여하고 있음을 의미한다. 즉 어떤 연관이 기록되었는지는 알지만, 그 연관들의 선후 순서는 반드시 알지 못한다. KDA에서의 망각은 이러한 순서 정보를 부분적으로 복원할 수 있다. 현재로부터 거리가 멀수록 Token이 겪는 감쇠 횟수가 많아져 그 신호가 더 약해질 수 있다. 따라서 망각 과정은 일종의 암묵적 상대 위치 인코딩으로 볼 수 있다.

(2) 이어서 Delta Rule을 수행한다: .

이는 현재 Value의 기록, 즉 현재 Token이 가져오는 원본 신규 지식이다. 이는 Value v_t와 주소 지정 모드 k_t를 결속시켜 상태에 기록할 준비를 한다.

이 항은 감쇠된 옛 상태 중 현재 Key 모드와 일치하는 부분의 내용을 빼낸다. 그 작용은 단순히 새 내용을 더하는 것이 아니라, 동일한 특징이 긴 문맥에서 끊임없이 반복 누적되는 것을 피하는 데 있다. 먼저 오래된 일치 내용을 삭제하거나 수정한 뒤 새로운 목표 Value를 기록하므로, 상태는 그만큼 더 정밀한 덮어쓰기 능력을 갖추게 된다.

표준 방식과의 수식 수준 비교 :

softmax attention KDA 메모리 cache = {(k_i, v_i)} token별 저장 S 단일 행렬 검색 Σ softmax(q·k) v (전역 정확) Sᵀ q (압축 근사) 복잡도 시간 O(T²), 메모리 O(T) 시간 O(T·d_k·d_v), 메모리 O(1)

4.1.2 비교 기준선

차원 softmax attention KDA 메모리 token별 KV 고정 행렬 S 갱신 없음(읽기 전용) 있음(선택적 기록) 망각 없음(softmax 정규화가 천연적으로 커버) 있음(channel 수준 decay)

변경이 가져오는 이점: O(1) 메모리, O(T) 시간; 대가: 검색이 압축 근사로 변함.

4.1.3 계산 복잡도 분석

KDA는 전형적인 "순환 시퀀스 모델"이다: 학습/prefill은 전체 시퀀스 chunkwise 순방향 (각 위치마다 출력해야 하며, loss 계산에 사용)을 사용하고; 추론 decode는 단일 스텝 순환 (자기회귀적으로 token을 하나씩 생성, 상태 전달)을 사용한다. 다리 역할은 chunkwise(블록 내 병렬, 블록 간 직렬)다.

- FLOPs(단일 token): KDA = O(d_k·d_v) (상태 갱신 + 출력 판독) vs softmax = O(T·d_k) (전체 이력과 비교 필요). T가 클 때 KDA가 완승한다.

- 메모리: KDA 상태 = d_k×d_v /head(고정); softmax KV = T×(d_k+d_v) /head(선형 증가).

- 파라미터: KDA는 순환 자체에 신규 파라미터 없음(투영 공유).

4.1.4 파라미터 민감도

- head_dim_k (기본값 128): 상태 정방행렬 크기와 "key 기억 방향"의 용량을 결정한다. 키우면 → 상태 용량↑이나 FLOPs↑; 너무 작으면 → 용량 부족.

- head_dim_v (기본값 128): value 정보 용량.

- 논문 경험: d_k=d_v=128은 Kimi Linear가 튜닝을 거친 후의 기본값이며, 96 head(d=7168)와 조합된다.

4.1.5 반대 측 논쟁

- 가장 강한 반대 이유: 고정 행렬의 압축은 필연적으로 정보 손실을 초래하며, 초장문 문맥의 정확한 검색(예: "40만 번째 token에서 언급된 어떤 구체적 숫자")은 실패할 것이다.

- 만약 실패한다면, 가장 가능성 높은 잘못된 가정: "압축 가능한 기억이 작업에 필요한 정보를 충분히 표현할 수 있다"는 가정. 정확한 token별 역추적이 필요한 작업(복잡한 agentic 궤적, 코드베이스 검색)에서 이 가정은 취약하다.

- 승리 조건: 기억이 근사 가능하고, 의미가 밀집되어 있으며, 문맥이 길고 token 의미 중복성이 높은 경우(대화, 장문 요약). 실패 조건: 정확한 위치 검색이 필요하거나 핵심 정보가 희소한 경우.

- 더 견고한 방법: 상태 확장(state expansion), 또는 KDA+MLA 혼합 비율을 높이는 것(3:1이 논문 최적이나, 작업에 민감함).

- 피해자: 장거리 정확 검색 작업 측; 수혜자: 1M 문맥 추론 처리량 측.

4.1.6 엔지니어링 구현

의사 코드(줄 단위):

for t in 1 ..T: state *= exp (g[:, :, t]) # Diag(alpha) S channel 감쇠 sk = einsum( 'bhd,bhdv->bhv' , k[:, :, t], state ) # S^T k 예측 delta = beta[:, :, t] * (v[:, :, t] - sk) # 잔차 × 학습률 state += k[:, :, t] * delta # rank-1 외적 갱신 out = einsum( 'bhd,bhdv->bhv' , q[:, :, t] , state ) # 판독

프로젝트 실제 코드( kda_ref.py 의 kda_recurrent ) :

def kda_recurrent(q, k, v, g, beta, *, initial_state =None, reset_mask =None, dtype=torch.float64): """Eq. 1 的字面递推。是其他所有 KDA 路径的 ground truth。""" qf, kf, vf, gf, bf = _as_float(q, k, v, g, beta, dtype=dtype) b, h, t, dk = qf.shape dv = vf.shape[- 1 ] state = (torch.zeros(b, h, dk, dv, dtype=dtype, device=q.device) if initial_state is None else initial_state.to(dtype).clone()) if t == 0 : return vf.new_zeros(b, h, 0, dv), state outputs = [] for i in range(t): if reset_mask is not None: r = reset_mask[:, i].view(b, 1 , 1 , 1 ) state = torch.where(r, torch.zeros_like(state), state) alpha = gf[:, :, i].exp() # 逐 channel 保留因子 state = alpha.unsqueeze(- 1 ) * state # Diag(alpha) S ki = kf[:, :, i] # 当前 key bi = bf[:, :, i].unsqueeze(- 1 ) # 当前写入强度 sk = torch.einsum( "bhd,bhdv->bhv" , ki, state) # S^T k delta = bi * (vf[:, :, i] - sk) # 写入方向 = 强度 * 残差 state = state + ki.unsqueeze(- 1 ) * delta.unsqueeze(- 2 ) # 外积更新 outputs.append(torch.einsum("bhd,bhdv->bhv", qf [:, :, i] , state)) return torch.stack(outputs, dim = 2 ), state

오픈소스 참조 구현과 비트 단위로 대조한다:

- 본 프로젝트의 kda_recurrent는 Kimi Linear 논문 Eq.1과 완전히 일치하며, kimi3-master의 naive_recurrent_kda와 대조했을 때 max diff = 0.0이다(동일한 공식, 동일한 dtype 경로).

- SGLang fused_recurrent_gated_delta_rule_fwd_kernel의 IS_KDA 분기: 논문의 KDA 점화식과 동등하며, 본 프로젝트의 CPU 버전은 그와 로직이 비트 단위로 정렬되어 있다(차이는 dtype/융합에만 있음).

4.1.7 자주 나오는 질문 & 논의

Q: in-context learning, 즉 암묵적 최적화

KDA는 토큰을 하나 처리할 때마다 그것으로 상태에 한 번의 망각이 포함된 온라인 업데이트를 수행한다. 상태는 일종의 소형 모델 파라미터에 해당하고, 현재 컨텍스트가 바로 그 학습 데이터다. 따라서 컨텍스트 학습은 다음과 같이 볼 수 있다: 토큰을 읽어들임 → 내부 상태에 작은 걸음의 최적화를 한 번 수행 → 업데이트된 상태로 다음 토큰을 처리. 감쇠는 온라인 학습에서의 정규화와 유사하여, 상태 파라미터가 무한히 커지는 것을 막는다.

Q: 왜 head별로 독립적인가

KDA는 각 head마다 독립적인 상태 행렬 S를 유지하며, 이는 세 군데에 나타난다:

- 초기화: 상태 텐서의 형상은 (b, h, dk, dv)이다. 제1차원이 head 수다. 각 head에 dk×dv 행렬 하나씩.

- 점화식 업데이트: 감쇠 alpha는 (b, h, dk)이다(head별·channel별). state의 업데이트는 head 차원을 따라 독립적으로 진행된다. head 사이에는 어떠한 교환도 없다. 이것이 "독립"의 핵심이다.

- 캐시: KDACache.state의 형상은 (batch, num_heads, head_dim_k, head_dim_v)이다. 즉 head마다 상태 행렬 하나.

서로 다른 head는 서로 다른 기억 패턴을 학습할 수 있다(어떤 것은 근거리를, 어떤 것은 장거리를, 어떤 것은 특정 내용을 주목한다). 만약 head 간에 상태를 공유한다면 이러한 분화는 지워질 것이다. 이것이 바로 HeadRMSNorm이 존재하는 이유이기도 하다. 각 head의 S가 독립적으로 진화하므로, 그들이 출력하는 스케일은 크게 다를 수 있고, head별로 정규화할 필요가 있다.

Q: 왜 beta는 head별 스칼라인가?

온라인 학습 관점에서 β_t는 곧 학습률이다: S_t = S_{t-1} + β_t · gradient. 만약 beta를 channel별 벡터로 만든다면 표현력은 더 강해지겠지만, 그 대가는 다음과 같다: 상태 업데이트에 더 많은 게이트가 필요하고, 메모리와 대역폭이 증가하며, chunkwise 분할 공식이 더 복잡해지고, kernel을 더 규칙적으로 만들기 어려워진다. head별 스칼라는 하나의 절충이다. 각 head가 이번 업데이트의 전체 크기를 결정하고, channel 사이의 차이는 Diag(α_t)가 담당하며, beta는 여전히 head级 학습률로 볼 수 있다.

Q: 왜 beta는 반드시 현재 토큰에만 의존해야 하는가?

만약 beta가 전체 히스토리에 의존한다면, 업데이트 진폭을 결정하기 위해 먼저 전체 점화식을 완료해야 하므로 chunk 내 병렬화가 깨진다. beta를 현재 토큰으로 즉시 예측하게 함으로써 일부 전역 적응 능력을 희생하는 대신, 다음을 얻는다: 단일 스텝으로 계산 가능, chunkwise로 전개 가능, 상태를 병렬로 귀결할 수 있음, kernel이 더 규칙적임. 이것은 병렬화의 전제이지 임의의 단순화가 아니다.

Q: 감쇠가 이미 있는데, 왜 Delta 소거가 또 필요한가?

이들은 역할 분담이 다르다. 감쇠는 오래된 상태를 천천히 잊는 것이고, Delta는 현재 Key를 따라 국소적 연관을 빠르게 덮어쓰는 것이다. 비유하자면, 감쇠는 방 전체의 소음을 점진적으로 줄이는 것 같고, Delta는 지우개로 특정 행 하나를 지우는 것 같다. 하나는 '시간'을, 다른 하나는 '방향'을 담당한다. 다만 감쇠에는 하한이 있어 한 번에 완전히 비울 수 없다. 이 빠른 삭제 능력이 부족한 부분을 Delta 항이 보충한다. 그러나 Delta는 보통 한 번에 한 방향만 지울 수 있으므로, '빠르면서도 선택적이고 임의의 부분공간까지 삭제할 수 있는' 능력을 한 단계에서 동시에 구현할 수는 없다.

Q: 점화식이 수치적으로 불안정해지지 않을까?

아니다. 중간 계산은 fp32/fp64로 강제되고, alpha는 항상 < 1이므로 상향 오버플로가 없다. chunkwise vs recurrent 대조 시험에서 chunk=1..64가 전부 일치한다.

Q: beta가 1에 가까워지면 무슨 일이 일어나는가?

해당 방향의 이전 기억을 완전히 덮어쓰며, '디스크에 덮어쓰기'와 비슷하다. diagnostics에는 beta 포화도 모니터링이 있다.

Q: 고정 state가 왜 KDA의 근본적 한계인가?

하나의 고정 state는 임의로 긴 역사를 담아야 한다. 역사 속의 독립적 사실이 많아질수록 더 쉽게 나타나는 현상은 다음과 같다. 서로 다른 두 역사가 같은 상태로 사상됨, 유사한 Key가 상태 방향을 공유함, 아주 오래된 정보가 감쇠됨, Query가 압축된 토큰 정체성을 되찾을 수 없음. 따라서 KDA의 고정 상태는 효율의 원천인 동시에 정보 용량의 상한이다.

Q: query는 이미 그어진 천장 아래에서만 방향을 고를 수 있다

KDA는 역사를 기록할 때부터 각 정보의 생존 정도를 결정한다. 미래의 Query는 남은 상태 안에서만 방향을 고를 수 있다. 만약 어떤 정보가 기록될 때 beta가 매우 작았고, 이후 여러 차례 감쇠되었으며, 다른 정보와 덮어쓰기가 발생했다면, 미래의 Query가 아무리 관련이 깊더라도 다시 읽어들일 원본 토큰이 없다.

유계 게이팅은 1/gamma를 데이터 무관한 하드 상한 exp(-tile*g_min)=exp(80)=5.54e34 < BF16 상한으로 바꾸어, 각 인과 tile(대각 포함)이 무조건적으로 조밀한 Tensor-Core matmul로 컴파일될 수 있게 한다 — 수치 안전성으로 전체 조밀 스케줄링을 확보하는 것이다.

KIMI-K3-Eq5

4.2.1 원리

직관: Kimi Linear의 망각 속도는 '밑 빠진 독'(softplus에 하한이 없음)이라, 16개의 토큰만 기록해도 float 언더플로까지 감쇠된다. K3는 망각 속도에 '하한' g_min을 두어 감쇠가 너무 빠르지 않고 역수가 오버플로하지 않도록 보장한다.

전문가 관점:

g_t = g_min · sigmoid ( exp (A_h) · z_t ) (Eq .5 )

- z_t = W_up W_down x_t + b_α : 저랭크 decay logits.

- A_h : head별 로그 스케일( a_log , 초기 0이므로 초기 exp(A)=1 ).

- sigmoid는 exp(A)·z를 (0,1)로 압축하고, 다시 g_min<0을 곱한다 → g ∈ (g_min, 0) .

하한 설계는 FlashKDA 조밀화의 핵심 전제다. g ≥ g_min이므로, 한 tile(16 token) 내 누적 log-decay ∈ (16·g_min, 0) 이고, 역수 1/γ ≤ exp(−16·g_min)=exp(80)≈5.5e34 로, BF16 최댓값 3.4e38이 이를 담을 수 있다. 하한이 없으면 1/γ가 오버플로하여 조밀화할 수 없다.

유계 감쇠 = '수치적 조밀화 가능성'을 모델 설계의 하드 제약으로 삼는 것으로, 알고리즘-하드웨어 협동 설계의 전형이다. 즉 수학을 먼저 정하고 하드웨어 적합성을 찾는 것이 아니라, 수학적 형태를 하드웨어 정밀도가 직접 결정하는 것이다.

4.2.2 비교 기준

차원 Kimi Linear(무계) K3(유계 Eq.5) log-decay 범위 (−∞, 0) (g_min, 0) 대각 tile 명시적 position-pair 경로 조밀 matmul 수치 위험 언더플로/역수 오버플로 제어 가능

4.2.4 계산 복잡도 분석

- 추가 FLOPs 없음(sigmoid는 값싼 원소별 연산).

- 파라미터: a_log (head당 1개) → 새로 추가되는 h개의 스칼라로, 무시할 수 있음.

4.2.5 파라미터 민감도

- g_min (기본값 −5.0): 키우면(예: −3) 더 보수적인 감쇠로 수치가 더 안전하지만 망각이 느려짐; 줄이면(예: −8) 감쇠가 더 빨라져 tile 내 exp(−16·8)=exp(128)이 BF16에서 오버플로하고 safe_chunk_size가 10으로 떨어짐(수치 실측은 §4.2.7 참조).

- alpha_init (기본값 0.99): 초기 retention 인자로, alpha_bias = log(log(α_init)/g_min / (1−…)) 를 통해 모델이 '거의 항등 감쇠'에서 출발하게 하여 초기 학습 시 상태가 빠르게 붕괴하는 것을 막는다.

4.2.6 엔지니어링 구현

의사 코드:

z = alpha_up(alpha_down(x)) + alpha_bias # (B,T,H·Dk) scale = exp(a_log) # (H,) g = g_min * sigmoid(scale * z.view(B,T,H,Dk))

프로젝트 실제 코드( kda.py 의 log_decay ):

def log_decay ( self, x: Tensor ) -> Tensor: """Eq. 5：g = g_min * sigmoid(exp(A_h) z)，범위는 엄격히 (g_min, 0) 안에 있다. 이 하한은 이 구현이 SGLang/FLA의 lower_bound 메커니즘과 정렬될 수 있게 하는 핵심이다: log-decay에 하한이 있기 때문에 비로소 chunkwise에 필요한 1/gamma가 오버플로되지 않고, 조밀화(稠密化)가 성립한다. """ b, t, _ = x.shape z = self.alpha_up(self.alpha_down(x)) + self.alpha_bias # (B, T, H*Dk) z = z.view(b, t, self.num_heads, self.head_dim_k) # head별로 다시 분해 scale = self.a_log.exp().view( 1 , 1 , self.num_heads, 1 ) # exp(A_h) 브로드캐스트 return self.cfg.g_min * torch.sigmoid(scale * z)

safe_chunk_size 함수( kda_ref.py )는 g_min에 근거해 현재 dtype의 최대 안전 chunk를 계산한다(기본 g_min=−5일 때 bf16/fp32는 16 = 논문 tile을 반환).

def safe_chunk_size ( g_min: float , dtype: torch.dtype, *, headroom: float = 1e3 ) -> int : """"그 누적 역수 감쇠가 여전히 dtype에 담길 수 있는" 최대 tile 길이를 구한다.""" limit = _DTYPE_MAX[dtype] / headroom return max ( 1 , int (math.log(limit) / -g_min))

실측 수치 :

- max_reciprocal_decay(-5, 16) = 5.54e+34 （= exp(80)）

- safe_chunk_size(-5, bf16) = 16 （정확히 논문 tile!）

- safe_chunk_size(-5, fp32) = 16

- safe_chunk_size(-5, fp16) = 1 （fp16 상한 65504가 너무 작아 tile=1이어야 담긴다）

- safe_chunk_size(-8, bf16) = 10 （g_min이 더 음수일수록 안전 chunk는 더 작아진다）

4.2.7 수치 안정성 검증 실험

실패 case: 무계(無界) 감쇠에서의 chunkwise 오버플로

softplus 무계 감쇠로 chunk=16을 돌려, 1/γ가 inf로 오버플로되는 것을 관찰:

게이팅 g.min 1/gamma max 출력 유한? 유계 g_min·sigmoid（K3） -4.80 6.15e+22 ✅ True 무계 softplus z~N(0,1) -3.22 1.36e+08 ✅ True 무계 softplus z+3 -6.18 6.83e+25 ✅ True 무계 softplus z+6 -9.18 inf ❌ False（2048개의 NaN） 무계 softplus z+9 -12.18 inf ❌ False（2048개의 NaN）

fp32 오버플로 임계는 대략 z≈6 부근이다（ exp(16*6)이 이미 3.4e38을 초과한다）. 이것은 데이터 의존적이다——같은 tile=16이라도 z가 조금만 커지면 붕괴한다.

chunk=64 —— 유계 버전조차 오버플로（safe_chunk_size 경계를 입증）

chunk=64일 때 exp(64*5)=exp(320)은 fp32 상한을 훨씬 넘어서고, 유계 버전 역시 마찬가지로 inf/NaN이다. 이는 safe_chunk_size(-5, fp32)=16의 경계가 진짜임을 정확히 보여준다——안전 tile을 벗어나면 반드시 붕괴한다.

4.2.8 논의

Q1: 왜 하한이 있는 유계 감쇠가 필요한가?

만약 감쇠가 0에 무한히 가까워질 수 있다면, 한 걸음 만에 전체 상태를 거의 비워버릴 수 있다: S_t ≈ 0 . 그 결과는: 그래디언트가 갑자기 사라질 수 있고; 장거리 정보가 예기치 않게 지워질 수 있으며; 저정밀도에서 수치 불안정이 쉽게 나타날 수 있고; 서로 다른 token 사이의 상태 스케일 변화가 너무 커질 수 있다. 따라서 감쇠가 어떤 안전 하한보다 작아지지 않도록 제한할 필요가 있다. 그것은 "한 걸음에 완전히 잊어버리는" 능력을 희생하는 대신 상태의 연속성과 수치적 안전을 얻는다.

Q2: 왜 감쇠를 log 공간에서 계산해야 하는가

여러 token의 감쇠는 연속 곱해진다: α_{s→t} = ∏_{i=s+1}^{t} α_i . 만약 1보다 작은 수를 직접 많이 곱하면 쉽게 언더플로되어 0이 된다. 그래서 log 공간으로 바꾼다: log α_{s→t} = Σ_{i=s+1}^{t} log α_i . 곱셈이 덧셈이 되어 수치 범위를 더 쉽게 통제할 수 있다.

Q3: 왜 Kimi Linear는 무계 softplus를 쓰는가?

먼저 공식부터 풀어보자

Kimi Linear（KDA v1）의 게이팅:

g = -exp (A) * softplus (z) （하한 없음: z -> +inf일 때 softplus (z) -> z, g는 밑이 없음） alpha = exp (g) channel별 보존 인자, 0 < alpha <= 1

chunkwise 병렬 공식（Eq.3-4）에서 핵심량은 "chunk 시작점부터 위치 r까지의 누적 감쇠 연속 곱"이다:

gamma^{1->r} = prod( s = 1 ..r) alpha_s = exp(sum g_s) <= 1 khat_r = k_r / gamma^{ 1 ->r} （Eq. 4의 K/Gamma 재스케일링）

KIMI-K3-Eq3

KIMI-K3-Eq4

두 방향에서 동시에 문제가 생긴다

① 누적 감쇠 곱의 언더플로 : gamma^{1->r} = exp(sum g_s) . g = -softplus(z)는 하한이 없기 때문에, z가 조금만 커지면（모델이 "빠른 망각"을 학습하면） sum g_s가 매우 음수가 되어 exp가 곧바로 0으로 떨어진다（언더플로）.

② 그 역수 상향 오버플로: chunkwise에서는 khat_r = k_r / gamma 를 계산해야 하는데, 이는 k_r * (1/gamma) 와 같다. 1/gamma = exp(-sum g_s) 이며, sum g_s 가 매우 음수일 때 1/gamma 는 무한히 증가한다 — float 상한을 넘으면 inf 가 된다.

그 결과는 다음과 같다.

- 명시적 position-pair 경로: 즉 chunk 내 위치를 둘씩 명시적으로 짝지어 계산하는 방식이다(position-pair 는 "row/col 이 모두 chunk 내 구체적 위치"인 대각 블록을 가리킨다). 쌍마다 어텐션/기입 관계를 개별 계산하며, FlashAttention 식의 희소/명시적 경로이다.

- 인과 삼각 분해 시, 대각 tile 은 "token 과 그 자신 및 같은 chunk 내 다른 위치"의 상호작용을 처리해야 한다. 만약 1/gamma 가 오버플로하면, 대각 tile 을 일반적인 조밀 행렬곱에 안전하게 넣을 수 없다(Tensor-Core 안에서 중간값이 폭발한다).

- 그래서 Kimi Linear 는 어쩔 수 없이 우회해야 했다. 대각 tile 은 조밀 matmul 을 쓰지 않고, 명시적 position-pair 계산 경로로 처리한다(각 위치 쌍을 개별 계산하여 1/gamma 를 한꺼번에 증폭시키지 않는다). 대가는 텐서 코어 가속을 잃어 대각 tile 이 느려지는 것이다.

즉, 무계 softplus 는 g 를 밑 빠진 독으로 만들고 → 1/gamma=exp(-sum g) 는 상한이 없으며(데이터 의존적이라 z 가 조금만 커져도 inf) → 대각 tile 은 조밀 matmul 을 안전하게 쓸 수 없어 명시적 position-pair 경로로 되돌아갈 수밖에 없다(느리다).

Q4: K3 는 어떻게 느린 경로를 제거하는가?

초기 chunkwise 계산은 감쇠 수치가 극단적으로 되는 것을 막기 위해 대각 tile 을 특수 처리해야 했다.

K3 의 사고방식은 근원에서 log‑decay 를 제한하는 것이다: 감쇠 logit → 하한 g_min 설정 → 역수 재스케일링 인자가 항상 표현 가능 → 대각 및 비대각 tile 이 통일된 행렬 형식을 사용.

만약 g_min=-5 라면, 가장 극단적인 감쇠도 표현하기 어려운 수치 영역에 들어가지 않는다는 뜻이다. 이로써 별도의 느린 보정 경로를 피할 수 있고, 대각 tile 도 규칙적인 GEMM 에 더 가까워지게 하며, Tensor Core 계산을 통일하고, 약간의 정밀도 제약을 대가로 더 높은 하드웨어 활용률을 얻는다. 이는 "수학적으로 인과 관계를 소멸시킨" 것이 아니라, 인과 관계의 수치적 표현을 충분히 안정화하여 동일한 행렬 kernel 로 처리할 수 있게 한 것이다.

Q5: 왜 "모든 인과 tile(대각 포함)이 조밀 Tensor-Core matmul 을 쓸 수 있는가?

인과 분해 후, A = Tril[(Q*Gamma)(K/Gamma)^T] 의 각 대각 tile 은 블록 내 matmul 을 한 번씩 계산해야 하며, 중간값에 K/Gamma (즉 1/gamma 증폭)가 포함된다. 1/gamma <= 5.54e34 가 데이터 무관하며 BF16 에 담기는 경성 상한이라면, 그렇다면:

- 컴파일러/커널은 무조건 (어떤 입력, 어떤 token 에 대해서도) 이 대각 tile 을 조밀 Tensor-Core matmul 로 컴파일할 수 있으며, 런타임에 "여기서 오버플로가 나지 않을까"를 판단할 필요가 없다;

- 명시적 position-pair 경로로 되돌아갈 필요 없이, 대각 tile 도 일반 tile 과 똑같이 텐서 코어 가속을 쓴다.

이것이 K3 가 Kimi Linear 보다 더 "조밀"하고 더 빠를 수 있는 이유다 — 수치적 안전성이 전면적 조밀 스케줄링을 사들인 것이다.

Q6: 왜 유계 감쇠는 데이터 의존적 상한을 데이터 무관 경성 상한으로 바꾸는가

K3 의 log_decay:

g = g_min * sigmoid (exp(A_h) * z) (Eq. 5 , g_min = - 5 )

sigmoid 는 출력을 (0,1) 로 압축하고, 다시 g_min<0 을 곱하면 → g ∈ (g_min, 0) 이 엄격히 성립하며, 입력 z 와 무관하다.

그래서 하나의 tile(16 token)을 넘는 누적 log-decay:

tile 내 누적 log-decay = sum_{ 16 개 위치} g_s ∈ ( 16 *g_min, 0 ) = (- 80 , 0 ) 1 /gamma = exp (-sum g_s) <= exp ( 80 ) ≈ 5.5406 e34 （데이터 무관 경성 상한!）

BF16 최대 유한값 ≈ 3.3895e38 , FP32 ≈ 3.4028e38 . 5.54e34 < 3.39e38 → 담긴다. 이것이 exp(80)=5.5e34 가 여전히 BF16 범위 안에 있는 유래다.

Q7: a_log 와 dt_bias 의 역할

상태 공간 모델과 KDA 계열 게이팅에서 흔한 방식은 "훈련은 쉽지만 제약이 없는 파라미터"를 "안정성 요구를 충족하는 감쇠량"으로 매핑하는 것이다.

- a_log : 일반적으로 기초 시간 척도의 로그 파라미터를 나타내는 데 쓰인다. 로그로 저장하는 이점은 파라미터가 여러 자릿수에 걸쳐 포괄될 수 있고, 최적화 시 빠른 감쇠와 느린 감쇠를 더 쉽게 표현할 수 있으며, 지수 매핑을 통해 기초 속도가 양수임을 보장할 수 있다는 것이다. 직관적으로 A_log 는 "이 채널이 원래 얼마나 오래 기억하는가"를 결정한다.

- dt_bias : 각 입력 스텝 또는 시간 간격에 대한 오프셋이다. 현재 token 이 도착할 때 기초 시간 척도가 어떻게 조정되는지를 제어한다. 직관적으로: A_log 는 기본 기억 수명, dt_bias 는 현재 스텝의 시간 진행 오프셋, 입력 관련 게이트는 현재 token 을 더 빠르게 또는 더 느리게 갱신해야 하는지를 제어한다.

세 가지가 왜 하한 유계 감쇠를 구성할 수 있는가: 보통 A_log → 양의 기초 감쇠율; dt_bias + 현재 입력 → 양의 유효 스텝 길이; 유효 스텝 길이 × 감쇠율 → log‑decay 로 매핑 → 최소 log‑decay 하한 적용 → alpha 획득의 흐름을 거친다. softplus, exp 또는 동등한 양수값 매핑을 사용하면 시간 스텝 길이와 감쇠율에 불법 부호가 나타나지 않음을 보장할 수 있고, 여기에 log‑decay 에 하한 절단을 더하면 감쇠가 한 번에 비워지지 않고, 감쇠가 여전히 입력에 따라 변할 수 있으며, 장시간 재귀에서 더 쉽게 안정성을 유지하도록 보장할 수 있다.

4.2.9 반대 측 변론자

- 하한 g_min=−5 는 각 token 이 최소 e^{-5}≈0.0067 의 감쇠 비율을 유지한다는 뜻이다 — 초장기 컨텍스트에서는 정보를 "잊지 못함"이 잡음을 누적시킬 수 있다.

- 만약 실패한다면, 그 잘못은 "감쇠 하한이 recency 모델링을 크게 해치지 않는다"는 가정에 있다.

- g_min을 작게 조정하면(예: -8) 망각을 가속화할 수 있지만, safe_chunk_size가 10으로 떨어져 tile=16이면 오버플로 → 조밀화(densification)가 무효화된다(§4.2.10 그림 범례 좌하단 참조).

- 성공: 컨텍스트가 안정적으로 점진적 감쇠를 필요로 할 때; 실패: 강한 recency 편향이 필요한 작업.

- 견고한 대안: g_min을 학습 가능하게 설정.

4.2.10 그림 범례

plot_decay_stability

네 칸 읽는 법:

- 좌상단: 1/gamma 상한 vs tile 크기(로그). 유계 버전은 데이터와 무관한 직선 exp(-g_min*tile)이며, tile=16일 때 BF16 상한 아래에 있다; 무계 버전에는 이 직선이 없다(무한히 위로 커질 수 있다).

- 우상단: 실측 tile=16——무계 z+6/z+9는 inf로 폭발(빨간 막대), 유계 버전은 안전.

- 좌하단: safe_chunk_size vs g_min——g_min이 더 음수일수록 허용 chunk가 더 작아진다(-5→16, -8→10).

- 우하단: 실패 case——chunk=64일 때는 유계 버전조차 오버플로한다(안전 경계를 넘으면 반드시 붕괴).

설계 철학: KDA는 **"12288개의 channel 손잡이 + 96개의 head 온도"**——channel-wise는 세밀한 제어를, head-wise는 거친 조절을 제공하며, 저차원 분해가 파라미터 수를 제어 가능하게 만든다.

질문 답변 alpha_rank 크기 512, 구성 가능, 0 < rank <= k_dim을 만족해야 함 어떻게 설정하는가 KDAConfig에 정의되어 있으며, 표현력과 파라미터 수를 균형 있게 맞춤 어떻게 압축하는가 W_down: (dim, 512) → W_up: (512, 12288), LoRA 스타일 어떻게 매핑하는가 alpha_bias는 channel별(12288개), exp(A_h)는 head별(96개)의 스케일링 팩터

저차원 decay는 "LoRA 사고를 모델 아키텍처 내부에서 재사용"한 것이다——미세조정 시에만 압축하는 것이 아니라, 파라미터를 생성할 때 압축한다. 이는 파라미터 효율적 설계가 아키텍처 수준으로 일반화된 한 가지 예증이다.

KIMI-K3-Eq2

4.3.1 원리

직관: GDN은 하나의 head에 하나의 망각 손잡이만 쓰지만, KDA는 각 channel에 손잡이를 하나씩 준다(총 d_k개). 하지만 d_k개의 손잡이를 각각 x로부터 직접 생성하면 파라미터가 너무 많아지므로, 저차원 병목으로 압축한다.

z = W_up W_down x + b_α, W_down ∈ R^{rank×d}, W_up ∈ R^{d_k×rank} (Eq. 2 )

rank = alpha_rank (기본값 512)는 d_k = 96×128=12288보다 훨씬 작다. q = q_proj(x) 후에 SiLU + L2 정규화를 수행한다. 저차원은 decay 파라미터가 하나의 저차원 부분공간을 공유하게 하면서 동시에 channel 수준의 차별화를 유지한다——이는 "파라미터 효율로 세밀한 제어를 교환"하는 것이다.

"감쇠 패턴의 rank는 128(또는 구성된 alpha_rank)로 제한된다"는 것은 직관적으로 이렇게 이해할 수 있다: 모델이 각 head마다 완전히 독립적인 망각 규칙을 따로 학습하는 것이 아니라, 먼저 최대 alpha_rank 차원의 "망각 의도 공간"을 학습하고, 각 head가 이 공간에서 자신의 감쇠 방식을 읽어 들인다. 한 밴드 그룹이 128개의 손잡이를 공유하는 것과 같다: 각 head는 자기만의 볼륨을 가질 수 있지만, 모든 head의 변화는 동일한 기본 손잡이 세트에서 비롯된다; 파라미터가 더 적고, 구조가 더 안정적이며; 서로 다른 head의 망각 행동은 완전히 독립적이지 않다.

장점: 파라미터와 계산을 줄인다; 서로 다른 head가 통계적 규칙성을 공유하게 한다; 하드웨어 행렬 연산에 유리하다.

단점: 실제 작업이 rank가지 서로 독립적인 망각 패턴을 초과하여 필요로 하면, 저차원 공간이 병목이 될 수 있다; 서로 다른 head 사이에 상관관계가 지나치게 강해질 수 있다; 이는 표현력을 효율로 교환하는 전형적인 절충이다.

논리도

channel-wise decay

alpha_rank 설정

현재 구성 ( config.py:31 ):

alpha_rank: int = 512

어떻게 alpha_rank로 압축하는가? (저차원 분해)

alpha_rank

핵심 설계: alpha_up 초기 가중치를 0으로 설정( kda.py:144 )

self .alpha_up.weight.zero_() # 초기 log -decay는 정확히 편향과 같다

이는 다음을 의미한다:

- 훈련 초기: z = 0 + alpha_bias = alpha_bias, 모델은 "근사 항등 감쇠"에서 시작한다

- 훈련 과정: W_up이 점차 학습되어, W_down의 압축 표현을 통해 alpha_bias를 변조한다

- LoRA 스타일: 먼저 모델을 작동시킨 뒤, 점차 저차원 보정을 도입한다

어떻게 각 head에 매핑하는가? (브로드캐스트 + head별 변조)

이는 2단계 구조다:

1단계: channel-wise(공유 기저) ───────────────────────────────── z = W_up @ W_down @ x + alpha_bias │ ▼ ( B , T , H * Dk ) = ( B , T , 12288 ) │ ▼ view ( B , T , H , Dk ) = ( B , T , 96 , 128 ) │ └──► 각 channel은 자체적인 감쇠 편향을 가짐 (12288개의 "노브") 2단계: head-wise(헤드별 스케일링) ───────────────────────────────── a_log : ( H , ) = ( 96 , ) ←── 학습 가능한 파라미터, 0으로 초기화 │ ▼ exp ( A_h ) : ( H , ) ←── 초기값 1, 양수·음수 모두 가능 │ ▼ view ( 1 , 1 , H , 1 ) ←── ( B , T , H , Dk )로 브로드캐스트 │ ▼ scale * z : ( B , T , H , Dk ) │ ▼ g = g_min * sigmoid ( scale * z )

수학적 표현:

head h의 d번째 channel에 대하여:

g [b,t,h,d] = g_min * sigmoid ( exp(A_h) * z [b,t,h,d] ) = g_min * sigmoid ( exp(A_h) * (W_up @ W_down @ x + b) [b,t,h*Dk+d] )

직관적 이해:

- alpha_bias: 12288개 channel 각각의 "기준 감쇠 선호"(공유 파라미터)

- W_down/W_up: 입력 x에 따라 이러한 선호를 동적으로 변조(저랭크 압축)

- exp(A_h): 각 head의 "전역 온도"로, 해당 head 전체 감쇠의 민감도를 제어

4.3.2 비교 기준

차원 GDN KDA gate 1개의 스칼라/head d_k개/head 파라미터 출처 직접 투영 저랭크 투영

4.3.3 반대 측 논거

- 저랭크 병목이 decay의 다양성을 제한한다(모델이 중복된 rank 부분공간을 학습할 수 있음).

- 만약 실패한다면, rank가 channel 차이를 포착하기에 충분하다는 가정이 틀린 것이다.

- 성공 조건: channel 간 decay가 매우 상관관계가 높음. 실패 조건: 직교하는 channel 기억이 필요함.

- 견고한 대안: 직접 투영(파라미터는 많지만 더 유연함).

- 피해를 보는 쪽: 극단적 장문맥 기억 과제.

4.3.4 계산 복잡도

- 파라미터: d×rank + rank×d_k. 기본 rank=512: 4096×512 + 512×12288 ≈ 8.4M로, 비중이 작다.

4.3.5 파라미터 민감도

- alpha_rank(기본 512): 키우면→표현력↑파라미터량↑, 줄이면→표현 부족.

4.3.6 엔지니어링 구현

프로젝트 코드 kda.py에서 alpha_down/alpha_up/alpha_bias를 사용하며, alpha_up 가중치를 0으로 초기화한다(LoRA 관례). 초기화 reset_parameters에서는 방정식을 풀어 retention이 alpha_init에서 시작하도록 한다:

저랭크 decay logits(Eq. 2): z = W_up W_down x + b_alpha. self.alpha_down = nn.Linear(dim, cfg.alpha_rank, bias= False ) # rank로 압축 self.alpha_up = nn.Linear(cfg.alpha_rank, k_dim, bias= False ) # k_dim으로 다시 전개 self.alpha_bias = nn.Parameter(torch.empty(k_dim)) # channel별 편향 def reset_parameters ( self ) -> None : """모델이 처음에 "근사 항등 감쇠"의 안정 상태에 놓이도록 맞춤 초기화한다. - ``alpha_up`` 가중치를 0으로: 초기 log-decay가 정확히 편향과 같아진다(LoRA 관례: 먼저 up 투영을 꺼 두면 그래디언트는 여전히 그곳에 도달할 수 있고, down 투영이 0이 아니게 되면 곧바로 작동하기 시작한다). - ``g_min * sigmoid(z) = log(alpha_init)`` 을 풀어 z를 구하여, retention 인자가 ``cfg.alpha_init`` 에서 시작하도록 하면서 동시에 ``exp(g_min)`` 하한을 존중한다: ``alpha_bias = log( target/(1-target) ), target = log(alpha_init)/g_min`` """ with torch.no_grad(): self.alpha_up.weight.zero_() target = math.log(self.cfg.alpha_init) / self.cfg.g_min self.alpha_bias.fill_(math.log(target / ( 1.0 - target)))

오픈소스 참조와 대조: kimi3-master 역시 저랭크 decay projection(decay_down/up)을 사용한다.

4.3.7 논의

Q1: 왜 g_proj는 전랭크인데, 감쇠 게이트는 저랭크를 유지하는가

K3는 새로 추가한 출력 게이트 g_proj를 저랭크에서 전랭크로 확대했으며, 핵심 이유는 다음과 같다: Output gate는 최종 출력의 각 channel마다 "얼마나 쓸지"를 독립적으로 결정해야 한다. AttnRes의 value 경로는 정규화를 하지 않고, 블록 요약은 주로 합산에 의존한다. 그러므로 output gate의 크기는 해당 레이어가 심층 기억에서 차지하는 몫에 직접 영향을 미친다. 만약 출력 게이트가 너무 저랭크라면, 많은 channel이 같은 조절 방향을 공유할 수밖에 없어 세밀한 제어가 불가능하다.

감쇠 게이트는 다르다. 감쇠는 "이 토큰이 어떤 상태 방향을 장기적으로 유지하고 어떤 방향을 빠르게 망각하게 할지"를 기술한다. 이러한 "망각 의도"는 일반적으로 최종 출력 채널보다 훨씬 저차원이다. 따라서 먼저 저랭크 벡터로 망각 의도를 표현하고, 그다음 서로 다른 head / 채널이 선형으로 읽어내도록 할 수 있다.

정리:

게이트 제어 대상 독립성 요구 구조 Output gate 최종 출력 채널의 쓰기 진폭 높음, 채널별 조정 필요 전랭크 Decay gate 상태의 기억 수명 / 망각 의도 일부 구조 공유 가능 저랭크

Q2: 왜 diag(alpha)를 그래디언트 스텝 앞에 두는 것이 KDA가 더 강하다는 것을 의미하는가?

모든 좌표가 동일한 alpha를 공유하면, 상태 공간의 각 방향이 같은 속도로 업데이트되어 근사적으로 등방성이 된다. diag(alpha)는 각 채널이 자신만의 유효 학습률을 갖게 한다. 어떤 방향은 새로운 정보를 빠르게 받아들이고, 어떤 방향은 느리게 오래된 정보를 보존하며, 상태 다양체가 이방성적으로 업데이트될 수 있다. 따라서 KDA는 GDN보다 더 세밀한 시간 척도와 기억 분업을 가진다.

Q3: alpha_bias와 dt_bias의 차이

alpha_bias와 dt_bias는 동일한 파라미터가 아니지만, 유사한 "가산적 편향" 형태를 가진다.

Model-Viz에서 이들은 서로 다른 구현과 서로 다른 명명 체계에 속한다:

이름 소재 구현 더해지는 위치 주요 역할 alpha_bias src/models/llm/kimi_k3/modules/kda.py 의 KimiDeltaAttention 저랭크 decay logits에 더해짐: z = alpha_up(alpha_down(x)) + alpha_bias 각 key channel에 대해 학습 가능한 기초 감쇠 logits 설정 dt_bias src/models/llm/linear_attn/variants.py 의 kimi_linear_step 게이트 입력에 더해짐: z = gate_logit + dt_bias 전달되는 key의 편향 제공

alpha_bias의 계산 경로

KDA 모듈에서:

z = self.alpha_up(self.alpha_down(x)) + self.alpha_bias g = g_min * sigmoid(exp(A_log) * z) alpha = exp(g)

여기서:

- alpha_bias의 형태는 (H * Dk,) 이며, 즉 각 key channel마다 하나의 편향이다;

- 그것은 decay-logit z 에 직접 참여한다;

- alpha_bias는 초기화 시 alpha_init과 g_min에 따라 산출되어, 초기 retention factor alpha가 목표값에 가깝도록 한다;

- 그것은 KimiDeltaAttention의 지속적이고 학습 가능한 파라미터이다.

dt_bias의 계산 경로

선형 어텐션 변형에서:

z = gate_logit if dt_bias is not None: z = z + dt_bias g = g_min * sigmoid(exp(A_log) * z) # v2 alpha = exp(g)

- dt_bias는 (H, Dk) 또는 (Dk,) 일 수 있다;

- 그것은 kimi_linear_step()의 선택적 입력이다;

- 코드 자체는 해당 함수 내부에서 이를 생성하거나 초기화하지 않는다;

- 그것은 SGLang/KDA 참조 구현의 gate bias와 정렬하는 데 사용된다.

왜 혼동하기 쉬운가?

동일한 v2 게이팅 공식을 채택할 때, 둘 모두 다음과 같은 곳에 나타날 수 있다:

채널별 logits + 학습 가능하거나 외부에서 전달된 bias → sigmoid → g → exp (g)로 감쇠 인자 alpha를 얻음

따라서 수학적 작용 면에서 둘 모두 "감쇠 게이팅 입력을 바꾸는 가산적 편향"에 속한다. 그러나 이것은 이들이 기능적으로 유사하거나 대응될 수 있음을 설명할 뿐, 동일한 변수임을 설명하지는 않는다.

결론

더 정확히 말하면:

- alpha_bias는 Model-Viz KDA 메인 모듈에서 명시적으로 선언되고 초기화되는 decay-logit 파라미터이다;

- dt_bias는 선형 어텐션 step 함수가 받는 선택적 gate bias로, 주로 외부 참조 구현과 정렬하는 데 쓰인다;

- 둘은 어떤 수식 대조에서 유사한 역할을 담당할 수 있지만, 파라미터 귀속, 호출 인터페이스, 초기화 방식, 코드 위치가 모두 다르다;

- alpha_bias를 dt_bias로 직접 이름을 바꿀 수 없으며, 이를 근거로 두 구현이 완전히 동등하다고 단언할 수도 없다.

Q4: 왜 각 상태 원소마다 독립적인 망각 게이트를 직접 두지 않는가

"각 Head마다 스칼라 하나"에서 "각 Head의 각 Key 채널마다 하나의 값"으로 진화할 수 있다면, 자연히 다음과 같은 것을 계속 상정할 수 있다:

상태 행렬의 모든 원소가 독립적인 망각 계수를 갖고, 원소별 감쇠를 수행하도록 하는 것이다:

이 설계의 표현 능력은 더 강한데, 각 Value 행과 각 Key 열의 조합이 얼마나 보존할지를 독립적으로 결정할 수 있기 때문이다.

그러나 표현 능력만이 유일한 목표는 아니다. 원소별 망각은 뚜렷한 계산 및 하드웨어 비용을 초래한다:

- 각 Head, 각 Token의 게이트 값 개수가 head_dim에서 head_dim_v × head_dim_k 로 증가한다;

- 게이트 생성 비용이 더 높아진다;

- 상태 읽기/쓰기 데이터량이 더 커진다;

- 더 중요한 것은, 현재 업데이트에서 재사용 가능한 대수적 구조를 파괴한다는 점이다.

Key 채널별로 망각할 때는 모든 Value 행이 동일한 우측 곱 대각 행렬을 공유한다. 이렇게 하면 상태 업데이트가 여전히 "대각 행렬 더하기 랭크-원 보정"의 구조를 가져, 효율적인 블록 병렬 알고리즘을 설계할 수 있다.

만약 모든 원소가 독립적으로 망각한다면, 이러한 공유 구조는 대개 존재하지 않게 되어 기존의 Chunkwise 알고리즘을 그대로 적용할 수 없다.

그래서 여기서는 "표현 능력을 더 강하게 만들고 싶지 않다"는 것이 아니라, 표현 능력, 병렬 효율, 게이트 생성 비용, 메모리 읽기/쓰기 사이에서 절충하는 것이다. 대수적 구조가 병렬 상한을 결정하고, 하드웨어 비용이 사용할 수 있는 표현 능력을 제한한다.

HeadwiseRMSNorm은 RMSNorm의 한 변형으로, 각 head의 마지막 차원에 대해서만 정규화하고 head를 가로지르지 않는다. K3는 이를 사용해 KDA 출력 게이트의 첫 단계에서 각 head의 출력을 비교 가능한 스케일로 재조정한 뒤, 전랭크 sigmoid 게이트를 연결한다.‘

왜 head별로 정규화하고 전체 정규화가 아닌가: KDA의 각 head는 자신의 재귀 상태 행렬 S를 유지하며, head 간 출력 스케일은 크게 다를 수 있다(서로 다른 head는 서로 다른 주파수/내용에 주목한다). 만약 전체 정규화를 하면, 대규모 head가 정규화 인자를 지배하여 작은 head의 신호를 눌러 없애버린다. 반면 head별 정규화는 각 head의 출력이 모두 비교 가능한 스케일로 다시 조정되도록 보장하며, 이후 뒤의 전랭크 게이트에 의해 일괄적으로 변조된다.

4.4.1 전신

- LayerNorm: 먼저 평균을 빼고 그다음 분산을 축소한다. 결점: reduction이 한 번 더 있고, bias가 있다.

- RMSNorm: RMS로만 스케일링하고 re-centering을 제거한다. 결점: 전체 벡터에 걸쳐 정규화하므로, 벡터가 여러 head가 연결된 것이라면 큰 head가 정규화 인자를 지배한다.

- HeadwiseRMSNorm(K3): RMSNorm을 각 head 내부로 제한하여 "head 간 정규화 시 큰 head가 지배"하는 문제를 해결한다.

4.4.2 동시기 경쟁 제품

4.4.3 원리 해설

직관적 비유 관점

한 무리의 학생(여러 head)이 질문에 답한다고 상상해 보자. 각자의 음량(스케일) 차이가 매우 크다 —— 어떤 이는 조용히 속삭이고(작은 head), 어떤 이는 귀가 멍멍할 정도로 큰 소리로 말한다(큰 head). 만약 모두의 평균 음량을 함께 집계한 뒤 조정하면, 큰 목소리가 작은 목소리를 덮어버린다.

HeadwiseRMSNorm의 방식은 각자 개별적으로 음량을 표준값으로 맞추는 것(각자 RMS=1)으로, 모든 사람의 "음량"을 비교 가능하게 만든다. 이렇게 하면 이후 하나의 통일된 게이트(sigmoid)가 각자가 얼마나 읽을지를 공정하게 결정할 수 있으며, 큰 목소리에 의해 자연스럽게 눌리지 않는다.

준전문가 수준 관점

- 입력: x ∈ R^{B×T×H×D_h} (H=head 수, D_h=head 차원)

- 출력: y ∈ R^{B×T×H×D_h}

공식

rms (x[b,t,h]) = sqrt ( ( 1 /D_h) * sum_{j} x [b,t,h,j] ^ 2 ) # 각 head 독립 y [b,t,h,j] = x [b,t,h,j] / sqrt ( rms(x[b,t,h])^ 2 + eps ) * weight [j]

- 정규화는 오직 마지막 차원(D_h)에만 작용하며, head에 걸치지 않는다(전체 mean이 아닌 mean(-1)).

표준 RMSNorm과 비교

- 왜 head-wise인가: KDA의 각 head는 독립 상태 행렬 S를 유지하며, head 간 출력 스케일 차이가 클 수 있다(서로 다른 head는 서로 다른 주파수/내용에 주목한다). head 간 정규화는 큰 head가 정규화 인자를 지배하게 하여 작은 head를 눌러 없애버린다. 반면 head별 정규화는 각 head가 독립적으로 비교 가능하게 한다.

- 왜 출력 게이트에 쓰는가: K3의 KDA 출력 게이트 = sigmoid(W_g x) ⊙ RMSNorm_head(o) (Eq.6). 먼저 정규화하여 게이트가 읽는 입력 스케일이 안정적으로 유지되도록 보장해야, 게이트가 공정하게 변조할 수 있다.

- 왜 weight 아핀이 붙는가: 정규화 후 학습 가능한 weight를 곱하여, 각 head가 필요로 하는 스케일 자유도를 회복한다.

- 논문에서 미해결로 남긴 문제: head 간 스케일을 완전히 분리해야 하는가? 각 head마다 독립 weight를 두어야 하는가? 열린 문제다.

4.4.4 구현 & 문제

Q1: HeadRMSNorm과 RMSNorm의 코드 구조는 똑같지 않은가? 단지 dim 차원만 다를 뿐인가?

그렇다, 코드 구조는 거의 같다. 비교하면:

둘 다 동일한 rms_norm을 호출하며, 둘 다 오직 마지막 차원에만 정규화한다. 유일한 실질적 차이는 weight의 형태(dim vs head_dim)이며, weight의 형태는 다시 입력 텐서의 마지막 차원이 무엇인지에 의해 결정된다.

def rms_norm ( x: Tensor, weight: Tensor | None = None , eps: float = 1e-6 ) -> Tensor: """RMS 정규화의 핵심 계산(fp32 누적). 공식: ``y = x / sqrt(mean(x^2) + eps) * weight``, 여기서 ``mean``은 마지막 차원 기준이다. 세 가지 구현 세부 사항: 1. **fp32 승정밀도**: ``x.float()``로 입력을 먼저 fp32로 변환한 뒤 제곱하여 평균을 구한다. bf16/fp16에서는 ``x^2``가 정밀도를 심하게 잃거나 언더플로가 발생하므로, 평균 통계는 반드시 더 높은 정밀도로 계산해야 한다. 그렇지 않으면 정규화 인자 자체에 편향이 생긴다. 2. **먼저 원래 dtype으로 내린 뒤 weight를 곱함**: 가중치 곱은 정밀도를 낮춘 뒤에 수행하므로, 학습과 배포의 수치 동작이 일치한다(양자화 오차 특성이 재현 가능). 3. ``weight``는 채널별 아핀 파라미터이며 None일 수 있다(이때는 순수 정규화와 동등). """ dtype = x.dtype x32 = x. float () # fp32로 올려 통계량 계산 out = x32 * torch.rsqrt(x32. pow ( 2 ).mean(- 1 , keepdim= True ) + eps) # x / rms out = out.to(dtype) # 원래 dtype으로 복귀 if weight is not None : out = out * weight # 채널별 아핀 return out class RMSNorm (nn.Module): """표준 RMSNorm 레이어: 마지막 차원에 대해 단위 RMS 정규화 + 선택적 채널별 아핀. 전형적 사용법(pre-norm): attn_norm = RMSNorm(dim, eps) # model.py의 Kimi3Block h = attn_norm(x) # 정규화 후 어텐션/FFN으로 진입 저랭크 잠재 공간에도 사용된다: MLA의 ``RMSNorm(q_lora_rank)``, MoE의 ``RMSNorm(latent_dim)``. """ def __init__ ( self, dim: int , eps: float = 1e-6 , *, elementwise_affine: bool = True ): super ().__init__() self.dim = dim self.eps = eps # 채널별 학습 가능 스케일, 초기값은 전부 1(정규화 시작 시 순수 단위 스케일이라 초기화를 교란하지 않음). self.weight = nn.Parameter(torch.ones(dim)) if elementwise_affine else None def forward ( self, x: Tensor ) -> Tensor: return rms_norm(x, self.weight, self.eps) def extra_repr ( self ) -> str : return f"dim= {self.dim} , eps= {self.eps} " class HeadRMSNorm (nn.Module): """head-wise RMSNorm: 각 head별로 독립적으로 정규화한다. K3는 이를 KDA 출력 게이트의 첫 단계로 사용한다(kda.py:128 ``self.o_norm``). 입력 형상은 ``(batch, seq, num_heads, head_dim)``이며, 정규화는 마지막 차원 ``head_dim``에만 적용되고 **head를 가로지르지 않는다** — 이것이 바로 "head-wise"의 의미다. 왜 전체 정규화가 아니라 head별 정규화인가: KDA의 각 head는 자신의 순환 상태 행렬 S를 유지하며, head 간 출력 스케일 차이가 클 수 있다(서로 다른 head가 서로 다른 주파수/내용에 주목). 만약 전체를 정규화하면 대스케일 head가 정규화 인자를 지배하여 작은 head의 신호를 눌러버리게 된다. head별 정규화는 각 head의 출력이 비교 가능한 스케일로 재조정된 뒤, 뒤에서 전랭크 게이트에 의해 일괄 변조되도록 보장한다. 주의: 본 구현의 ``weight`` 형상은 ``(head_dim,)``이며 모든 head 간에 공유된다 (``(num_heads, head_dim)`` 가중치 변형은 ``models/llm/core/norms.py``의 ``HeadRMSNorm``을 참조, 인터페이스가 다르며 이 모듈에서는 중복 export하지 않음). """ def __init__ ( self, head_dim: int , eps: float = 1e-6 ): super ().__init__() self.head_dim = head_dim self.eps = eps self.weight = nn.Parameter(torch.ones(head_dim)) def forward ( self, x: Tensor ) -> Tensor: # x 형상 (..., head_dim); rms_norm은 마지막 차원만 정규화하고, weight는 head 축으로 브로드캐스트된다. return rms_norm(x, self.weight, self.eps) def extra_repr ( self ) -> str : return f"head_dim= {self.head_dim} , eps= {self.eps} "

그래서 진짜 문제는 이것이다: 호출 시 전달된 x의 마지막 차원이 무엇인가.

- RMSNorm(dim)은 (b, t, dim) 형상의 은닉 상태에 사용된다 → 마지막 차원은 dim(전체 은닉 차원)이다.

- HeadRMSNorm(head_dim)은 (b, t, num_heads, head_dim_v) 형상의 출력에 사용된다 → 마지막 차원은 head_dim_v이다.

"head-wise"라는 의미는 클래스 자체로 구현되는 것이 아니라, 호출 측이 텐서를 (..., num_heads, head_dim) 형태로 reshape한 뒤 mean(-1)이 자연스럽게 각 head 내부에서만 평균을 구하는 방식으로 구현된다. 즉 HeadRMSNorm은 본질적으로 "RMSNorm이지만 per-head 텐서에专门 적용되는 것"이다. 이름이 전달하는 것은 서로 다른 수학이 아니라 사용 장면이다.

Q2: 참고 코드와의 이점 및 차이

kimi-k3-pytorch-main과의 차이: kimi_primitives/headwise_rmsnorm.py의 HeadwiseRMSNorm 가중치 형상은 (num_heads, head_dim)(per-head affine)이고, 이 프로젝트는 (head_dim,)(head 간 공유)이다. 둘 다 오직 마지막 차원에 대해서만 정규화한다.

SGLang 교차 확인: SGLang은 FusedRMSNormGated(fla/fused_norm_gate.py)를 사용해 정규화 + sigmoid 게이트를 단일 커널로 융합한다(kimi_linear.py:318, o_norm = FusedRMSNormGated(head_dim, activation="sigmoid")). 이 프로젝트의 CPU 버전은 rms_norm + 독립적인 sigmoid(g_proj(x))이며, 의미는 일치한다(KDA 문서 참조).

Q3: out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1)에서 out은 per-head인가? 어디서부터 per-head인가?

per-head이며, 핵심은 252번째 줄에 있다:

out = out .transpose( 1 , 2 ).reshape(b, t, self .num_heads, self .head_dim_v) # ← per-head로 분해 out = self .o_norm( out .to(x.dtype)).reshape(b, t, -1 ) # ← per-head 정규화 후 다시 합침

per-head는 252번째 줄의 reshape(b, t, num_heads, head_dim_v)부터 시작된다. 그러나 그 이전에도 out은 전체 순환 과정에서 줄곧 per-head였으며, 단지 배열만 달랐다:

전체 흐름:

순환 출력 out: (b, h, t, dv) ← kda_chunkwise/kda_step에서 반환(per-head, head는 1번째 차원) ↓ transpose ( 1 , 2 ) → (b, t, h, dv) ↓ reshape (b, t, h, dv) → (b, t, h, dv) ← 여전히 per-head, 다만 (t, h) 배열로 변경 out. to (x.dtype) → o_norm → 마지막 차원 dv에 대해 정규화(각 head 독립) ↓ reshape (b, t, - 1 ) → (b, t, h*dv) ← head를 다시 합침, 더 이상 per-head 아님 out * sigmoid ( g_proj (x)) → 전랭크 게이트 o_proj → (b, t, dim)

더 이른 출처: per-head 형상은 사실 투영 단계에서부터 시작된다.

- kda.py:180-182: q/k/v는 먼저 넓은 벡터에서 view(b, t, num_heads, head_dim_k)로 per-head로 분해된 뒤, l2_normalize / silu가 진행된다.

출처 원문