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.

Towards Scalable and Stable Parallelization of Nonlinear RNNs

A fast algorithm for solving large scale nonlinear optimization problems using RNN

Were RNNs All We Needed?

An Optimized and Energy-Efficient Parallel Implementation of Non-Iteratively Trained Recurrent Neural Networks

Parallelizing non-linear sequential models over the sequence length

RWKV: Reinventing RNNs for the Transformer Era

Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems

Predict globally, correct locally: Parallel-in-time optimization of neural networks

P-ADMMiRNN: Training RNN with Stable Convergence via An Efficient and Paralleled ADMM Approach

Parallel Trust-Region Approaches in Neural Network Training: Beyond Traditional Methods

Parallelizing Linear Transformers with the Delta Rule over Sequence Length

Asynchronous Parallel Stochastic Quasi-Newton Methods

Optimization of Recurrent Neural Networks on Natural Language Processing.

Finetuning Pretrained Transformers into RNNs

Linearizing Large Language Models

Scalable Neural Network Kernels

FlashRNN: Optimizing Traditional RNNs on Modern Hardware

A fast neural hybrid Newton solver adapted to implicit methods for nonlinear dynamics

Automatic Task Parallelization of Dataflow Graphs in ML/DL models

BPPSA: Scaling Back-propagation by Parallel Scan Algorithm