Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

기술 요약: 신경 역학을 위한 모델 선택 기능이 포함된 계산 인지형 칼만 필터링

문제 정의
본 논문은 현대의 대규모 신경 기록 데이터에서 동적인 잠재 변수를 모델링할 때 발생하는 계산적 및 통계적 과제를 다룹니다. 베이지안 방법론(예: 가우스 과정 요인 분석, GPFA)은 명시적인 사전 확률과 불확실성 정량화를 제공하지만, 시간 해상도에 대해서는 세제곱(cubic)의 복잡도를, 공간 해상도(뉴런 수)에 대해서는 제곱(quadratic)의 복잡도를 가지므로 현대의 대규모 데이터셋에는 적용이 불가능합니다. 반대로, 과매개변수화된 딥러닝 모델(예: LFADS)은 확장성은 우수하지만, 견고한 불확실성 정량화가 부족하고 유도 편향(inductive bias)의 부재를 극복하기 위해 방대한 양의 시행(trial) 데이터를 필요로 합니다.

저자들은 신경과학(예: 인간 뇌 연구)에서 흔히 나타나는 특정 "규모 불균형 레짐(scale-imbalanced regime)"을 식별했습니다. 즉, 기록된 뉴런의 수가 실험 시행 횟수보다 훨씬 많은 상황입니다. 이 레짐에서 표준 딥러닝 모델은 데이터 부족으로 인해 성능이 저하되는 반면, 전통적인 베이지안 근사법은 자신의 계산적 제약으로 인해 발생하는 오류를 고려하지 못해 과도하게 확신하는("분산 결핍형") 불확실성 추정치를 내놓는 경향이 있습니다.

방법론: 계산 인지형 상태 공간 모델 (CASSM)
저자들은 모델 선택 기능을 포함하도록 계산 인지형 칼만 필터(CAKF)를 확장한 확률적 수치 프레임워크인 계산 인지형 상태 공간 모델(CASSM)을 제안합니다. 핵심 방법론은 다음과 같습니다:

1. 투영 및 계산 불확실성: CASSM은 "진정한" 사후 확률을 직접 계산하는 것이 계산적으로 불가능함을 인정합니다. 근사 오차를 무시하는 대신, 이를 "계산 불확실성"으로 명시적으로 모델링합니다. 이 방법은 학습된 정책 행렬 $S_k$를 사용하여 고차원 관측값($y_k \in \mathbb{R}^n$)을 저차원 공간($\check{y}_k \in \mathbb{R}^{\check{n}}$)으로 투영합니다. 이를 통해 혁신 행렬(innovation matrix) 역행렬 계산의 복잡도를 $O(n^3)$에서 $O(\check{n}^3)$으로 줄입니다.
2. 저계수 공분산 업데이트: 메모리 제약($O(d^2)$의 상태 공분산)을 해결하기 위해, CASSM은 전체 행렬 대신 "스키니(skinny)"한 저계수 제곱근 구조를 활용하여 공분산 업데이트를 저장합니다. 이는 메모리 요구량을 $O(dr)$로 감소시킵니다.
3. PCA 아날로그를 통한 모델 선택: 투영 정책 $S_k$는 정보 이론적 관점에서 해석됩니다. 정리 1(Theorem 1)은 잠재 사후 엔트로피를 최소화하는 최적의 행동(투영)이 혁신 공분사의 상위 고유벡터(주성분 분석, PCA와 유사)에 대응함을 입증합니다. 이 프레임워크는 기대 불확실성이 가장 큰 방향에 정렬되도록 이러한 행동을 학습합니다.
4. 새로운 학습 목적 함수: 저자들은 증거 하한(ELBO)에서 유도된 특정 손실 함수를 도입합니다. 투영 시 성능이 저하될 수 있는 표준 음의 로그 가능도(NLL)와 달리, CASSM 목적 함수는 데이터 적합 항과 발산(divergence) 항 사이의 균형을 맞춥니다. 이 발산 항은 근사 사후 확률이 정확한 필터링 분포로부터 벗어나는 정도를 벌칙(penalty)으로 부여하여, 계산적 근사가 증가함에 따라 불확실성 추정치가 적절히 넓어지도록 보장합니다.
5. 공간적 사전 확률: 가우스 과정 표현을 상태 공간 모델로 변환함으로써, CASSM은 공간 커널 $\Sigma^{(x)}$를 통해 뉴런의 좌표나 세포 유형과 같은 공간적 사전 확률을 명시적으로 인코딩할 수 있게 하며, 이는 GPFA나 LFADS와 같은 기존 방법론에서는 암묵적으로만 처리되었던 기능입니다.
6. 전처리 및 안정성: 파이프라인에는 스파이크 데이터에 대한 분산 안정화 변환(Ansarcme)이 포함되며, 저계수 업데이트의 절단(truncation) 과정에서 수치적 안정성을 보장하기 위해 미분 가능한 특이값 분해(Moore-Penrose 유사 역행렬 기반)를 사용합니다.

