Mechanistic Interpretability of RNNs emulating Hidden Markov Models

AI summary

Overview

Research area: Computational neuroscience and mechanistic interpretability of recurrent neural networks, specifically the reverse-engineering of RNNs trained to reproduce the output statistics of Hidden Markov Models.

Technical level: Advanced. The paper combines dynamical systems analysis (fixed points, Jacobian linearization, eigenvalues, perturbative expansions), optimal transport loss functions (Sinkhorn divergence), and neuroscience-style circuit dissection.

Scope: The paper trains vanilla RNNs to imitate the emission statistics of several families of Hidden Markov Models and then reverse-engineers the trained networks at three levels — global population dynamics, local state-space geometry, and single-neuron connectivity — to identify the mechanism by which continuous RNN dynamics produce discrete, stochastic state transitions.

What This Paper Is About

Hidden Markov Models (HMMs) describe natural behaviors as sequences of discrete latent states with stochastic transitions, while RNNs represent latent processes as continuous trajectories through state space, so the two appear conceptually incompatible. The authors ask whether RNNs can use continuous internal dynamics to generate the stochastic discrete-state transitions that HMMs implement, and if so, what mechanism the networks discover. To answer this, they build a training pipeline that fits vanilla RNNs directly to HMM outputs and then reverse-engineer the trained networks to reveal the underlying computations.

Key Contributions

1. A training paradigm for stochastic RNN behavior. Unlike previous work focused on deterministic, input-driven tasks, the authors introduce a training method that combines noise-driven RNNs with Sinkhorn divergence optimization, enabling the learning of the stochastic state transitions typical of HMMs.

2. A demonstration that RNNs can emulate HMM statistics. Vanilla RNNs accurately replicate the emission statistics of multiple HMM architectures, matching both transition dynamics and stationary distributions.

3. A multi-level mechanistic account of how RNNs generate discrete stochastic outputs, spanning global dynamics (noise-sustained orbital dynamics), local dynamics (clusters, kick-zones, and transitions), connectivity and single neurons (noise-integrating populations and kick neurons, validated by causal interventions), and the computational principle (self-induced stochastic resonance).

4. Evidence for a compositional dynamical primitive. Analyses across fully connected, cyclic, and linear-chain HMM architectures show the same dynamical motif is reused modularly, suggesting a compositional principle by which RNNs emulate complex discrete latent dynamics.

Main Findings

• RNNs replicate HMM emission statistics. Trained RNNs matched their target HMMs on four metrics: Euclidean distance between Sinkhorn-aligned output sequences, the 3×3 transition matrix of successive emissions, marginal observation frequencies, and observation volatility. Higher input dimensionality improved convergence, while increasing hidden size beyond 150 yielded only marginal benefits.

• Without input, activity collapses to a single fixed point. When initialized from random hidden states and receiving no external input, the trained networks converged to a single fixed point, with no evidence of separate attractors that might be expected to represent the discrete latent states of the target HMMs.

• Stochastic input sustains orbital dynamics. Under Gaussian input, activity evolves uni-directionally along stable closed orbits. The stochastic input pushes activity outward while the recurrent component pulls it back. Orbit radius scales linearly with input variance (tested at σ² ∈ {0.1, 1.0, 2.0, 3.0, 4.0}), a relationship confirmed by a perturbative analysis.

• Three local dynamical zones. Segmenting state space by residency time (RT) revealed three regimes with distinct signatures: clusters (RT > 8, frequent logit-gradient sign changes in the range 5–20, essentially only contracting eigenvalues), kick-zones (2 ≤ RT ≤ 8, moderate sign changes around 2–4, a few unstable directions), and transitions (RT < 2, fewer than 1 sign change, nearly deterministic directed flow). Similar zones appeared across all RNNs regardless of the target HMM architecture.

• Clusters are noise-sensitive; transitions are not. Under noise conditions ranging from identical (γ = 0) to partially resampled (γ = 0.5) to fully independent (γ = 1), trajectories initialized in transition regions showed minimal variability, whereas trajectories from cluster regions diverged — some crossing the kick-zone and others returning to the cluster.

