[论文解读] 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

一句话总结

本文提出Jacobian Switching Linear Dynamical Systems（JSLDS），一种联合训练框架，通过在训练过程中直接学习固定点和局部线性近似，实现对循环神经网络（RNNs）的逆向工程。通过将可学习的辅助函数与基于RNN一阶泰勒展开的切换线性系统相结合，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逆向工程中后训练固定点分析的局限性，如固定点选择不明确和线性近似中的误差累积问题。
• 开发一种方法，实现在无需训练后数值优化的前提下，对全状态空间连续映射固定点至局部线性动力学。
• 通过正则化RNN训练过程，使网络的非线性动力学能被固定点附近线性化动力学之间的切换良好近似。
• 将切换线性动力系统（SLDS）推广至连续切换点流形，同时在不同切换间共享参数。
• 使该框架可成功应用于复杂架构（如LFADS），以分析神经元放电数据。

提出的方法

• 联合训练标准RNN与一种新型切换线性动力系统（SLDS），后者基于RNN更新函数在学习到的固定点处的一阶泰勒展开。
• 引入一个可学习的辅助函数，用于预测任意状态空间中的给定状态对应的RNN固定点或慢变点。
• 定义联合损失函数，使SLDS能近似RNN的动力学，同时确保切换发生在辅助函数预测的固定点处。
• 使用RNN更新函数的雅可比矩阵作为SLDS中的线性转移矩阵，使局部线性动力学能反映RNN在固定点附近的动态行为。
• 通过使用一组共享的线性参数并根据辅助函数输出条件性激活，实现在不同切换之间的参数共享。
• 在GRUs和LFADS等架构中端到端应用该框架，实现对单次试验神经放电数据中神经动力学的直接分析。

实验结果

研究问题

• RQ1 联合训练的SLDS模型是否能在无需后训练固定点优化的前提下，准确近似RNN的非线性动力学？
• RQ2 可学习的辅助函数是否能将RNN状态空间中的任意点映射到对应的固定点或慢变点，从而实现全局动力学解释？
• RQ3 与JSLDS框架联合训练是否能将RNN正则化为更优描述为在固定点附近线性化动力学之间切换的解？
• RQ4 JSLDS框架是否能成功集成到LFADS等复杂架构中，以分析真实神经数据？
• RQ5 与标准逆向工程方法相比，JSLDS在多大程度上提升了线性近似的准确性？

主要发现

• JSLDS框架成功消除了对后训练固定点优化的需求，实现了固定点到线性动力学的直接、全状态空间映射。
• 该模型将SLDS推广至连续切换点流形，同时保持参数数量恒定，实现了高效且可扩展的动力学近似。
• 在合成任务中，JSLDS在显著降低误差的前提下，实现了对RNN动力学的高精度重构，优于基于标准线性化的逆向工程方法。
• 在单次试验的灵长类运动皮层放电数据上，LFADS-JSLDS混合模型揭示了一个单一、与条件无关的线性系统主导运动动力学，其特征值和子空间与jPCA结果高度重叠。
• 在3位记忆任务中，联合训练过程改变了RNN的固定点结构，提升了线性化动力学的近似效果，并提供了向切换线性解正则化的证据。
• 该方法在多种任务和架构（包括GRUs和LFADS）中均表现出鲁棒性，且JSLDS的特征值与神经数据分析技术（如jPCA）保持一致对齐。