Emergent mechanisms for long timescales depend on training curriculum and affect performance in memory tasks

TL;DR

It is suggested that adapting timescales to task requirements via recurrent interactions allows learning more complex objectives and improves the RNN's performance.

Abstract

Recurrent neural networks (RNNs) in the brain and in silico excel at solving tasks with intricate temporal dependencies. Long timescales required for solving such tasks can arise from properties of individual neurons (single-neuron timescale, $τ$, e.g., membrane time constant in biological neurons) or recurrent interactions among them (network-mediated timescale). However, the contribution of each mechanism for optimally solving memory-dependent tasks remains poorly understood. Here, we train RNNs to solve $N$-parity and $N$-delayed match-to-sample tasks with increasing memory requirements controlled by $N$ by simultaneously optimizing recurrent weights and $τ$s. We find that for both tasks RNNs develop longer timescales with increasing $N$, but depending on the learning objective, they use different mechanisms. Two distinct curricula define learning objectives: sequential learning of a single-$N$ (single-head) or simultaneous learning of multiple $N$s (multi-head). Single-head networks increase their $τ$ with $N$ and are able to solve tasks for large $N$, but they suffer from catastrophic forgetting. However, multi-head networks, which are explicitly required to hold multiple concurrent memories, keep $τ$ constant and develop longer timescales through recurrent connectivity. Moreover, we show that the multi-head curriculum increases training speed and network stability to ablations and perturbations, and allows RNNs to generalize better to tasks beyond their training regime. This curriculum also significantly improves training GRUs and LSTMs for large-$N$ tasks. Our results suggest that adapting timescales to task requirements via recurrent interactions allows learning more complex objectives and improves the RNN's performance.

Figures (27)

• Figure 1: Schematics of network structure and timescales. a. An outline of the network. A binary sequence is given as input to a leaky RNN, with each neuron's intrinsic timescale being a trainable parameter $\tau$. The illustration shows the $N$-parity task with readout heads for $N=3$ and $N=4$. b. An illustration of the manifestation of different timescales (single-neuron and network-mediated) on the autocorrelation (AC) of a network neuron (see also Fig. \ref{['suppfig:ac-fits']}).
• Figure 2: Schematic description of the tasks and curricula a. An outline of the network and tasks. In both tasks, the network receives a binary input sequence, one bit at each time step. b. In the single-head curriculum, only one read-out head is trained at each curriculum step, while in the multi-head curriculum, a new read-out head is added at each step without removing the older heads.
• Figure 3: Training performance depends on the curriculum. a. Accuracy of training the networks ($N$-parity task) without a curriculum increases slowly, especially when $N>10$. For each $N$, 5 models are independently trained for 50 epochs or until reaching $>98$% accuracy. b. Multi-head (dashed) trained networks are solving larger $N$s than single-head (solid) within the same training time (colors in c). c. The maximum trained $N$ for each task/curriculum at the end of training (1000 epochs or solving $N=101$, whichever comes first). Gray lines - mean value across 4 networks.
• Figure 4: Importance of single-neuron timescales for different curricula. a. The maximum $N$ solved in the $N$-bit parity task after 1000 epochs (reaching an accuracy of $98\%$). X-label indicates training constraints: $\tau = 1, 2$ or 3 - fixed $\tau$ with only weights being trained, "Trained" allows training of $\tau$. In the single-head curriculum, models rely on training $\tau$, whereas in the multi-head curriculum, $\tau$ fixed $\tau = 1$ is as good as training $\tau$. Horizontal bars - mean. b. The mean and standard deviation (STD) of the trained $\tau$s increase with $N$ in single-head networks. In contrast, in multi-head networks, the mean $\tau$ decreases towards $1$, and the STD remains largely constant. The mean and STD are computed across neurons within each network (up to the maximum $N$ shared between all trained networks). Shading - variability across 4 trained networks.
• Figure 5: The emergence of network-mediated timescales depends on the curriculum. a. Example average ACs of all the neurons within a single-head network, $N$-parity task. The ACs of individual neurons' activity decay slower with increasing $N$. b. Distributions of the network-mediated timescales $\tau_\textrm{\small{net}}$ for single and multi-head networks solving $N$-parity task for $N = 5$ and $N = 30$. The distribution becomes broader for higher $N$. c, d. The mean and STD of the network-mediated timescale $\tau_\textrm{\small{net}}$ increase with $N$ in both tasks. The mean and STD are computed across neurons within each network. Shades - variability across 4 trained networks.