FG$^2$-GDN: Enhancing Long-Context Gated Delta Networks with Doubly Fine-Grained Control

1. 문제 정의 (Problem)

최근 Transformer 아키텍처의 대안으로 등장한 선형 어텐션 (Linear Attention) 메커니즘은 추론 시 $O(L^2)$의 복잡도를 $O(L)$로 줄여 긴 컨텍스트 처리를 가능하게 합니다. 특히 **Gated DeltaNet (GDN)**과 **Kimi Delta Attention (KDA)**은 '델타 규칙 (Delta Rule, 온라인 경사 하강법 기반)'을 도입하여 단순한 가법 업데이트보다 우수한 연관 기억 (Associative Recall) 성능을 보여주었습니다.

그러나 기존 모델들의 한계는 다음과 같습니다:

• 학습률의 스칼라 제약: KDA 는 '기억 소멸 (forgetting)'을 채널 단위 (channel-wise) 로 세밀하게 제어할 수 있게 했지만, '기억 쓰기 (writing/update)'를 담당하는 학습률 $\beta_t$는 여전히 **스칼라 (scalar)**로 고정되어 모든 차원이 동일한 강도로 업데이트됩니다.
• 차원별 적응 부재: 이는 모델이 특정 컨텍스트에서 각 특징 차원 (feature dimension) 의 중요도에 따라 적응적으로 학습 강도를 조절하는 능력을 제한합니다. 즉, 모든 차원이 동일한 '쓰기 (write)' 강도를 공유하게 되어 정밀한 정보 저장이 어렵습니다.

2. 제안 방법론 (Methodology)

저자들은 이러한 한계를 극복하기 위해 **FG2-GDN (Fine-Grained Gated Delta Network)**을 제안했습니다. 핵심 아이디어는 델타 업데이트의 학습률을 스칼라에서 채널 단위 벡터로 확장하는 것입니다.

핵심 기술 요소:

1. 채널 단위 적응 학습률 (Channel-wise Adaptive Learning Rates):

   • 기존 KDA 의 스칼라 $\beta_t$를 벡터 $\beta_t \in \mathbb{R}^{d_k}$로 대체합니다. 이는 SGD 에서 AdaGrad나 Adam 과 같은 좌표별 (per-coordinate) 적응적 최적화 기법을 델타 규칙에 적용한 것과 유사합니다.
   • 각 특징 차원이 독립적인 적응 역학을 가지도록 하여, 중요한 차원은 강하게 학습하고 덜 중요한 차원은 약하게 학습할 수 있게 합니다.

2. 대칭적 스케일링 및 구조 유지:

   • 단순히 행렬 곱을 변경하면 GDN/KDA 의 효율적인 병렬 학습을 가능하게 하는 Rank-1 Householder 변환 구조가 깨질 수 있습니다.
   • 이를 해결하기 위해 $\beta_t$를 키 ($k_t$) 와 값 ($v_t$) 벡터 모두에 대칭적으로 적용합니다 ($\tilde{k}_t = \sqrt{\beta_t} \odot k_t$, $\tilde{v}_t = \sqrt{\beta_t} \odot v_t$).
   • 이를 통해 업데이트 규칙은 여전히 Diagonal-Plus-Low-Rank (DPLR) 구조를 유지하며, 기존 KDA 의 Chunkwise Parallel Training 알고리즘을 그대로 재사용할 수 있습니다.

3. FG2-GDN+ (Decoupled Variant):

   • 더 나아가 키 (Key) 의 '지우기 (erasure)' 강도와 값 (Value) 의 '쓰기 (write)' 강도를 분리 제어하기 위해 $\beta^k_t$와 $\beta^v_t$를 별개의 벡터로 정의한 변형 모델을 제안했습니다.
   • 이는 기존 정보를 얼마나 강하게 지울지와 새로운 정보를 얼마나 강하게 쓸지를 독립적으로 조절할 수 있게 합니다.

