Drift-Diffusion Matching: Embedding dynamics in latent manifolds of asymmetric neural networks

TL;DR

The results extend attractor neural network theory beyond equilibrium, showing that asymmetric neural populations can implement a broad class of dynamical computations within low-dimensional manifolds, unifying ideas from associative memory, nonequilibrium statistical mechanics, and neural computation.

Abstract

Recurrent neural networks (RNNs) provide a theoretical framework for understanding computation in biological neural circuits, yet classical results, such as Hopfield's model of associative memory, rely on symmetric connectivity that restricts network dynamics to gradient-like flows. In contrast, biological networks support rich time-dependent behaviour facilitated by their asymmetry. Here we introduce a general framework, which we term drift-diffusion matching, for training continuous-time RNNs to represent arbitrary stochastic dynamical systems within a low-dimensional latent subspace. Allowing asymmetric connectivity, we show that RNNs can faithfully embed the drift and diffusion of a given stochastic differential equation, including nonlinear and nonequilibrium dynamics such as chaotic attractors. As an application, we construct RNN realisations of stochastic systems that transiently explore various attractors through both input-driven switching and autonomous transitions driven by nonequilibrium currents, which we interpret as models of associative and sequential (episodic) memory. To elucidate how these dynamics are encoded in the network, we introduce decompositions of the RNN based on its asymmetric connectivity and its time-irreversibility. Our results extend attractor neural network theory beyond equilibrium, showing that asymmetric neural populations can implement a broad class of dynamical computations within low-dimensional manifolds, unifying ideas from associative memory, nonequilibrium statistical mechanics, and neural computation.

Figures (15)

• Figure 1: Embedding dynamics in latent manifolds of RNNs.$a)$ Our framework focuses on RNNs with low-rank connectivity of the form $W=\Gamma W_s$, where $\Gamma \in \mathbb{R}^{n \times k}$ with $k\ll n$. We train these autonomous systems to encode arbitrary stochastic dynamics in a learnt affine subspace. $b)$ The training of the RNN can be reformulated as the training of a two-layer perceptron to approximate a target vector field. $c)$ Once the parameters of the RNN have been fitted, it can be integrated over time as an autonomous circuit which produces stochastic dynamics. Here we plot traces from 5 of 64 neurons in a network trained to approximate the van der Pol oscillator. When we project the dynamics into the learnt subspace, we recover a trajectory from the stochastic van der Pol system, which has been embedded into the RNN.
• Figure 2: Embedding attractors in neural networks.$a)$ Traces from 5 randomly selected neurons in a 64-neuron RNN trained to embed a stochastic VDP oscillator in a latent subspace. In the subspace, we see the characteristic VDP limit-cycle. $b-c)$ Traces from a 512-neuron RNN trained to embed the LA, and a 1024-neuron RNN trained to embed the DA.
• Figure 3: Symmetric-asymmetric decomposition of RNNs.$a)$ Decomposition of RNN trained to encode the VDP. The symmetric dynamics descend the contoured energy function. This appears to keep the process near the limit cycle, whilst the asymmetric component provides rotational dynamics. $b-c)$ Decompositions of RNNs trained to encode the LA and DA. For the DA, the symmetric dynamics keep the process near the attractor, whilst the asymmetric component provides rotational dynamics. For the LA, it appears the symmetric component pushes the dynamics away from the origin, and the asymmetric component provides a restoring force. $d-e)$ This is confirmed by plotting isopotentials of the energy functions for the LA and DA respectively.
• Figure 4: Reconstructing latent subspaces from trajectories. $a)$ Trajectories from a low-rank RNN trained to embed a $k-$dimensional system only have $k$ non-zero principal components. $b)$ With PCA, we can identify the affine subspace in which the dynamics are constrained. $c)$ The PCA representation is equivalent to original coordinates up to an affine linear transformation. Given $(\Gamma, \mathbf{b})$ can solve for this transformation to recover the original coordinates.
• Figure 5: 'Tilting' the energy landscape. Our approach uses network inputs to 'tilt' the energy landscape, encouraging the process to push a trajectory into a particular energy minimum, by making it a global minimum. This is a soft version of the child's 'labyrinth' game, where the board can be tilted such that the ball rolls into a specific hole -- though in the actual game, one tries rather hard to avoid them.