This repository contains the code for my MEng thesis, To attract or to oscillate: Validating dynamics with behavior. All code is written in Python, and recurrent neural network (RNN) models are written and trained in JAX, Flax, and Optax.

RNNs can learn two distinct dynamical systems to compute modular arithmetic, specifically $a + b + c \equiv 0\ (\text{mod } 3)$. One dynamical system is characterized by a lattice of fixed-point attractors, termed the attractive mechanism, and the other is characterized by a limit cycle, termed the oscillatory mechanism. These two systems exhibit unique psychometric curves when computing modular arithmetic for ambiguous stimuli and unique generalization curves when trained on partial datasets. Fundamentally, they differ in how they encode stimuli: the attractive mechanism encodes stimuli as vectors, while the oscillatory mechanism encodes stimuli as angles.

In the results/ and scripts/ folders, there are two subfolders:

• development/ - Jupyter notebooks used for the development of the src/ python package
• experiments/ - Jupyter notebooks used to train models and generate the results discussed in the thesis

JAX is an incredibly powerful Python package. In the context of training RNNs, JAX's scan function is significantly faster than using a for loop. I was able to train 16,128 RNNs on the MIT SuperCloud HPC in about 60 hours. Given that training and analyzing RNNs has become routine in computational neuroscience, I'm betting that the field will shift to using JAX, Flax, and Optax.