• Kick neurons and noise-integrating populations. In the RNN trained on the two-state HMM, two separate triplets of neurons had pre-activations strongly negative within clusters, near zero in kick-zones, and positive during transitions, placing them near the ReLU threshold. Within each triplet, cells connect with mutual excitation while inhibiting the opposing triplet. Two larger neuronal groups (approximately 70 neurons each) form recurrent excitatory loops internally and project inhibitory connections to each other, interacting with the kick neurons through structured excitatory and inhibitory connections. Each triplet is linked to one transition direction.

• Causal interventions confirm the circuit's role. Ablation (μ = 0) of kick neurons or of the noise-drive from the corresponding integrating population trapped trajectories within their current cluster and eliminated critical eigenvalue pairs, consistent with reversion to a single fixed-point regime. Enhancement (μ = 2) produced overshoots beyond the target cluster while preserving the number of critical eigenvalue pairs. Control manipulations on neurons not part of the noise-integrating populations preserved cluster switching.

• Training shifts the network into a stochastic resonance regime. Across training epochs, the network first learns a single fixed point, then becomes unstable, and finally transitions to orbital dynamics. This transition coincides with eigenvalues moving near the imaginary axis, a characteristic double-descent in the loss curve, and the emergence of the expected second-order perturbation vector near the transition. The authors interpret the resulting behavior as self-induced stochastic resonance (SISR): slow noise integration in cluster zones combined with fast kick-neuron-driven resets produces quasi-periodic switching without any external periodic input.

• The mechanism generalizes compositionally. For linear-chain HMMs with more latent states (M ∈ {2, 3, 4, 5}), RNNs did not form additional slow regions along the orbits; instead they captured finer emission discretizations by modulating the alignment of readout axes with the plane of the orbital dynamics. For fully connected and cyclic architectures, RNNs developed multiple orbits connecting distinct pairs of slow regions, decomposable into the same fundamental motif.

Methodology in Plain English

The authors took a standard, vanilla RNN and fed it random Gaussian noise at every time step. The RNN's hidden state was updated with a ReLU nonlinearity, then projected linearly onto three output logits. Those logits were converted into categorical samples via Gumbel-Softmax (temperature set to 1), which acts as a differentiable stand-in for sampling.

The target sequences came from HMMs with three architectures: linear-chain models with 2, 3, 4, or 5 latent states over an observation alphabet of {1, 2, 3}; a 3-state fully connected model; and a 4-state cyclic model. Because the targets are probabilistic rather than deterministic, the authors used the Sinkhorn divergence — an optimal transport distance that finds a smoothed coupling between predicted and target sequences — as the loss function instead of a simple regression error.

They ran a systematic sweep: for the linear-chain HMMs, three hidden sizes (50, 150, 200) crossed with four input dimensionalities (1, 10, 100, 200), giving 48 unique configurations, each trained with three independent random seeds for a total of 144 models. The fully connected and cyclic architectures used hidden size 150 and input dimensionality 100 with three seeds each, adding six models. Training used 30,000 sequences of fixed length (100 for M = 2, 30 for M = 3 and 4, 40 for M = 5, 30 for the fully connected model, 40 for the cyclic model), batched at 4096 with the Adam optimizer at a learning rate of 0.001, gradient clipping at 0.9 for linear-chain and 0.3 for the others, converging typically within 1000 epochs. Each model trained in roughly 5–20 minutes on an NVIDIA RTX 4090 GPU.

To reverse-engineer the networks, the authors projected hidden states onto principal components, linearized the dynamics via Jacobians (using Möbius-transformed eigenvalues to count unstable directions), computed residency times and logit-gradient sign-change counts across state space, and ran causal interventions that scaled either kick-neuron activity or their noise-drive by a factor μ.

Why This Matters

