From Gradients to Riccati Geometry: Kalman World Models for Single-Pass Learning

1. 문제 제기 (Problem)

현대 머신러닝은 역전파 (Backpropagation) 와 확률적 경사 하강법 (SGD) 및 그 변형에 기반한 역모드 자동미분 (Reverse-mode Automatic Differentiation) 에 의해 주도되어 왔습니다. 그러나 이 방식은 대규모 모델과 온라인 학습 환경에서 다음과 같은 구조적, 계산적 한계를 노출시킵니다.

• 메모리 오버헤드: 역전파는 역방향 전파를 위해 순방향 통과 중의 모든 중간 활성화 (Intermediate Activations) 를 저장해야 하므로, 모델의 깊이와 시퀀스 길이에 비례하는 막대한 메모리가 필요합니다.
• 배치 의존성: 역전파는 일반적으로 배치 (Batch) 기반의 최적화를 전제로 하며, 데이터가 순차적으로 유입되는 온라인 학습이나 분산 환경에서 동기화 비용과 불안정성을 초래합니다.
• 불확실성 부재: 1 차 최적화 방법은 매개변수 공간의 불확실성을 명시적으로 표현하지 못하며, 이는 continual learning(연속 학습) 에서 catastrophic forgetting(치명적 망각) 을 유발합니다.

이 논문은 역전파를 대체할 수 있는 통제 이론 (Control Theory) 과 베이지안 추정 에 기반한 새로운 학습 패러다임을 제안합니다.

2. 방법론 (Methodology)

저자는 신경망을 정적 함수 근사기가 아닌 확률적 동적 시스템 (Stochastic Dynamical System) 으로 재해석하고, 칼만 필터 (Kalman Filter) 를 학습 알고리즘의 핵심으로 도입합니다.

2.1 칼만 월드 모델 (Kalman World Models, KWM)

• 상태 공간 표현: 신경망의 매개변수 $\theta$와 은닉 상태 $x$를 동적 상태 벡터 $z_t = [x_t, \theta_t]^T$로 정의합니다.
• 학습의 재정의: 학습을 정적 손실 함수의 최소화 (Deterministic Optimization) 가 아닌, 관측 데이터 $y_t$에 기반한 순차적 베이지안 상태 추정 (Sequential Bayesian State Estimation) 으로 변환합니다.
• 업데이트 규칙: 경사 하강법 대신 칼만 이득 (Kalman Gain) 을 이용한 재귀적 보정을 사용합니다.
  $$\theta_{t+1} = \theta_t + K_t (y_t - \hat{y}_t)$$
  여기서 $(y_t - \hat{y}_t)$는 혁신 (Innovation) 이라 불리는 오차 신호이며, $K_t$는 공분산 전파를 통해 동적으로 계산된 가중치입니다.

2.2 칼만 이득과 자연 경사 (Natural Gradient) 의 동치성

• 핵심 이론: 국소 선형 가우시안 관측 모델 하에서, 칼만 필터의 매개변수 업데이트는 피셔 정보 행렬 (Fisher Information Matrix) 로 사전 조건화 (Preconditioning) 된 자연 경사 하강법 (Natural Gradient Descent) 과 수학적으로 동치임을 증명합니다.
• 리카티 기하학: 매개변수 공분산 행렬 $P_t$는 피셔 정보 행렬의 역행렬 ($F^{-1}$) 역할을 하며, 리카티 방정식 (Riccati Equation) 을 통해 시간에 따라 진화합니다. 이는 학습을 유클리드 공간이 아닌 정보 기하학 (Information Geometry) 의 리만 다양체 상에서의 강하로 해석하게 합니다.

2.3 활성화 수준의 혁신 보정 (Activation-Level Innovation Correction)

• 변환기 (Transformer) 적용: LLM 의 내부 활성화 (Hidden States) 를 고정된 값이 아닌 잠재 상태 (Latent State) 로 간주합니다.
• 온라인 보정: 토큰이 생성될 때, 예측된 토큰과 실제 관측된 토큰 사이의 오차 (혁신) 를 이용해 활성화 상태를 실시간으로 보정합니다.
  $$h_t \leftarrow h_t + K_t (y_t - \hat{y}_t)$$
• 이 방식은 전역 가중치를 수정하지 않으면서도 분포 변화 (Distribution Shift) 에 대한 적응력과 환각 (Hallucination) 억제를 가능하게 합니다.

2.4 쿠퍼만 (Koopman) 리프트 (Lifted) 월드 모델

