Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

技術要約：神経ダイナミクスのための計算量を考慮したモデル選択付きカルマンフィルタリング

問題提起
本論文は、現代の大規模な神経記録における動的な潜在変数のモデリングに伴う、計算量および統計的な課題に対処している。ガウス過程因子解析（GPFA）などのベイズ的手法は、明示的な事前分布と不確実性の定量化を提供するが、時間解像度に対しては3次の計算量、空間解像度（ニューロン数）に対しては2次の計算量を持つため、現代のデータセットには適用不可能である。対照的に、過学習気味のディープラーニングモデル（LFADSなど）はスケーラビリティに優れるが、強力な誘導バイアスを欠くため、不確実性の定量化が堅牢性に欠け、多くの試行回数を必要とする。

著者らは、神経科学（例：ヒトの脳研究）で頻繁に見られる特定の「スケール不均衡レジーム（scale-imbalanced regime）」（記録されたニューロン数が実験試行回数を大幅に上回る状況）を特定している。このレジームでは、標準的なディープラーニングモデルはデータの不足により性能が低下する一方で、従来のベイズ近似は、それら自身の計算上の制約によって導入される誤差を考慮できないことが多く、結果として過信的（分散が不足した状態）な不確実性推定に陥る。

手法：計算量を考慮した状態空間モデル（CASSM）
著者らは、モデル選択を含む計算量を考慮したカルマンフィルタ（CAKF）を拡張した、確率的な数値フレームワークである「計算量を考慮した状態空間モデル（CASSM）」を提案する。コアとなる手法は以下の通りである：

1. 投影と計算上の不確実性： CASSMは、「真の」事後分布を直接計算することが計算上不可能であることを認めている。近似誤差を無視するのではなく、「計算上の不確実性」を明示的にモデル化する。本手法は、高次元の観測値（$y_k \in \mathbb{R}^n$）を、学習された方策行列 $S_k$ を用いて低次元空間（$\check{y}_k \in \mathbb{R}^{\check{n}}$）へと投影する。これにより、イノベーション行列の逆行列計算の複雑さを $O(n^3)$ から $O(\check{n}^3)$ へと削減する。
2. 低ランク共分散更新： メモリ制約（状態共分散に対する $O(d^2)$ のコスト）に対処するため、CASSMは低ランクのダウンデイト構造を利用する。共分散の更新としてフル行列ではなく、「細長い（skinny）」低ランク平方根を保持することで、メモリ要件を $O(dr)$ に削減する。
3. PCAに類似したモデル選択： 投影方策 $S_k$ は、情報理論的な観点から解釈される。定理1は、潜在事後分布のエントロピーを最小化する最適なアクション（投影）が、イノベーション共分散のトップ固有ベクトルに対応すること（主成分分析（PCA）に類似）を確立している。本フレームワークは、最大の期待不確実性が存在する方向に一致するように、これらのアクションを学習する。
4. 新規の学習目的関数： 著者らは、証拠下限（ELBO）から導出された特定の損失関数を導入している。投影下でも劣化しやすい標準的な負の対数尤度（NLL）とは異なり、CASSMの目的関数は、データ適合項とダイバージェンス項のバランスをとる。このダイバージェンス項は、近似事後分布が正確なフィルタリング分布から逸脱することを罰することで、計算近似が増大するにつれて不確実性の推定が適切に拡大するように機能する。
5. 空間的事前分布： ガウス過程表現を状態空間モデルへと翻訳することにより、CASSMは空間カーネル $\Sigma^{(x)}$ を介して、ニューロンの座標や細胞型などの空間的事前分布を明示的に符号化することを可能にする。これは、GPFAやLFADSのような従来の手法では暗黙的であった特徴である。
6. 前処理と安定性： パイプラインには、スパイクデータに対する分散安定化変換（Anscombe変換）が含まれ、低ランク更新の切り捨て中の数値的安定性を確保するために、微分可能な特異値分解（ムーア・ペンローズ擬似逆行列に基づく）を利用している。

主な貢献

• CASSMフレームワーク： スケール不均衡レジーム（試行数 $\ll$ ニューロン数）に特化して設計された、スケーラブルで確率的な数値的手法。
• 不確実性較正を伴う扱いやすい推論： 高次元の状態空間においても扱いやすい推論を実現しつつ、計算近似による誤差を考慮した較正済みの不確実性を提供し、他の近似手法に典型的な過信を回避する。
• 新規の損失関数： 正確なフィルタリング分布からのダイバージェンスを罰することで、効率的なハイパーパラメータ最適化とモデル選択を可能にし、モデルが情報量の多いデータ投影を学習できるようにする。
• 明示的な空間モデリング： ニューロンに対する空間的事前分布を状態空間モデル内に明示的に組み込む手法を提供し、空間構造を暗黙的に扱う従来の潜在変数モデルと差別化する。

結果
著者らは、合成データおよび実世界のデータセットの両方でCASSMを評価している：

• 合成Lorenzデータ： CASSMはニューロン数に対して線形のスケーリングを示し、大規模なスケール（$10^4$ ニューロン）において、計算時間の面でGPFAおよびbGPFAを凌駕する。スケール不均衡レジームにおいて、CASSMはLFADSと比較して競争力のある予測性能（MSE）を維持する一方、LFADSは試行数が限られると急速に性能が低下する。
• 不確実性の較正： 被覆率（Coverage）の実験により、CASSMは、過信傾向にあるbGPFAやGPFAと比較して、有意に精度が高い（較正された）95%信頼区間を提供することが示された。
• 実世界の神経データ（Neural Latents Benchmark）： Area 2 BumpやDMFCなどのデータセットにおいて、CASSMはLFADSやbGPFAと同等の性能を示す。 「MC Maze」データセット（140ニューロン、1720試行）では、CASSMはLFADSと同等のコー・スムージング・スコアを達成した。
• ゼブラフィッシュ・データセット： 約88,000個のニューロンとわずか96回の試行を含む大規模なデータセットにおいて、bGPFAはメモリ不足エラーにより失敗する。CASSMは、このデータを正常に処理でき、空間的事前分布（CASSM-SP）を利用した場合、LFADSと統計的に同等の性能を発揮しつつ、較正された不確実性を提供する。

意義と主張
本論文は、CASSMが研究者に対し、データセットの特性に基づいて動的な潜在変数モデルを選択するための「ロードマップ」を提供すると主張している。著者らは、試行数がニューロン数を上回る場合にはディープラーニング手法（LFADS）が優れているが、大規模神経科学に典型的なスケール不均衡レジームにおいては、ベイズ的手法が不可欠になると論じている。CASSMは、ディープネットワークのように計算量的にスケーラブルでありながら、ベイズ推論の誘導バイアスと不確実性の定量化を保持することで、このギャップを埋めるものである。著者らは、CASSMがこの特定のレジームに高度に特化しており、線形手法であるため、基礎となるダイナミクスが高度に非線形である場合や、近似のランクが不適切に選択された場合には苦戦する可能性があることを謙虚に述べている。しかし、対象となるレジームにおいて、CASSMは、データ要求量の多いディープネットワークに対して、優れた不確実性の較正を備えた競争力のある代替手段を提供する。