Detecting Invariant Manifolds in ReLU-Based RNNs

TL;DR

The paper addresses the challenge of understanding RNN dynamics by focusing on invariant manifolds for ReLU-based PLRNNs. It introduces a semi-analytical algorithm that computes stable and unstable manifolds within each linear subregion and composes them across regions to reconstruct global manifolds, enabling basin delineation and detection of homoclinic/heteroclinic intersections indicative of chaos. Validation spans classical chaotic and multistable systems (Duffing, Lorenz-63, multi-choice decision tasks) and extends to an empirical cortical neuron recording, demonstrating the method's capacity to reveal underlying dynamical mechanisms from data. This work advances explainability in DS reconstruction with RNNs by providing a scalable, geometry-aware tool for probing the state-space skeleton that governs RNN behavior and computational capabilities.

Abstract

Recurrent Neural Networks (RNNs) have found widespread applications in machine learning for time series prediction and dynamical systems reconstruction, and experienced a recent renaissance with improved training algorithms and architectural designs. Understanding why and how trained RNNs produce their behavior is important for scientific and medical applications, and explainable AI more generally. An RNN's dynamical repertoire depends on the topological and geometrical properties of its state space. Stable and unstable manifolds of periodic points play a particularly important role: They dissect a dynamical system's state space into different basins of attraction, and their intersections lead to chaotic dynamics with fractal geometry. Here we introduce a novel algorithm for detecting these manifolds, with a focus on piecewise-linear RNNs (PLRNNs) employing rectified linear units (ReLUs) as their activation function. We demonstrate how the algorithm can be used to trace the boundaries between different basins of attraction, and hence to characterize multistability, a computationally important property. We further show its utility in finding so-called homoclinic points, the intersections between stable and unstable manifolds, and thus establish the existence of chaos in PLRNNs. Finally we show for an empirical example, electrophysiological recordings from a cortical neuron, how insights into the underlying dynamics could be gained through our method.

Figures (10)

• Figure 1: One-dimensional illustration of the iterative procedure for computing stable manifolds with subregion boundaries indicated in purple-dashed. Step 1: The stable manifold (black) is initialized using the stable eigenvector $E^{\{+1\}}$ (blue) of the saddle point (green), and sample points (orange) are placed along it. Step 2: These points are propagated until they enter another linear subregion. Step 3: A new segment of the manifold is determined. Step 4: Repeating this process iteratively reconstructs the full global structure of the stable manifold. For a trained model example see Fig. \ref{['fig: evolution2']}.
• Figure 2: A) Top-left: Relative proportion of subregions with a positive determinant of the Jacobian ($\text{det}(\boldsymbol{J})>0$), $\mathcal{S}_+$, as a function of latent space dimensionality $M$ for different proportions $|\mathcal{S}_\text{reg}|/|\mathcal{S}|$. Medians $\pm$ interquartile range are shown. Top-right: Runtime and reconstruction quality (bottom) as a function of latent space dimensionality $M$ for different proportions $|\mathcal{S}_\text{reg}|/|\mathcal{S}|$ of subregions for which invertibility was enforced by regularization ($\lambda=0.1 \exp({M})$ in Eq. \ref{['eq:Regularization']}). Means across $100$ different training runs on the Lorenz-63 system $\pm$ SD are shown. Reconstruction quality was assessed through (dis-)agreement in attractor geometry (bottom-left; $D_{\text{stsp}}$, see Appx. \ref{['sec:Dstsp']}) and $20$-step-ahead prediction error (bottom-right). B) Top: $20$-step-ahead prediction error, PE(20), as a function of the number of training epochs when the invertibility regularization, Eq. \ref{['eq:Regularization']}, was turned off (blue) vs. on (orange), for a damped nonlinear oscillator. Bottom: Convergence to a predefined performance criterion ($\text{PE}(20) \leq 0.5$) was significantly faster with the regularization turned on vs. off. Median across $20$ trained models, error bands = interquartile range. C) Same as B for Lorenz-63 system.
• Figure 3: Model validation. A) Two examples of saddle points (half-green) with stable (gray solid lines) and unstable (red solid lines) manifolds determined by our algorithm, and points (black/ red dots, respectively) sampled by the analytical resp. backward/ forward map, showing that these all fall onto the analytically determined manifolds. B) Basins of attraction (light green, confirmed by sampling initial conditions and tracing their trajectories) of a stable fixed point (green dot) delineated by the stable manifold (black) of a 4-cycle (left) or 3-cycle (right).
• Figure 4: A) Reconstruction of the Duffing system by a shPLRNN ($M=2$, $H=10$). Trajectories drawn from the actual Duffing system in blue, identified fixed points in green, and in gray the stable manifold of the saddle in the center separating the two basins of attraction as determined by our algorithm. B) $3$d subspace of the state space of an ALRNN ($M=15$, $P=6$) trained on a 2-choice decision making task, with true trajectories in blue. Two point attractors (green) were identified, with the stable manifold (black/gray) of a saddle (half-green) in the center separating the two basins. The stable manifold (basin boundary) consists of different planar pieces -- note that for visualization these are projected down from a truly $15$d space into a $3$d subspace (accounting for some of the 'folded' appearance). C) Reconstruction for a shPLRNN ($M=3$, $H=20$) trained on the Lorenz63 system, with true trajectories in blue. The Lorenz63 system has two saddle-spirals in the center of the two lobes and a saddle at the bottom, which were correctly located by the shPLRNN. In black and red are the stable and unstable manifolds, respectively, of the saddle as identified by our algorithm (left), while on the right as computed by numerical continuation of the original Lorenz-63 system. The close agreement indicates the shPLRNN has correctly recovered the state space structure of the underlying system, although having been trained on trajectories from the actual attractor only. D) ALRNN ($M=25$, $P=6$) trained on electrophysiological recordings from a cortical neuron. Left: Time series of membrane voltage (true: gray, model-simulated: black); right: $2$d projection of the ALRNN's state space with stable manifold of a saddle (black) separating the basins of attraction of a stable fixed point (green) and the 38-cycle (orange dots) corresponding to the spiking process. Note that the true stable manifold is a $24$d curved object, which for visualization purposes is represented here by a locally linear approximation in the shown $2$d subspace.
• Figure 5: A) Stable (black) and unstable (red) manifolds of a saddle point (green dot) and their homoclinic intersections as identified by our algorithm. B) Structure of the chaotic attractor caused by these homoclinic intersections. C) Bifurcation diagram as a function of bias parameter $h_1$. D) Lyapunov exponents across the $h_1$-range for which the chaotic attractor exists.

Theorems & Definitions (8)

• Definition 1: Un-/stable manifold
• Definition 2: Basin of attraction
• Definition 3: Homoclinic orbit
• Definition 4: Heteroclinic orbit
• Example 1: 2d manifold in $\mathbb{R}^3$
• Remark I.1
• Proposition I.2
• proof