• 비선형성 해결: 기존 확장 칼만 필터 (EKF) 의 반복적 국소 선형화 (Jacobian 계산) 한계를 극복하기 위해 쿠퍼만 연산자 (Koopman Operator) 이론을 도입합니다.
• 선형 공간으로의 매핑: 비선형 동적 시스템을 고차원 특징 공간으로 "리프트 (Lift)"하여, 해당 공간에서는 시스템의 진화가 정확히 선형이 되도록 합니다.
• 이를 통해 칼만 필터링을 근사 없이 정확하게 수행할 수 있으며, 학습의 핵심을 가중치 최적화가 아닌 선형 연산자 추정으로 전환합니다.

3. 주요 기여 (Key Contributions)

1. 칼만 - 자연 경사 동치성 증명: 칼만 필터 업데이트가 피셔 정보 기반의 자연 경사 하강법과 동치임을 수학적으로 증명하여, 제어 이론과 정보 기하학을 통합했습니다.
2. 학습의 필터링 재해석: 학습을 역전파가 아닌 순차적 베이지안 추정으로 재정의하여, 온라인 학습과 불확실성 정량화를 자연스럽게 지원합니다.
3. 활성화 수준 보정 메커니즘: Transformer 아키텍처의 내부 상태를 칼만 필터의 관측 대상으로 삼아, 역전파 없이도 실시간으로 표현을 보정하는 새로운 방식을 제시했습니다.
4. 확장성 및 안정성 분석:
   • 구조화된 공분산: 블록 대각 (Block-diagonal), 저랭크 (Low-rank), 크로네커 인수분해 (Kronecker-factored) 등을 통해 $O(p^2)$의 복잡도를 $O(pr)$로 줄여 대규모 모델 적용 가능성을 제시했습니다.
   • 수렴성 보장: 강한 볼록성 (Strong Convexity) 하의 전역 수렴, 시간 변화 야코비안 하의 지수적 안정성 (UES), 저랭크 근사 하의 강건한 안정성 등을 이론적으로 증명했습니다.

4. 실험 결과 (Results)

논문은 Sequential MNIST, WikiText-2 언어 모델링, 온라인 적응 벤치마크에서 KWM 을 평가했습니다.

• 성능: 정적 분포 (Stationary Distribution) 조건에서 Adam, K-FAC 등 기존 최적화 기법과 유사한 정확도 및 퍼플렉시티 (Perplexity) 를 달성했습니다.
• 강건성 (Robustness):
  • 분포 변화 (Distribution Shift): 토큰 드롭아웃 및 어휘 변화와 같은 시나리오에서 KWM (특히 활성화 보정 적용 시) 이 기존 방법보다 훨씬 안정적인 퍼플렉시티를 유지했습니다.
  • 연속 학습 (Continual Learning): Permuted MNIST 및 스트리밍 언어 모델링에서 치명적 망각 (Catastrophic Forgetting) 을 약 30% 감소시켰습니다. 이는 공분산 행렬이 매개변수의 불확실성을 표현하여, 확신이 높은 매개변수는 크게 업데이트되지 않도록 자동 조절하기 때문입니다.
• 수렴 특성: 학습률 초기화에 덜 민감하며, 공분산 진화가 리카티 해에 수렴하는 모습을 보였습니다.

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 역전파의 독점적 지배에 도전하며, 신경망을 리카티 기하학 (Riccati Geometry) 에 기반한 확률적 동적 추정기 (Stochastic Dynamical Estimator) 로 재정의합니다.

• 패러다임 전환: 학습을 "손실 최소화"가 아닌 "불확실성 하의 상태 추정"으로 바라봄으로써, 온라인 학습, 불확실성 정량화, 그리고 분포 변화에 대한 적응력을 내재화합니다.
• 이론적 통합: 제어 이론 (칼만 필터), 정보 기하학 (자연 경사), 그리고 현대 딥러닝 아키텍처 (Transformer, Koopman Operator) 를 하나의 수학적 프레임워크로 통합했습니다.
• 미래 전망: 대규모 언어 모델 (LLM) 에 대한 역전파의 메모리 병목 현상을 해결할 수 있는 잠재력을 제공하며, 단일 패스 (Single-Pass) 학습과 실시간 적응형 AI 시스템의 새로운 길을 제시합니다.

요약하자면, 이 연구는 기울기 (Gradient) 대신 혁신 (Innovation) 과 리카티 기하학을 통해 학습을 재구성함으로써, 더 강건하고 적응적이며 이론적으로 엄밀한 머신러닝의 새로운 기반을 마련했습니다.