cs.LGFeb 20, 2026
SaveLeakage and Second-Order Dynamics Improve Hippocampal RNN Replay
Organizations: Department of Electrical and Computer Engineering, Rice University, USA · Mila - Quebec AI Institute & Université de Montréal, Canada
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 & tables
Figure 1: Underdamped dynamics accelerate offline replay, adaptation slows it. Here we simulate a noisy RNN that optimally path-integrates an Ornstein-Uhlenbeck process from its velocity . We assume minimizes the loss in Equation 8 and thus evolves according to its score function (Equations 11 and 15 ), 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 in different ways: underdamped sampling accelerates towards , decreasing the distance between them, while adaptation slows convergence of towards , 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 (note that is equivalent to unmasked training, see Equation 13 ). 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 Figure 1 , 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 ( and ), varying friction and adaptation strength (see Appendix A for details). While the two mechanisms both generally increase this distance, underdampening ( ) decreases it if adaptation is nonzero. Like in Figure 3 , 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 ( ) not only shortens this reach time, but makes it smaller than that of awake paths, temporally compressing awake activity. See Appendix F Figure 9 for mean reach times. We do not include unbiased rat trajectories because they do not have defined endpoints, but we do confirm in Appendix F Figure 10 that underdampening quickens them.
Figure 6: On average, underdampening increases exploration via path length. In here and Figure 7 we simulate replay paths for more time than awake paths, varying friction and adaptation strength. Shown above, underdampening ( ) increases the average length of replay paths. It may do so via increased replay speeds ( Appendix F Figure 10 ) or via additional path transitions ( Figure 7 ); the latter is how adaptation increases path length.
Figure 7: On average, underdampening increases exploration via regions visited. Another relevant aspect of exploration is transitions not present in awake activity. (a) Here we quantify these transitions via regions visited . Each region is an area where every point is closest to the same endpoint. In our T-maze and triangle tasks, awake paths go from starting points to endpoints and stay there, visiting two regions total (starting and ending). Meanwhile, adaptation ( ) can make replay paths visit multiple endpoints (regions visited ¿ 2). (b) When we use adaptation and underdampening ( ), we see that replay paths might visit (i) more, (ii) fewer, or (iii) the same regions (see Appendix F Figure 12 for more details). (c) However, on average, across several different trained models, underdampening increases regions visited, thus increasing adaptation.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Example decoded 2D replay trajectories in unbiased (left) and biased (right) rat tasks. Like in Figure 3 , trajectories get darker over time.
Figure 9: This is another version of Figure 5 , but now with mean reach times also shown.
Figure 10: Underdampening ( ) increases mean displacement of replay trajectories, especially at early timesteps. Similarly to ( McNamee et al., 2021 ) , here we analyze the mean displacement , or distance, of replay trajectories from their starting points as a function of time. For reference, mean displacement over time is also plotted for awake trajectories. In exploration experiments, we simulate replay for the duration of awake trajectories. Higher mean displacements over time generally correlate with increased path length and exploration. Note that awake activity is always calculated with , hence why awake statistics are the same across plots within the same task.
Figure 11: Underdampening ( ) increases replay trajectory variances. Like in Figure 10 , here we analyze the variance of trajectories at each timestep. Trajectories are 2D, so the variances plotted above are simply the averages of the variances along each coordinate. Underdampening increasing variance is another confirmation that underdampening complements adaptation-induced exploration.
Figure 12: This is another version of Figure 7 , but this time with the specific distributions of regions visited shown as bar plots above pairs of replay trajectory sets. Recall that replay paths start at yellow circles and end at symbols. In (a), underdampening clearly increases regions visited for several trajectories: instead of just going from to ( , replay paths also continue onwards to traverse between and . Meanwhile in (b), underdampening decreases regions visited for many replay trajectories: instead of traversing the whole triangle (4 regions visited), many replay paths stop in the region defined by . As for (c), underdampening did not affect regions visited since all paths started at , visited , and then returned to .
Figure 13: Mild underdampening can improve replay fidelity in the presence of adaptation in tanh RNNs. This is another version of Figure 4 , for T-maze and triangle tasks, but with tanh RNNs.
Figure 14: Underdampening increases speed in tanh RNN replay. This is another version of Figure 9 for T-maze and triangle tasks, but now with tanh RNNs.
Figure 15: In tanh RNNs, underdampening may increase or decrease exploration via path length. This is an alternate version of Figure 6 but with tanh RNNs. Underdampening has mixed effects on tanh RNN replay path length because of tanh saturation.