기술 요약: Gated DeltaNet-2

문제 정의

선형 어텐션 메커니즘은 무제한 캐시를 고정 크기의 재귀 상태로 대체함으로써 표준 소프트맥스 어텐션에 대한 매력적인 대안을 제시하며, 이를 통해 시퀀스 혼합을 선형 시간으로, 디코딩을 상수 메모리로 줄입니다. 그러나 델타 규칙 모델(DeltaNet, Gated DeltaNet, Kimi Delta Attention/KDA 등) 에서는 근본적인 병목 현상이 남아 있습니다. 이러한 모델들은 타겟팅된 덮어쓰기를 통해 기억 관리 능력을 향상시키지만, 두 가지 서로 다른 연산을 동시에 제어하기 위해 스칼라 게이트($\beta_t$) 에 의존합니다.

1. 삭제 (Erasing): 키 측에서 읽은 이전 콘텐츠의 얼마나 많은 부분을 제거할지 결정합니다.
2. 기록 (Writing): 값 측에서 들어오는 새로운 콘텐츠의 얼마나 많은 부분을 커밋할지 결정합니다.

본 논문은 이러한 두 가지를 단일 스칼라 결정으로 취급하는 것이 모델링의 제약이라고 주장합니다. 삭제는 키 축 (기존 읽기에서 어떤 좌표를 제거할지 결정) 에 작용하는 반면, 기록은 값 축 (들어오는 값의 어떤 좌표를 커밋할지 결정) 에 작용합니다. 이러한 결정을 결합하면 기존 연관성을 뒤섞지 않고 압축된 기억을 선택적으로 편집하는 모델의 능력이 제한되며, 특히 많은 연관성이 동일한 유한 상태 공간을 위해 경쟁하는 장문맥 시나리오에서 더욱 그렇습니다.

방법론: Gated DeltaNet-2

저자들은 델타 규칙 내에서 삭제 및 기록 연산을 분리한 재귀 어텐션 레이어인 Gated DeltaNet-2를 제안합니다.

1. Gated Delta Rule-2

핵심 혁신은 스칼라 게이트를 두 개의 독립적인 채널별 게이트로 분리하는 것입니다.

• 삭제 게이트 ($b_t \in [0, 1]^{d_k}$): 키 차원에 적용되는 채널별 벡터입니다. 이는 현재 읽기 ($e_t = b_t \odot k_t$) 의 좌표에 가중치를 부여하여 이전 상태의 어떤 부분이 삭제될지 결정합니다.
• 기록 게이트 ($w_t \in [0, 1]^{d_v}$): 값 차원에 적용되는 채널별 벡터입니다. 이는 들어오는 값 ($z_t = w_t \odot v_t$) 의 좌표에 가중치를 부여하여 새로운 값의 어떤 부분이 커밋될지 결정합니다.

상태 업데이트 방정식은 다음과 같습니다:
$$S_t = (I - k_t(b_t \odot k_t)^\top) D_t S_{t-1} + k_t(w_t \odot v_t)^\top$$
여기서 $D_t = \text{Diag}(\alpha_t)$는 KDA 에서 계승된 채널별 감쇠를 나타냅니다.

이 공식은 이전 모델들을 일반화합니다.

• $b_t$와 $w_t$가 동일한 스칼라 $\beta_t$에 묶인 경우, KDA를 복원합니다.
• 감쇠 $\alpha_t$도 스칼라에 묶인 경우, Gated DeltaNet을 복원합니다.

2. Fast-Weight Update 관점

저자들은 업데이트를 지역 회귀 손실 (local regression loss) 에 대한 온라인 최적화 단계로 해석합니다. 모델은 읽기 값 ($\bar{S}_{t-1}^\top e_t$) 을 게이트된 목표 ($z_t$) 로 수정하는 연관성 편집을 조건으로 하여, 새로운 상태와 감쇠된 기억 사이의 거리를 최소화합니다. 이 관점은 업데이트가 게이트된 이전 읽기를 뺀 후 게이트된 새로운 값을 추가함을 명확히 합니다.

3. 효율적인 청크별 병렬 학습

학습 효율성을 유지하기 위해 저자들은 누적 채널별 감쇠를 랭크-1 삭제 인자로 흡수하는 WY 알고리즘(KDA 와 유사) 을 유도했습니다.

