Non-stationary recurrent neural networks for reconstructing computational dynamics of rule learning

Updating behavioral policies in novel environments or upon shifting action-outcome contingencies is essential for survival. A number of studies have identified the medial prefrontal cortex (mPFC) as a key player in cognitive flexibility, rule learning and uncertainty detection. Animals learning a new behavioral paradigm, or switching between rules, show abrupt changes in their performance, accompanied by sharp transitions in neuronal population dynamics. However, the precise mechanisms regulating behavioral flexibility remain unclear. Here we aim to extract the computational mechanisms behind these behavioral and neurophysiological phenomena through dynamical systems reconstruction, employing an interpretable non-stationary piecewise-linear recurrent neural network (PLRNN) for this purpose. Our approach enables to model neuronal adaptation to changing environmental contingencies with trial-specific connectivity matrices regularized by smoothness and continuity priors. We train our model to reconstruct the non-stationary neuronal dynamics from multiple-single unit (MSU) recordings of the rodent's mPFC on a trial-by-trial basis while animals performed a probabilistic rule-shifting task. After model training, 1) the non-stationary PLRNN can accurately generate a variety of single-unit firing rate profiles, 2) task-related events can be decoded as well from the PLRNN’s latent space as from the original MSU activity, 3) change points identified in the PLRNN-generated activity and the original MSU recordings across trials tightly correlate, and 4) PLRNN-simulated trial-to-trial trajectories for both rules closely match those directly obtained from the data. Thus, after PLRNN training on the neural recordings, its behavioral and dynamical characteristics closely mimicked those observed in the real data. Moreover, we show that trial-specific connectivity matrices allow for highly accurate decoding of rule-type, but are not influenced by other task events. Hence, changes in the animal’s behavior appear to be based on alterations in the underlying functional connectivity. In conclusion, the non-stationary PLRNN offers a novel framework for investigating time-variant, neuro-dynamical phenomena during learning, plasticity, and development.