RNNs perform task computations by dynamically warping neural representations

AI summary

Overview

Research area: Interpretability of neural networks and computational neuroscience, combining Riemannian geometry with the theory of dynamical systems (recurrent neural networks, RNNs).

Technical level: Advanced. The paper relies on differential geometry (pullback metrics, tangent spaces, Gaussian curvature, geodesics), differential equations, and stochastic differential equations. The motivation is stated accessibly, but the derivations require graduate-level mathematics.

Scope: The paper develops a Riemannian-geometric framework that derives the manifold topology and time-varying geometry of a dynamical system's state from the manifold of its inputs, and uses it to explain how RNNs solve contextual decision-making and sequential working-memory tasks.

What This Paper Is About

Most work on the "geometry of neural representations" assumes static, time-independent inputs, while most work on computation-through-dynamics treats nonlinear dynamical systems as black boxes analysed only through fixed points and linearisations around them. This paper asks how the topology and geometry of a dynamical system's activity manifold relate to the manifold of its time-varying inputs, and how that geometry reveals the computation the system performs. The authors answer this by deriving the pullback Riemannian metric on the state manifold of a general class of dynamical systems, and then showing that RNNs solve tasks by dynamically warping their representations of task variables.

Key Contributions

1. A topology result linking inputs to states. The authors prove that if the inputs to a dynamical system lie on an m-dimensional manifold, the system's state is constrained to a manifold that is at most m + 1-dimensional, and give conditions under which that manifold is an immersed submanifold of the state space.
2. A pullback metric for dynamical systems. They derive the pullback of the metric from neural state space to the input manifold for a general class of dynamical systems, defined via an adjoint differential equation, together with a specific closed form for RNNs.
3. Demonstration in contextual decision-making. In an RNN trained on a contextual evidence integration task, they show the neural manifold becomes warped over time so as to compress irrelevant input information while preserving relevant input information.
4. Demonstration in sequential working memory. In an RNN trained to store and recall two sequentially presented inputs, they show activity lies on a hyper-torus whose intrinsic and extrinsic geometry dynamically warp to retrieve different memories at different time points.

Main Findings

• Static networks warp representations around decision boundaries. In a feedforward network with one hidden layer of three neurons trained to map inputs on a circle to binary outputs (θ ∈ [0, π) mapped to 1 and θ ∈ [π, 2π) mapped to −1), the single metric entry G_θθ = ‖dz/dθ‖² showed that space was stretched near the decision boundaries. An untrained network did not show this effect.
• Input topology constrains state topology. Theorem 3.1 states that inputs u ∈ M with M an m-dimensional manifold imply x(t) ∈ N = P(M × ℝ) for a projection map P; Corollary 3.2 states N is at most m + 1-dimensional. The relevant quantity is the input manifold dimension m, not the ambient input dimension d. Table 1 illustrates this with working memory (m = 2), decision making (m = 1) and timing (m = ∞), with corresponding d values of 3, 4 and 2, and navigation with d = 2.
• The metric is given by an adjoint system. For RNNs of the form ẋ = Wφ(x) − x + Bu(t), the metric is G_RNN with blocks built from ‖f‖², f^T A and A^T A, where A follows Ȧ = (Wφ′(x) − I)A + B du_κ/dκ(t) with A(0) = 0.
• Attractor systems have rank-deficient metrics. In a classic three-unit attractor RNN with two correlated, constant-in-time inputs lying on a line manifold, the metric at t = 0 is rank-deficient with a single non-zero entry G_tt = f(x, u(0))^T f(x, u(0)) and zero G_tu = G_ut = G_uu entries. As t → ∞, G_tt → 0 while G_uu > 0, reflecting convergence to a line of fixed points whose position depends on the input.
• Contextual decision-making compresses irrelevant inputs. In the stochastic contextual integration task (two noisy inputs, readout at t = 10, u_i ∈ [−0.2, 0.2]), the neural activity lies on a three-dimensional manifold for each of the two contexts, generated by the two input dimensions plus time. Initially the diagonal metric entries were roughly equal; near readout the time component approached zero and the component for the irrelevant input dropped substantially depending on context.
• The manifold becomes effectively lower-dimensional. Two eigenvalues of the metric — corresponding to time and the irrelevant input — decayed to zero, so the manifold became closer to one-dimensional by decision time and became pseudo-Riemannian at large times.
• Weight changes and dynamical modes align with the relevant input. The change in weights ΔW = W_K − W_0 was low-rank. Tangent vectors for both inputs were initially partially aligned with the leading right singular vector of the weight matrix; over time alignment increased for the relevant input and decreased for the irrelevant input.
• Neuronal embedding is structured but degenerate. Individual neuron responses were approximately Gaussian distributed along the tangent vectors. Participation in the relevant-input and time directions was correlated, whereas participation in the irrelevant-input and time directions was uncorrelated.
• Working memory uses a curved hyper-torus. In the sequential working-memory task (two inputs on a circular manifold, delay h ∈ [0.5, 3.5] ended by a go pulse), PCA revealed a toroidal manifold with dynamic geometry. Gaussian curvature was non-flat with both positive and negative values that were highly non-uniform over the torus (examined at t = 2 and at t = 6 output onset).
• Warping, not just realignment, underlies recall. During the input period the torus' shape formed, it remained stable over the delay, and during retrieval the RNN selectively compressed the stimuli that were not immediately being retrieved. The decoder was selectively aligned to the tangent basis vector encoding the relevant stimulus, while the subspace spanned by the top two principal components stayed stable across the delay.
• More memories need higher-dimensional manifolds. The authors prove that during the delay the activity lies on a hyper-torus whose dimension equals the number of stored items, and that to optimally encode stimuli this torus must be embedded in state space rather than merely immersed. Geodesic distances on the hyper-torus showed compression along irrelevant directions at different retrieval times.

