기술 요약: OSDN: 선형 어텐션에서 증명 가능한 온라인 전조건화를 통한 델타 규칙 개선

1. 문제 제기

선형 어텐션과 상태 공간 모델 (SSM) 은 소프트맥스 어텐션에 대한 $O(N)$ 추론 복잡도와 상수 메모리 대안을 제공하여 긴 시퀀스 처리에 적합합니다. 그러나 소프트맥스 기반 트랜스포머에 비해 **맥락 내 연관 회상 (in-context associative recall)**에서 종종 어려움을 겪습니다.

델타 규칙 (Delta Rule)(DeltaNet 에서 구현된 바와 같이) 은 토큰별 회귀 손실 $f_t(S) = \frac{1}{2}\|S k_t - v_t\|_F^2$에 대한 온라인 경사 하강 (OGD) 의 단일 단계로 재귀 상태 업데이트를 해석함으로써 이 격차를 해소하려 합니다. 업데이트 규칙은 $S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)k_t^\top$이며, 여기서 $\beta_t$는 스칼라 게이트입니다.

핵심 한계: 스칼라 게이트 $\beta_t$는 모든 키 차원에 적용되는 균일한 학습률로 작용합니다. 이는 내부 목적 함수의 **특성별 곡률 (feature-wise curvature)**을 무시합니다. 연관 회상 작업에서 서로 다른 키 (예: 빈번한 대 희귀한 식별자, 안정적인 대 고엔트로피 값) 는 서로 다른 경험적 곡률을 보입니다. 단일 스칼라 단계 크기는 타협을 강요합니다: 한 방향의 구식 메모리를 덮어쓰기 위해 이를 증가시키면 이미 잘 보정된 다른 방향을 과도하게 수정할 수 있습니다. 표준 적응형 옵티마이저 (Adam 등) 는 이를 해결하기 위해 대각 전조건화를 사용하지만, 이러한 메커니즘을 재귀 레이어에 적용하는 것은 일반적으로 밀집된 2 차 상태가 필요하거나 DeltaNet 을 효율적으로 만드는 하드웨어 친화적 청크별 병렬화 파이프라인 (WY 변환 등) 을 깨뜨립니다.

2. 방법론: 온라인 스케일드 델타넷 (OSDN)

저자들은 OSDN을 제안하며, 이는 하이그래디언트 피드백을 통해 온라인으로 학습되는 대각 전조건부 (diagonal preconditioner) $D_t = \text{diag}(d_t)$로 스칼라 게이트를 보강합니다.

2.1. 대수적 동치성 및 하드웨어 보존

OSDN 의 핵심 통찰은 그래디언트 스텝을 우측 전조건화하는 것이 쓰기 측 키의 특성별 스케일링과 대수적으로 동치라는 것입니다.
표준 업데이트:
$$S_t = S_{t-1} - \beta_t \nabla f_t(S_{t-1}) D_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)(d_t \odot k_t)^\top$$
전조건부 키 $\tilde{k}_t = d_t \odot k_t$를 정의하면 업데이트는 다음과 같이 됩니다:
$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)\tilde{k}_t^\top$$
이 변환을 통해 OSDN 은 DeltaNet 의 청크별 WY 병렬 파이프라인을 엄격하게 보존할 수 있습니다. 구현은 청크 그람 행렬 및 청크 내 점수에서 저장 키 $K$를 $\tilde{K}$로 대체하기만 하면 되며, $O(K)$ 전조건부 벡터 이상의 고차원 상태 오버헤드가 발생하지 않습니다.

2.2. 분리된 하이그래디언트 피드백

전조건부 $d_t$는 온라인 스케일드 그래디언트 방법 (OSGM) 프레임워크 [18] 에서 유도된 **하이그래디언트 대리 손실 (hypergradient surrogate)**에 대한 온라인 경사 하강 (OGD) 을 통해 업데이트됩니다.
전조건부에 대한 대리 손실은 다음과 같습니다:
$$h_t(d_t) = \frac{f_t(S_t) - f_t(S_{t-1})}{\|\nabla f_t(S_{t-1})\|_F^2}$$
중요하게도, 내부 회귀 손실 $f_t$의 **정확한 2 차 구조 (exact-quadratic structure)**로 인해, 이 대리는 키 $k_t$와 스칼라 게이트 $\beta_t$에만 의존하는 **폐쇄형 표현 (closed-form expression)**을 허용하며, 고차원 상태 $S_t$, 값 $v_t$, 또는 잔차 $u_t$로부터 분리됩니다.
$$h_t(d_t) = \frac{(1 - \beta_t \langle d_t, k_t^2 \rangle)^2 - 1}{2 \|k_t\|_2^2}$$
이 분리는 2 단계 구현을 가능하게 합니다:

1. 1 단계: 전조건부 키 시퀀스 $\{\tilde{k}_t\}$를 계산하고 $d_t$를 업데이트하기 위해 키 스트림을 가볍게 $O(K)$ 스캔합니다.
2. 2 단계: $K$ 대신 $\tilde{K}$를 사용하는 표준 청크별 DeltaNet 패스입니다.

2.3. 적응형 전조건부 망각 (APF)

키 분포가 이동하는 비정상적인 언어 컨텍스트를 처리하기 위해 저자들은 APF를 도입합니다. 이는 메모리 상태 $S_t$가 아닌 전조건부 상태 $d_t$에만 적용되는 학습된 토큰별, 헤드별 유지 게이트 $r_{t,h}$를 추가합니다:
$$\bar{d}_{t+1} = r_{t,h} d_t + \eta \nabla_d h_t(d_t)$$
이는 2 단계 스캔에 필요한 아핀 재귀 구조를 유지하면서 메타 옵티마이저가 구식 보정을 "망각"할 수 있게 합니다.