주요 기여

• CASSM 프레임워크: 규모 불균형 레짐(시행 횟수 $\ll$ 뉴런 수)에 특화되어 설계된 확장 가능한 확률적 수치 방법론입니다.
• 불확실성 보정이 포함된 트랙터블 추론: 대규모 상태 공간에서도 실행 가능한 추론을 제공하는 동시에, 계산적 근사 오차를 고려한 보정된 불확실성을 제공하여 다른 근사법에서 나타나는 과도한 확신 문제를 방지합니다.
• 새로운 손실 함수: 정확한 필터링 분포로부터의 발산을 벌칙으로 부여함으로써 효율적인 하이퍼파라미터 최적화와 모델 선택을 가능하게 하여, 모델이 유익한 데이터 투영을 학습할 수 있도록 합니다.
• 명시적 공간 모델링: 뉴런에 대한 공간적 사전 확률을 상태 공간 모델 내에 명시적으로 통합하는 기술을 제공하여, 공간 구조를 암묵적으로 처리하는 기존 잠재 변수 모델들과 차별화합니다.

결과
저자들은 합성 및 실제 데이터셋 모두에서 CASSM을 평가했습니다:

• 합성 로렌츠(Lorenz) 데이터: CASSM은 뉴런 수에 대해 선형적 확장을 보여주며, 대규모 환경($10^4$ 뉴런)에서 GPFA 및 bGPFA보다 계산 시간이 우수합니다. 규모 불균형 레짐에서 CASSM은 LFADS와 비교하여 경쟁력 있는 예측 성능(MSE)을 유지하는 반면, LFADS는 데이터가 제한될 때 급격히 성능이 저하됩니다.
• 불확실성 보정: 커버리지 실험 결과, CASSM은 과도하게 확신하는 경향이 있는 bGPFA 및 GPFA에 비해 훨씬 더 잘 보정된 95% 신뢰 구간을 제공합니다.
• 실제 신경 데이터 (Neural Latents Benchmark): Area 2 Bump 및 DMFC와 같은 데이터셋에서 CASSM은 LFADS 및 bGPFA와 경쟁할 만한 성능을 보입니다. "MC Maze" 데이터셋(뉴런 140개, 시행 1720회)에서 CASSM은 LFADS와 대등한 공동 평활화(co-smoothing) 점수를 달성했습니다.
• 제브라피시 데이터셋: 약 88,000개의 뉴런과 단 96회의 시행만 있는 대규모 데이터셋에서 bGPFA는 메모리 부족 오류로 실패합니다. CASSM은 이 데이터를 성공적으로 처리하며, 공간적 사전 확률(CASSM-SP)을 사용할 경우 LFADS와 통계적으로 유사한 성능을 보이면서도 보정된 불확실성을 제공합니다.

의의 및 주장
본 논문은 CASSM이 연구자들에게 데이터셋의 특성에 따라 어떤 동적 잠재 변수 모델을 선택해야 하는지에 대한 "로드맵"을 제공한다고 주장합니다. 저자들은 딥러닝 방법(LFADS)이 뉴런 수보다 시행 횟수가 많을 때 우수하지만, 대규모 신경과학 연구에서 전형적인 규모 불균형 레짐에서는 베이지안 방법이 필수적이라고 주장합니다. CASSM은 딥 네트워크처럼 계산적으로 확장 가능하면서도 베이지안 추론의 유도 편향과 불확실성 정량화 능력을 유지함으로써 이 간극을 메웁니다. 저자들은 CASSM이 이 특정 레짐에 매우 특화되어 있으며 선형 방법론이기에, 기저 역학이 매우 비선형적이거나 근사의 계수(rank)를 잘못 선택할 경우 어려움을 겪을 수 있다는 점을 겸허히 명시합니다. 그러나 대상 레짐 내에서 CASSM은 데이터 집약적인 딥 네트워크에 대한 경쟁력 있는 대안이자 우수한 불확실성 보정 능력을 제공합니다.