Methodology in Plain English

The authors treat a dynamical system as a map that takes a whole input time-course and returns a trajectory of states. They formalise the collection of all possible input time-courses as a low-dimensional manifold, and the resulting set of possible trajectories as another manifold inside the state space.

To characterise that second manifold, they adapt a standard tool from deep-learning geometry: the pullback metric. Instead of measuring distances between input points directly, the metric measures how much the network's internal state changes when the input changes slightly. The key technical step is that, for a dynamical system, this metric cannot be computed at a single time point alone; it is generated by an adjoint differential equation that propagates sensitivities forward through time alongside the state itself. This yields an explicit metric matrix for RNNs.

They then train small RNNs on two standard neuroscience tasks — a contextual evidence integration task modelled with stochastic differential equations, and a sequential working-memory task with a variable delay — and inspect the metric, its eigenvalues, its geodesics, and the Gaussian curvature of the resulting manifold. They complement this with analyses of the trained weight matrices and of how individual neurons are embedded along the manifold's tangent directions.

Why This Matters

Impact on research. The framework offers an exact, nonlinear characterisation of a dynamical system's activity manifold, going beyond the dominant practice of describing RNN computation through fixed points and local linearisations. Because the input manifold does not depend on the trained parameters, any change in the geometry is attributable to the computation the network learned. The authors also note that their intrinsic-dimension constraint is complementary to the embedding-dimension constraint familiar from low-rank RNN work — an RNN with rank-r weights and d-dimensional inputs confines the state largely to an (r + d)-dimensional subspace, whereas this work constrains the intrinsic dimension of the nonlinear state manifold. The paper's discussion flags a degeneracy: the same metric can arise from different architectures.

Potential applications (implied by the paper's framing rather than demonstrated in it):

• Interpreting trained recurrent models in computational neuroscience, where RNNs defined by differential equations are used routinely to investigate neural computations underlying behavioural tasks.
• Analysing diffusion models based on stochastic differential equations, which the paper cites as achieving state-of-the-art performance for generative modelling of images.
• Diagnosing and designing models elsewhere that dynamical systems are central, including physics, chemistry and genomics, where the paper notes methods to infer dynamical systems from data are now central.
• Guiding architecture design: the paper suggests choosing the input manifold and weight rank jointly could let practitioners target specific intrinsic and embedding dimensionalities.

Industry relevance. The results bear on model interpretability and controllability for sequence models and for continuously-defined generative models — specifically, on predicting and verifying which input features a trained system amplifies, compresses, or discards as a function of time and context. The paper reports no deployed system, benchmark comparison, or dataset scale, so these are directions the framework enables rather than product-level results.

Future Directions

• Resolving the geometry–architecture degeneracy. The authors state in the Limitations that the relationship between geometry, dynamics and connectivity is degenerate, since the same metric can arise through different network architectures. Working out what additional constraints distinguish architectures remains open.
• Unifying intrinsic and embedding dimensionality. The paper proposes combining its intrinsic-dimension constraints with low-rank RNN embedding constraints to design networks with specified intrinsic and embedding dimensionalities.
• Extending beyond the settings already handled. The supplementary material (S7) extends the results to dynamical systems defined on non-Euclidean spaces, to initial states that themselves vary on a low-dimensional manifold, and to stochasticity in the dynamics; how far these extensions reach in practice is not settled in the truncated text.
• Applying the metric to other task families and model classes. The paper demonstrates the approach on two tasks (contextual decision-making and sequential working memory); whether the warping picture extends to the other task types listed in Table 1 (timing, navigation) and to other dynamical systems models is left open. The Discussion text is truncated in the provided content, so further listed limitations are not reported here.

Target Audience

This paper is aimed at researchers in computational neuroscience and machine-learning interpretability who study recurrent networks and computation-through-dynamics, and at theorists comfortable with differential geometry and dynamical systems. Readers interested in neural manifold analyses, attractor-based accounts of RNN computation, or the geometry of representation learning will find the framework most directly useful; readers without a differential-geometry background will need to rely on the conceptual explanations and figures rather than the derivations.

Authors’ abstract

Analysing how neural networks represent data features in their activations can help interpret how they perform tasks. Hence, a long line of work has focused on mathematically characterising the geometry of such "neural representations." In parallel, machine learning has seen a surge of interest in understanding how dynamical systems perform computations on time-varying input data. Yet, the link between computation-through-dynamics and representational geometry remains poorly understood. Here, we hypothesise that recurrent neural networks (RNNs) perform computations by dynamically warping their representations of task variables. To test this hypothesis, we develop a Riemannian geometric framework that enables the derivation of the manifold topology and geometry of a dynamical system from the manifold of its inputs. By characterising the time-varying geometry of RNNs, we show that dynamic warping is a fundamental feature of their computations.