3. 이론적 보장

이 논문은 내부 손실의 정확한 2 차 특성에 기반하여 두 가지 수렴 보장을 확립합니다:

1. 집합 한계 초기하 수렴 (Theorem 5.1): 이상적인 우측 뉴턴 비교기 ($D^\star = \Sigma_k^\dagger$) 에 대해, 온라인 학습자가 아선형 후회 (sublinear regret) 를 달성한다면, 글로벌 목적 함수의 비최적성은 초기하 속도인 $O((C/T)^T)$로 수렴합니다. 이는 표준 기하 (선형) 수렴보다 훨씬 빠릅니다.
2. 알고리즘 로컬 잔차 수축 (Theorem 5.2): 구현된 대각 업데이트의 경우, 토큰 로컬 잔차 비율의 곱 $\prod_{t=1}^T \frac{f_t(S_t)}{f_t(S_{t-1})}$是有界입니다. 대각 비교기가 게이트된 뉴턴 조건 ($\beta_t \langle d^\star, k_t^2 \rangle = 1$) 을 만족하면, 수축 속도도 표준 OGD 후회 한계 하에서 특히 $O((1/\sqrt{T})^T)$인 초기하적입니다.

4. 실험 결과

실험은 340M 및 1.3B 파라미터 규모에서 DeltaNet, Gated DeltaNet (GDN), Kimi Delta Attention (KDA) 백본을 대상으로 수행되었습니다.

4.1. 맥락 내 회상 (340M 규모)

• 성능: 바닐라 OSDN 은 기준선 DeltaNet 대비 JRT 스타일 맥락 내 회상을 32% 개선했습니다 (0.150 $\to$ 0.198). 향상은 키가 반복되는 반복 컨텍스트 작업 (예: "JRT-twice" 변형) 에 집중되었으며, 이를 통해 전조건부가 보정될 수 있었습니다.
• APF 영향: OSDN-APF 는 **17%**의 향상 (0.176) 을 유지했으며, DeltaNet 변형 중 WikiText/LAMBADA 퍼플렉시티가 가장 낮았습니다 (30.99 대 32.00).
• 일반화: 더 넓은 벤치마크 (Commonsense, LongBench) 에서 OSDN 변형은 기준선과 동등한 수준을 유지하여, 이 메커니즘이 보편적인 성능 부스터가 아니라 표적화된 것임을 확인했습니다.

4.2. 1.3B 파라미터로 확장

• 잔차 수축: 메커니즘 수준의 신호가 규모에서 증폭되었습니다. 기하 평균 잔차 비율 ($q_{geo}$) 이 0.432 에서 0.265 로 하락 (39% 감소) 하여 340M 에서 본 개선의 거의 두 배에 달했습니다.
• 하류 동등성: 내부 잔차 수축이 대폭 개선되었음에도 불구하고, 하류 지표 (WikiText/LAMBADA 퍼플렉시티, Commonsense, LongBench) 는 1.3B DeltaNet 기준선과 동등한 수준을 유지하여, 이 메커니즘이 일반 능력을 저하시키지 않고 효과적으로 이전됨을 입증했습니다.

4.3. 효율성

• 처리량: OSDN 변형은 최소한의 추론 오버헤드를 발생시켰습니다. 340M 에서 처리량 변화는 $\pm 2.2\%$ 이내였습니다. 1.3B 에서 OSDN-APF 는 기준선보다 6.8% 느렸으며, 이는 전조건부를 위한 추가 순차 스캔 때문입니다.
• 상태 오버헤드: $O(K)$ 전조건부 벡터로 인해 영구 재귀 상태는 $\le 0.05\%$만 증가했습니다.

5. 중요성 및 주장

이 논문은 OSDN 을 1 차 스칼라 업데이트와 적응형 최적화 사이의 격차를 해소하는 최소적이고 표적화된 델타 규칙 확장으로 위치시킵니다.

• 메커니즘 수준 개선: 주요 기여는 잔차 수축 비율의 직접적인 진단으로 검증된 온라인 전조건화를 통한 연관 회상의 증명 가능한 메커니즘 수준 개선입니다.
• 하드웨어 호환성: 다른 2 차 또는 밀집 전조건화 방법과 달리 OSDN 은 청크별 병렬 파이프라인과 $O(K)$ 상태 복잡도를 보존하여 대규모 배포에 실용적입니다.
• 확장성: 결과는 온라인 전조건화의 이점이 십억 파라미터 규모에서 이전되고 증폭됨을 보여주며, 선형 어텐션의 "곡률 불일치"가 아키텍처 개편 없이 해결될 수 있는 근본적인 병목 현상임을 시사합니다.
• 한계: 저자들은 OSDN 이 보편적인 벤치마크 개선 도구가 아님을 명시적으로 지적합니다. 향상은 반복되는 키가 포함된 검색 작업에 집중됩니다. 이론적 보장은 내부 손실의 정확한 2 차 가정과 단조 하강에 의존하며, 이는 이상화된 조건입니다. 이 방법은 모든 맥락 내 학습 한계를 해결한다고 주장하지 않으며, 구체적으로 델타 규칙의 쓰기 단계 기하학을 다룹니다.