• 감쇠 흡수: 채널별 감쇠는 삭제 ($\bar{E}$) 및 기록 ($Z$) 행렬로 정규화되어, 재귀를 순수 비대칭 델타 업데이트로 표현할 수 있게 합니다.
• 삼각형 풀이: 청크별 계산은 보조 행렬 $Y$(삭제 측) 와 $U$(기록 측) 를 계산하기 위해 하삼각 시스템 $(I+T)^{-1}$을 푸는 과정을 포함합니다.
• 게이트 인식 역전파: 중요한 기술적 기여는 역전파의 유도입니다. 스칼라 게이트 모델에서는 게이트 인자를 내적에서 분리할 수 있지만, Gated DeltaNet-2 는 채널별 게이트 ($b_t$ 및 $w_t$) 를 보조 행렬에 대한 기울기 누적에 직접 포함시켜야 합니다. 이는 효율적인 청크별 구조를 깨뜨리지 않으면서 올바른 기울기 흐름을 보장합니다.

4. 아키텍처

이 모델은 표준 Transformer 스타일 블록의 토크 믹서로 구현됩니다. 재귀 전용 및 하이브리드 구성 모두를 지원합니다. 하이브리드 변형은 정확한 지역 상호작용을 처리하기 위해 슬라이딩 윈도우 어텐션 (SWA) 과 Gated DeltaNet-2 믹서를 번갈아 배치하며, 재귀 레이어는 장거리 의존성을 관리합니다.

주요 결과

이 모델은 100B FineWeb-Edu 토큰으로 학습된 1.3B 파라미터로 평가되었으며, Mamba-2, Gated DeltaNet, KDA, Mamba-3(SISO/MIMO), 그리고 Transformer 들과 비교되었습니다.

• 언어 모델링 및 추론: Gated DeltaNet-2 는 퍼플렉시티 (WikiText, LAMBADA) 및 제로샷 상식 추론 벤치마크 (PIQA, HellaSwag 등) 에서 재귀 및 하이브리드 설정 모두에서 가장 강력한 전반적인 결과를 달성했습니다.
• 장문맥 검색 (RULER): 가장 큰 개선은 다중 키 헤이스트랙 (MK-NIAH) 작업에서 관찰되었습니다. 이러한 시나리오에서는 고정 크기 상태가 경쟁하는 연관성을 분리해야 하므로, Gated DeltaNet-2 는 모든 베이스라인을 능가했으며 특히 문맥 길이가 증가함에 따라 (최대 8K 토큰) 더 두드러진 성능을 보였습니다.
• 실제 세계 검색: 추출 및 질문 답변 (SQuAD, TriviaQA, DROP) 과 관련된 작업에서 하이브리드 Gated DeltaNet-2 모델은 가장 높은 평균 정확도를 달성하여 노이즈가 많은 실제 세계 컨텍스트에서의 견고함을 입증했습니다.
• 효율성: H100 GPU 에서의 처리량 분석에 따르면, Gated DeltaNet-2 는 추가적인 채널별 게이트로 인해 KDA 대비 소규모의 상수 오버헤드만 발생시키면서 시퀀스 길이가 증가함에 따라 거의 평탄한 확장성을 유지합니다.

중요성 및 주장

본 논문은 활성 기억 편집의 분리가 고정 상태 재귀의 주요 압력 지점인 많은 압축된 연관성 간의 간섭을 직접적으로 해결한다고 주장합니다.

• 모델링 유연성: 삭제와 기록을 분리함으로써 모델은 감쇠를 통해 광범위한 컨텍스트를 지우고, 삭제 게이트를 통해 특정 낡은 연관성을 제거하며, 기록 게이트를 통해 필요한 값 채널만 삽입할 수 있습니다. 이러한 유연성은 모델이 KDA 와 Gated DeltaNet 을 특수한 경우로 복원하면서도 필요시 더 복잡한 채널 구조를 학습할 수 있게 합니다.
• 실증적 검증: 제거 연구 (Ablation studies) 는 두 게이트 모두 성능에 기여하며, 특히 채널별 삭제 게이트 ($b_t$) 가 언어 모델링 및 검색에서 대부분의 개선을 차지함을 확인했습니다.
• 실용성: 이 연구는 고급 기억 제어 메커니즘이 대규모 학습에 필요한 병렬 학습 효율성을 희생하지 않고 선형 어텐션에 통합될 수 있음을 보여줍니다.

저자들은 Gated DeltaNet-2 가 "재귀 어텐션 최전선"에서의 한 걸음 전진을 나타내며, 장문맥 시나리오에서 압축된 기억을 관리하기 위한 더 표현력 있고 효율적인 메커니즘을 제공한다고 결론지었습니다.