Gated RNN Numerics

This small project contains numerical experiments and modules used to reproduce and explore results from "Theory of Gating in Recurrent Neural Networks" (K. Krishnamurthy et al., Phys. Rev. X).

The main results are in the notebook Numerics.ipynb: gated RNN integrations, stability spectra, fixed-point solvers and Lyapunov estimates.

What you'll find here

• Example code to integrate a gated recurrent network with three variables (hidden state, update gate, output gate).
• Scripts to compute stability spectra and Jacobians around fixed points.
• Utilities to visualize time series, spectra, and phase diagrams.

Model

The repository simulates a gated RNN with $N$ units whose dynamics follows:

• Internal State (hidden $h$):

$\dfrac{dh}{dt} = \sigma_z(z) \odot \left(-h + J^h (\phi(h) \odot \sigma_r(r)) \right)$

• Update Gate ($z$):

$\tau_z \dfrac{dz}{dt} = -z + J^z \phi(h)$

• Output Gate ($r$):

$\tau_z \dfrac{dr}{dt} = -r + J^r \phi(h)$

Main developments

Open Numerics.ipynb and run the cells. The notebook demonstrates:

• Integrating gated RNN dynamics with random couplings.
• Plotting sample unit time-series and distributions.
• Computing Jacobian spectra and scanning parameters (gains, gate strengths)
• Solving DMFT-like fixed point equations and plotting transitions.

Quick example

import torch
import Dynamics as dyn
# Define Parameters
# Parameters of the integration
N_samples = 400
dt_samples = 1.0
T_discard = 0
dt_int = 0.1

#Parameters of the network
N = 1000 # Number of units
gains = [3.0 , 0. , 0.] # [gh , az , ar]
biases = [0.0 , 0.0 , 0.0] # [bh,bz,br]
taus = [1.0 , 1.0] # [tau_z,tau_r]
Js = (1/N**0.5)*torch.randn(3,N,N) # [Jh,Jz,Jr] Random Couplings

#Initial Conditions
x0 = torch.randn(3,N)

# Run the dynamics and save the time series
X = dyn.gated_recurrent_dynamics(N,x0,Js,gains,biases,
                                    taus,N_samples,dt_samples,
                                    T_discard,dt_integration=dt_int,show=True)

# The output is a tensor of shape [3,N_samples,N] where the first dimension 
# corresponds to the variables h,z and r, the second dimension corresponds to
# the time samples and the third dimension corresponds to the units.
# X = [h,z,r] where h.shape = z.shape = r.shape = [N_samples,N]

print(f'The shape of X is: {X.shape}')

Output:

Running Integrator ...
End of Integration
Integration Time 5.57 [seg] = 0.0928 [min]

The shape of X is: torch.Size([3, 400, 1000])

References

K. Krishnamurthy et al., "Theory of Gating in Recurrent Neural Networks", Phys. Rev. X (link: https://journals.aps.org/prx/abstract/10.1103/PhysRevX.12.011011)