4. 하이브리드 아키텍처:

   • 순수 선형 모델의 한계를 보완하기 위해, MLA (Multi-head Latent Attention) 레이어와 FG2-GDN 레이어를 3:1 비율로 교차 배치하는 하이브리드 구조를 사용합니다.

3. 주요 기여 (Key Contributions)

• 이중 미세 조정 (Doubly Fine-Grained Control): KDA 가 '기억 소멸'을 채널 단위로 세분화했다면, FG2-GDN 은 '기억 쓰기 (학습률)'까지 채널 단위로 세분화하여 이중 미세 조정을 달성했습니다.
• 효율성 유지: 새로운 벡터 학습률 도입에도 불구하고, DPLR 구조를 유지하여 기존 KDA 와 동일한 병렬 학습 알고리즘을 사용할 수 있으며, 프리필 (prefill) 처리량 오버헤드는 5% 미만으로 매우 낮습니다.
• 성능 향상: 합성 및 실세계 벤치마크에서 GDN 과 KDA 대비 연관 기억 (Associative Recall) 및 긴 컨텍스트 이해 능력이 크게 향상됨을 입증했습니다.

4. 실험 결과 (Results)

340M 및 1.3B 파라미터 규모에서 실험을 수행한 결과는 다음과 같습니다.

• 언어 모델링 (Language Modeling):

  • FG2-GDN 은 1.3B 규모에서 LAMBADA 퍼플렉시티 (Perplexity) 를 13.80 (KDA) 에서 13.09로 낮추었으며, 7 가지 상식 추론 벤치마크에서 평균 정확도도 가장 높았습니다.
  • 이는 미세 조정된 쓰기 메커니즘이 짧은 컨텍스트 언어 이해에 특히 유리함을 시사합니다.

• 긴 컨텍스트 평가 (Long-Context Evaluation):

  • RULER 벤치마크: 단일 키 검색 (Single-key retrieval) 및 다중 키 검색 (Multi-key retrieval) 작업에서 GDN 과 KDA 를 크게 상회하는 성능을 보였습니다. 특히 FG2-GDN+ 는 16K 시퀀스 길이에서도 가장 높은 정확도를 유지하며 외삽 (Extrapolation) 능력이 뛰어났습니다.
  • LongBench: 문서 QA, 요약, 퓨샷 (Few-shot) 학습, 코드 완성 등 다양한 긴 컨텍스트 작업에서 평균 점수가 가장 높았습니다 (18.3%).
  • 해석: 차원별 학습률 부여는 키 - 값 바인딩 (Key-Value Binding) 을 더 단단하게 만들어, 쿼리 시의 변형에도 강인한 검색을 가능하게 합니다. 또한, 키와 값의 강도를 분리 (FG2-GDN+) 하면 복잡한 다중 연관 관계를 관리하는 데 도움이 됩니다.

• 효율성:

  • 시퀀스 길이가 길어질수록 MLA 대비 1.7~1.9 배의 속도 향상을 보였으며, KDA 대비 처리량 오버헤드는 미미했습니다.

5. 의의 및 결론 (Significance)

이 논문은 선형 재귀 모델 (Linear Recurrent Models) 의 **적응성 (Adaptivity)**을 한 단계 발전시켰습니다.

• 이론적 의의: 델타 규칙 (온라인 학습) 에 좌표별 적응적 학습률 개념을 성공적으로 통합하여, 단순한 가중치 업데이트를 넘어 차원별 특성에 맞는 동적 메모리 관리를 가능하게 했습니다.
• 실용적 의의: 계산 비용 증가 없이 (Near-zero cost) 긴 컨텍스트 처리 능력을 획기적으로 개선하여, 에이전트 (Agentic) 및 멀티모달 시나리오와 같은 긴 컨텍스트가 필요한 응용 분야에서 Transformer 의 대안으로서의 가능성을 높였습니다.

요약하자면, FG2-GDN은 "기억을 지우는 것"뿐만 아니라 "기억을 쓰는 것"까지도 세밀하게 제어함으로써, 긴 문맥에서의 정보 저장 및 검색 효율성을 극대화한 차세대 선형 어텐션 모델입니다.