The work addresses a conceptual gap between two dominant modeling frameworks in neuroscience: HMMs, which impose discrete latent states, and RNNs, which operate in continuous state space. It shows that unconstrained RNNs can recover discrete and continuous latent structure directly from data without any imposed topological priors, and that they naturally converge toward structured, biologically plausible circuit motifs. The identified motif — slow noise integration coupled with fast kick-triggered resets — offers a reusable "dynamical primitive" for generating probabilistic discrete transitions, which is a novel interpretation of how RNNs can perform internally driven, stochastic computation.

Potential applications suggested by the paper's framing:

• Modeling the latent structure of unconstrained, natural behaviors recorded in freely behaving animals, where HMMs have been used to decompose behavior into elemental motifs.
• Generating testable hypotheses about how biological neural circuits implement discrete behavioral modes through continuous internal dynamics.
• Interpreting latent dynamics inferred from large-scale electrophysiology and imaging recordings of thousands of neurons.
• Providing a mechanistic account for how stochasticity in neural circuits, rather than being a nuisance, could support probabilistic computation.

Industry relevance: The paper does not report industry benchmarks or commercial applications. Its relevance is primarily to basic and computational neuroscience, and to the machine learning interpretability community interested in how trained recurrent networks represent discrete structure.

Future Directions

• Scaling to more complex HMM structures. The authors state that their results lay the foundation for scaling to more complex latent structures, which would test how the compositional dynamical motif extends beyond the architectures studied here. They already note that convergence becomes less consistent as the number of hidden states in the target HMM increases.

• Testing whether biological circuits use similar motifs. A central open question is whether the slow-integration plus fast-kick mechanism is employed by biological neural circuits during naturalistic behavior observed in natural settings.

• Determining the generality of the mechanism across RNN architectures. The paper notes prior work showing that fixed-point topology can be invariant across implementations such as vanilla RNN, GRU, and LSTM, but this study trained only vanilla RNNs, leaving open whether the same motif emerges in other architectures.

• Clarifying how emission probabilities are shaped through learning. The authors propose that the oscillation period, linked to emission probabilities, "can be shaped through learning by adjusting the effective time-constant of the slow subsystem," which invites further investigation.

Target Audience

This paper benefits most: computational neuroscientists who use RNNs to model neural population dynamics and want to know what computations trained networks actually implement; machine learning researchers working on mechanistic interpretability of recurrent models, especially those interested in probabilistic or stochastic computation rather than deterministic tasks; researchers applying HMMs and related discrete-state models to animal behavior who want a continuous-dynamics alternative; and dynamical systems researchers interested in noise-induced phenomena such as self-induced stochastic resonance in trained networks.

Authors’ abstract

Recurrent neural networks (RNNs) provide a powerful approach in neuroscience to infer latent dynamics in neural populations and to generate hypotheses about the neural computations underlying behavior. However, past work has focused on relatively simple, input-driven, and largely deterministic behaviors - little is known about the mechanisms that would allow RNNs to generate the richer, spontaneous, and potentially stochastic behaviors observed in natural settings. Modeling with Hidden Markov Models (HMMs) has revealed a segmentation of natural behaviors into discrete latent states with stochastic transitions between them, a type of dynamics that may appear at odds with the continuous state spaces implemented by RNNs. Here we first show that RNNs can replicate HMM emission statistics and then reverse-engineer the trained networks to uncover the mechanisms they implement. In the absence of inputs, the activity of trained RNNs collapses towards a single fixed point. When driven by stochastic input, trajectories instead exhibit noise-sustained dynamics along closed orbits. Rotation along these orbits modulates the emission probabilities and is governed by transitions between regions of slow, noise-driven dynamics connected by fast, deterministic transitions. The trained RNNs develop highly structured connectivity, with a small set of "kick neurons" initiating transitions between these regions. This mechanism emerges during training as the network shifts into a regime of stochastic resonance, enabling it to perform probabilistic computations. Analyses across multiple HMM architectures - fully connected, cyclic, and linear-chain - reveal that this solution generalizes through the modular reuse of the same dynamical motif, suggesting a compositional principle by which RNNs can emulate complex discrete latent dynamics.