Leakage and Second-Order Dynamics Improve Hippocampal RNN Replay

TL;DR

This work proposes the first model of temporally compressed replay in noisy path-integrating RNNs through hidden state momentum, connects it to underdamped Langevin sampling, and shows that, together with adaptation, it counters slowness while maintaining exploration.

Abstract

Biological neural networks (like the hippocampus) can internally generate "replay" resembling stimulus-driven activity. Recent computational models of replay use noisy recurrent neural networks (RNNs) trained to path-integrate. Replay in these networks has been described as Langevin sampling, but new modifiers of noisy RNN replay have surpassed this description. We re-examine noisy RNN replay as sampling to understand or improve it in three ways: (1) Under simple assumptions, we prove that the gradients replay activity should follow are time-varying and difficult to estimate, but readily motivate the use of hidden state leakage in RNNs for replay. (2) We confirm that hidden state adaptation (negative feedback) encourages exploration in replay, but show that it incurs non-Markov sampling that also slows replay. (3) We propose the first model of temporally compressed replay in noisy path-integrating RNNs through hidden state momentum, connect it to underdamped Langevin sampling, and show that, together with adaptation, it counters slowness while maintaining exploration. We verify our findings via path-integration of 2D triangular and T-maze paths and of high-dimensional paths of synthetic rat place cell activity.

Figures (15)

• Figure 1: Underdamped dynamics accelerate offline replay, adaptation slows it. Here we simulate a noisy RNN $r(t)$ that optimally path-integrates an Ornstein-Uhlenbeck process $s(t)$ from its velocity $s'(t)$. We assume $r(t)$ minimizes the loss in \ref{['eq:loss_upper']} and thus evolves according to its score function$\nabla_{r(t)} \log p(r(t))$ (Equations \ref{['eq:rnn_langevin']} and \ref{['eq:score_ou_main']}), performing a variant of Langevin sampling when no input is given. Above, we compare three modifiers of RNN activity: the default (no modification, a.k.a. overdamped), our proposed underdamped (momentum), and adaptation (negative feedback) dynamics. Each modifier affects the replay distribution $p(r(t))$ in different ways: underdamped sampling accelerates $p(r(t))$ towards $p(s(t))$, decreasing the distance between them, while adaptation slows convergence of $p(r(t))$ towards $p(s(t))$, increasing this distance.
• Figure 2: Leakage helps path-integration. Here we train RNNs on two tasks, ablating the leakage term. Leakage helps training, especially when losses increase with the masking difficulty $k$ (note that $k=1$ is equivalent to unmasked training, see Equation \ref{['eq:masked_training']}). Means are solid, standard deviations are faint.
• Figure 3: Underdampening and adaptation counter each other. Here we show replay from RNNs trained to path-integrate in T-maze or triangular environments. (a) Awake paths in each task form a mixture of Ornstein-Uhlenbeck processes, one for each direction of travel. Awake paths reach their endpoints and stay there. In (b) and (c) are mean replay paths (darkening over time) from T-maze and triangle tasks, simulated for the same time as awake paths. Standard deviations are faint, ideal path means are dashed. Like in \ref{['fig:1d_mean_std']}, adaptation slows convergence towards endpoints, underdampening quickens it. Also, the two mechanisms induce deviations that negate each other.
• Figure 4: Underdampening improves replay fidelity in the presence of adaptation. We compute the Wasserstein distance (dissimilarity) between awake and replay path distributions ($p( \left\{ \bm{s}(t) \right\}_{t=0}^T )$ and $p( \left\{ \bm{r} (t) \right\}_{t=0}^T )$), varying friction and adaptation strength (see \ref{['appendix:methods']} for details). While the two mechanisms both generally increase this distance, underdampening ($\lambda_v < 1$) decreases it if adaptation is nonzero. Like in \ref{['fig:paths']}, underdampening counters adaptation-induced deviations.
• Figure 5: Underdampening temporally compresses replay. We calculate how long it takes awake and replay trajectories to reach their endpoints. Underdampening ($\lambda_v < 1$) not only shortens this reach time, but makes it smaller than that of awake paths, temporally compressing awake activity. See \ref{['appendix:more_results']}\ref{['fig:reachtime_heatmaps_all']} for mean reach times. We do not include unbiased rat trajectories because they do not have defined endpoints, but we do confirm in \ref{['appendix:more_results']}\ref{['fig:mean_displacement']} that underdampening quickens them.

Theorems & Definitions (14)

• Definition 2.1
• Definition 2.2
• Lemma 2.5
• Theorem 2.7
• Theorem 2.8
• Theorem 3.2
• Remark 3.3
• Remark 3.4
• Remark 3.6
• Proposition 4.1