[論文レビュー] Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems

Jimmy T. H. Smith, Scott W. Linderman+1 more|arXiv (Cornell University)|Nov 1, 2021

Neural dynamics and brain function参考文献 46被引用数 7

ひとこと要約

この論文では、固定点と局所線形近似を訓練中に直接学習することで、再帰的ニューラルネットワーク（RNN）を逆設計する共同訓練フレームワークであるヤコビアンスイッチング線形動的システム（JSLDS）を導入する。一階テイラー展開に基づくスイッチング線形システムと学習可能な補助関数を統合することにより、JSLDSは訓練後の固定点最適化を不要とし、状態空間全体にわたる動的挙動の近似を可能にするとともに、固定点まわりの線形化された動的挙動によってRNNを正則化する。

ABSTRACT

Recurrent neural networks (RNNs) are powerful models for processing time-series data, but it remains challenging to understand how they function. Improving this understanding is of substantial interest to both the machine learning and neuroscience communities. The framework of reverse engineering a trained RNN by linearizing around its fixed points has provided insight, but the approach has significant challenges. These include difficulty choosing which fixed point to expand around when studying RNN dynamics and error accumulation when reconstructing the nonlinear dynamics with the linearized dynamics. We present a new model that overcomes these limitations by co-training an RNN with a novel switching linear dynamical system (SLDS) formulation. A first-order Taylor series expansion of the co-trained RNN and an auxiliary function trained to pick out the RNN's fixed points govern the SLDS dynamics. The results are a trained SLDS variant that closely approximates the RNN, an auxiliary function that can produce a fixed point for each point in state-space, and a trained nonlinear RNN whose dynamics have been regularized such that its first-order terms perform the computation, if possible. This model removes the post-training fixed point optimization and allows us to unambiguously study the learned dynamics of the SLDS at any point in state-space. It also generalizes SLDS models to continuous manifolds of switching points while sharing parameters across switches. We validate the utility of the model on two synthetic tasks relevant to previous work reverse engineering RNNs. We then show that our model can be used as a drop-in in more complex architectures, such as LFADS, and apply this LFADS hybrid to analyze single-trial spiking activity from the motor system of a non-human primate.

研究の動機と目的

• 訓練後の固定点解析の限界、たとえば固定点選択の曖昧さや線形近似における誤差蓄積を解消するため。
• 訓練後における数値最適化に依存せずに、固定点から局所線形動的挙動への連続的かつ状態空間全域へのマッピングを可能にする手法を開発するため。
• RNNの訓練を、非線形動的挙動が固定点まわりの線形化された挙動の間をスイッチングすることでよく記述される解に正則化するため。
• パラメータ共有を用いてスイッチング点の連続的多様体に一般化されたスイッチング線形動的システム（SLDS）を構築するため。
• LFADSのような複雑なアーキテクチャへのこのフレームワークの応用を可能にするため。

提案手法

• RNNの更新関数の学習済み固定点まわりの一次テイラー展開を用いる、新しいスイッチング線形動的システム（SLDS）と標準RNNを共同で訓練する。
• 任意の状態空間内の状態に対してRNNの固定点またはスローポイントを予測する学習可能な補助関数を導入する。
• SLDSがRNNの動的挙動を近似するよう正則化するとともに、スイッチングが予測された固定点で発生することを保証するための共同損失関数を定義する。
• RNNの更新関数のヤコビアンをSLDSにおける線形遷移行列として使用することで、固定点近傍での局所線形動的挙動がRNNの挙動を反映するようにする。
• 補助関数の出力をもとに条件付きで活性化される一組の線形パラメータを用いることで、スイッチ間でのパラメータ共有を実現する。
• GRUやLFADSのようなアーキテクチャにエンドツーエンドで適用し、単一試行の神経パルス発火データにおける神経動的挙動の直接的分析を可能にする。

実験結果

リサーチクエスチョン

• RQ1 共同訓練されたSLDSモデルは、訓練後の固定点最適化を必要とせずに、RNNの非線形動的挙動を正確に近似できるか？
• RQ2 学習可能な補助関数は、RNNの状態空間の任意の点を対応する固定点またはスローポイントにマッピングでき、全体的な動的挙動の解釈を可能にするか？
• RQ3 JSLDSフレームワークによる共同訓練は、固定点まわりの線形化された挙動の間をスイッチングする解にRNNを正則化するか？
• RQ4 LFADSのような複雑なアーキテクチャにこのJSLDSフレームワークを効果的に統合できるか？
• RQ5 JSLDSは、標準的な逆設計手法と比較して、線形近似の精度をどの程度向上させるか？

主な発見

• JSLDSフレームワークは、訓練後の固定点最適化の必要性を完全に排除し、固定点から線形動的挙動への状態空間全域への直接的マッピングを可能にした。
• スイッチング点の連続的多様体への一般化を実現しながらも、パラメータ数を一定に保ち、効率的かつスケーラブルな動的挙動近似を実現した。
• 合成タスクにおいて、標準的な線形化に基づく逆設計手法と比較して、顕著に低い誤差でRNNの動的挙動を正確に再構築した。
• 単一試行のサル運動皮膚のパルス発火データにおいて、LFADS-JSLDSハイブリッドは、条件に依存しない1つの線形システムが運動動的挙動を支配することを明らかにした。固有値と部分空間はjPCAの結果と非常に高く一致した。
• 3ビットメモリタスクにおいて、共同訓練プロセスがRNNの固定点構造を変化させ、線形化された動的挙動の近似精度が向上し、スイッチング線形解への正則化の証拠が得られた。
• GRUやLFADSを含む多様なタスクとアーキテクチャにおいて、動的挙動の捉え方が一貫して良好であり、JSLDSの固有値と神経データ解析手法（jPCAなど）との間で一貫した整合性が確認された。