Traveling Waves Encode The Recent Past and Enhance Sequence Learning
Authors: T. Anderson Keller
Lyle Muller
Terrence Sejnowski
Max Welling
Venue: ICLR 2024 poster (2024)
Keywords: RNNs
Traveling Waves
Memory
Sequence Modeling
Primary area: applications to neuroscience & cognitive science
TL;DR: We demonstrate that traveling waves are an efficient and effective mechanism for storing long term memories in recurrent neural networks.

Abstract: Traveling waves of neural activity have been observed throughout the brain at a diversity of regions and scales; however, their precise computational role is still debated. One physically inspired hypothesis suggests that the cortical sheet may act like a wave-propagating system capable of invertibly storing a short-term memory of sequential stimuli through induced waves traveling across the cortical surface, and indeed many experimental results from neuroscience correlate wave activity with memory tasks. To date, however, the computational implications of this idea have remained hypothetical due to the lack of a simple recurrent neural network architecture capable of exhibiting such waves. In this work, we introduce a model to fill this gap, which we denote the Wave-RNN (wRNN), and demonstrate how such an architecture indeed efficiently encodes the recent past through a suite of synthetic memory tasks where wRNNs learn faster and reach significantly lower error than wave-free counterparts. We further explore the implications of this memory storage system on more complex sequence modeling tasks such as sequential image classification and find that wave-based models not only again outperform comparable wave-free RNNs while using significantly fewer parameters, but additionally perform comparably to more complex gated architectures such as LSTMs and GRUs.

Published as a conference paper at ICLR 2024
TRAVELING WAVES ENCODE THE RECENT PAST AND ENHANCE SEQUENCE LEARNING
Anderson Keller"
Lyle Muller
The Kempner Institute for the Study of Natural and Artificial Intelligence Harvard University, USA
Department of Mathematics Western University, CA
Terrence Sejnowski
Max Welling
Computational Neurobiology Lab Salk Institute for Biological Studies, USA
Amsterdam Machine Learning Lab University of Amsterdam, NL
ABSTRACT
Traveling waves of neural activity have been observed throughout the brain at a diversity of regions and scales; however, their precise computational role is still debated. One physically inspired hypothesis suggests that the cortical sheet may act like wave-propagating system capable of invertibly storing a short-term memory of sequential stimuli through induced waves traveling across the cortical surface, and indeed many experimental results from neuroscience correlate wave activity with memory tasks. To date, however, the computational implications of this idea have remained hypothetical due to the lack a simple recurrent neural network ar- chitecture capable of exhibiting such waves. In this work, we introduce model fill this gap, which we denote the Wave-RNN (wRNN), and demonstrate how such architecture indeed efficiently encodes the recent past through a suite of synthetic memory tasks where wRNNs learn faster and reach significantly lower error than wave-free counterparts. We further explore the implications of this memory storage system more complex sequence modeling tasks such as sequential image clas- sification and find that wave-based models not only again outperform comparable wave-free free RNNs while using significantly fewer parameters, but additionally per- form comparably to more complex gated architectures such as LSTMs and GRUs.
INTRODUCTION
Since the carliest neural recordings (Caton 1875 Beck 1890), neural oscillations and the spatial organization of neural activity have persisted as topics of great interest in the neuroscience community, Consequently, plethora of hypothesized potential functions for these widely observed phenomena have been put forth in the literature. For example: Pitts & McCulloch (947)) propose that alpha oscillations perform 'cortical scanning" akin radar; Milner (1974) suggests that synchrony may serve to segregate the visual scene and 'bind' individual features into cohesive objects; (Engel et al. 2001) suggest brain-wide oscillations serve as 'top-down' stored knowledge and contextual influence for local processing; RRaghavachari et al. (2001) show theta is consistent with gating mechanism for human working memory; (uusáái et al. (2033) posit that oscillations form hierarchical system that offers syntactical structure for spike traffic; [nd[L et al. (2012) & de Mooij-van Malsen et al. (2023) implicate oscillatory coherence with the transfer of information on working memory tasks.
Recently, the advents of high density mlett-eleectrode arrays and high resolution imaging have led to the discovery that many of the oscillations observed in the brain are better described as traveling waves of activity rather than precise zero-lag synchrony (Muller et al. 2018). For example, alpha and theta oscillations have been measured to precisely correspond to traveling waves in both the
cortex and hippocampus of humans (Lubenov & Siapas 2009 Dozano-Soldeville & VanRullen 2019 Zhang et al. 2018). Furthermore, has become increasingly clear that wave-like dynamics are
"Correspondence: T.Anderson. TAndderoonKelleree Keller@ gmail.com, work completed while part of the UvA-Bosch Delta Lab.

---

Published as a conference paper at ICLR 2024
prevalent throughout the brain, from local (Muller et al., 2014) to global (Muller al. 2016 scales, and across virtually all brain regions measured (Townsend et al. 2015 Gu et al. |02 Pang et al.
2023). Research has correlated wave activity with range of functional roles rivaling oscillations
in diversity and number, including: perceptual awareness (Davis et al. 2020); attentional scanning
(Fries
(023); information transfer (Rubino et al. (2066); motor control sequencing and topographic
organization (Takahashi et al. 2011); integrating visual sensory signals with saccadic eye movements
(Zanos et al.
2013); and coordinating and reconfiguring distinct functional regions (Xu et al. 2023).
Most relevant to the present study, significant work has shown correlations between wave dynamics and memory. For example: King Wyart (202T) provide evidence that traveling waves propagating from posterior to anterior cortex serve to encode sequential stimuli in an overlapping 'multiplxedd manner; Sauseng et al. (2002) show that traveling theta oscillations propagate from anterior to posterior regions during retrieval from long-term memory; and Zabeh et al. (2023) show traveling beta waves in frontal and parietal lobes encode memory of recent rewards. As described by "Muller et al. 2), one way to understand the relation between traveling waves and memory comes from an analogy to physical wave systems. Specifically, non-dissipative wave-propagating systems can be shown to contain all information about past disturbances in 'time-reversible' manner. In contrast, in a static wave-free 'bump' system. onset time of stimuli are tied to position, and therefore information about the sequential order of multiple inputs at the same position is lost. This is visualized in Figure I and is demonstrated experimentally in Section 3 Although an imperfect analogy to the brain, such ideas have inspired neuroscientists to suggest that the cortical surface may act like such a wave- propagating system to efficiently encode recent sequential inputs as form of working memory. In this work, we propose a complimentary mechanistic understanding of the relation between traveling waves and memory by suggesting that a wave-propas unve-propggatin hidden state can be thought of like a register or 'stack' to which inputs are sequentially written or 'pushed', The propagation of waves then serves to prevent overwriting of past input, allowing for superior invertible memory storage.
In concert with these theories of potential function, an almost equal number of arguments have been put forth suggesting that traveling waves are "epiphenominal" or merely the inert byproduct of more fundamental causal neural processes without any causal power of their own, One driving factor behind the continuation of this con- troversy is a lack of sufficiently capable models which permit the study of the computational role of waves in task-relevant manner. In the computational neuroscience literature, there exist large-scale spiking neural network (SNN) models which exhibit traveling waves of activity (Davis et al., 2021): 1 however, due to the computational complexity of SNNS, such models are still not feasible to train on real world tasks. These models therefore lack the ability to demonstrate the implications of traveling waves on memory. In the machine learning community, Keller & Welling (2023) recently introduced the Neural Wave Machine which exhibits traveling waves of activity in its hidden state in the service of sequence modeling; however, due to its construction as a network of coupled oscillators, it is impossible disassociate the memory-enhance tributions of oscillations (as described by Rusch & Mishra (202aa) from those of traveling waves.
Input
Channel 0:
Channel 1:
Channel 2:
E-0
t-1
t+2
E+3
t+4
Static Bump System
E-0
E+1
C+2
E+3
E-4
Figure 1: Three binary input signals (top), a corresponding wave-RNN hid- den state (middle), and wave-free static bump system (bottom). At each timestep we are able decode both the onset time and channel of each input from the wave- RNN state. In the wave-free system, rel- ative timing information lost for inputs on the same channel, hindering learning and recall for sequential inputs.
con-
In this work, we derive and introduce a simple non- coniiiiilytt recurrent neural network (RNN) architecture which exhibits traveling waves in its hidden state and subsequently study the performance implications of these waves when compared directly with identical wave-free 'bump system' RNN counterparts. Ultimately, we find that wave-based models are able to solve sequence modeling tasks of significantly greater length, learn faster, and reach lower error than non-wave counterparts, approaching performance of more complicated state of the art architectures - thereby providing some of the first computational evidence for the benefits of wave-based memory systems. We therefore present our model as both a proof of concept of the computational role of wave-based memory systems, well starting point for the continued development of such systems in the future.
2

---

Published as a conference paper at ICLR 2024
2 TRAVELING WAVES IN RECURRENT NEURAL NETWORKS
In this section, we outline how to integrate traveling wave dynamics into a simple recurrent neural network architecture and provide preliminary analysis of the emergent waves.
Simple Recurrent Neural Networks. In order to reduce potential confounders in our analysis of the impact of waves memory, we strive study the simplest possible architecture which exhibits traveling waves in its hidden state. To accomplish this, we start with the simple recurrent neural network (sRNN) defined as follows. For an input sequence {xo}T=o with X E R2 and hidden state ho = 0 & h, E RN, an sRNN is defined with the following recurrence: h = o (h, + Vx, + b) where the input encoder and recurrent connections are both linear, i.. V E NRd and U E RNXX where N is the hidden state dimensionality and a is nonlinearity. The output of the network is then given by another linear map of the final hidden state: y Why, with W R
Wave-Recurrent Neural Networks. To integrate wave dynamics into the sRNN, we start with the simplest equation which encapsulates our goal, the one-dimensional one-way wave equation:
@h(z, t) ah(z,t) Ōr
Ôt
Where t our time coordinate, æ defines the continuous spatial coordinate of our hidden state, and y is the wave velocity. We can see that if we discretize this equation over space and time, with timestep At, defining h(x,t) t) = hī as the activation of the z'th neuron at timestep f, this is equivalent to multiplication of the hidden state vector with the following circulant matrix:
1
v
0
1 1-
v
v
0
he+1= Eh,
where
Σ
0 1 1-1/ v/
where v/ = PÂt. A common linear operator which has similar circulant structure to L convolution. Specifically, assuming single channel length 3 convolutional kernel = [0, v2 v/. we see the following equivalence: u h = Σh,-1, where defines circular convolution over the hidden state dimensions N. Intuitively this can be thought of as a ring of neurons with shared recurrent local connectivity. We therefore propose to define the Wave-RNN (wRNN) with the following recurrence:
h o x h, + Vx, + b)
(3)
In practice we find that increasing the number of channels helps the model to learn significantly faster and reach lower error. To do this, we define u E Rxx where is the number of channels, and is the kernel size, and we reshape the hidden state from single N dimensional circle to separate L dimensional circular channels (e.g. E R For the activation σ, we follow recent work which finds that linear and rectified linear activations have theoretical and empirical benefits for long-sequence modeling (Le et al. 2015 Orvieto et al., (0023). Finally, similar to the prior work with recurrent neural networks (Gu et al. (022), we find careful initialization can be crucial for the model to converge more quickly and reach lower final error. Specifically, we initialize the convolution kernel such that the matrix form of the convolution is exactly that of the shift matrix Σ for each channel separately, with v = 1. Furthermore, we find that initializing the matrix V to be all zero except for single identity mapping from the input to single hidden unit to further drastically improve training speed. Intuitively, these initalizations combined can be seen to support a separate traveling wave of activity in each channel, driven by the input at single source location. Pseudocode detailing these initalizations can be found in Appendix B and an ablation study can be found in Table
Baselines. In order to isolate the effect of traveling waves on model performance, we desire to pick baseline models which are similar to the Wave-RNN as possible while not exhibiting traveling waves in their hidden state. To accomplish this, we rely on the Identity Recurrent Neural Network simple
(iRNN) of et al. 2015). This model is nearly identical to the Wave-RNN, constructed RNN with σ = ReLU, U. but uses an identity initialization for U. Due to this initialization, the iRNN can be seen to be nearly analagous to the bbmmp system' described in Figure [ where activity does not propagate between neurons but rather stays as a static/decaying 'bump'. Despite its simplicity the iRNN is found to be comparable to LSTM networks on standard benchmarks, and thus represents the
3

---

Published as a conference paper at ICLR 2024
Wave RNN (WRNN)
Fiddan Stare (arito)
Activatioon Magnittde
Space Freq. lopoluntt
Normalized Power
Time (steps)
Freq. layo step)
Simple RNN (IRNN)
Activation Algnituse
Space Froo. (cyc/unt)
Normalzed Power
110
Time Istaps)
405
150
Time Freq, layo step)
Figure 2: Visualization of hidden state (left) and associated 2D Fourier transform (right) for a wRNN (top) and iRNN (bottom) after training the SMNIST task. We the wRNN exhibits clear flow of activity across the hidden state (diagonal bands) while the iRNN does not. Similarly, from the 21) space-time fourier transform, diagonal corresponding to the wave propagation velocity of unit/step (Mahmoud et al. 1988).
the wRNN exhibits significantly higher power along the
ideal highly capable simple recurrent neural network which is comparable to the Wave-RNN. We note that, in the original work, we allow all parameters of the matrix U be optimized.
Visualization of Traveling Waves, Before we study the memory capabilities of the wRNN, we first demonstrate that the model does indeed produce traveling waves within its hidden state. To do this, in Figure 2 for the best performing (wRNN & iRNN) models on the Sequential MNIST task of the proceeding section, we plot in the top row the activations of our neurons (vertical axis) over time (horizontal axis) as the RNNs process sequence of inputs (MNIST pixels). As can be seen, there are distinct diagonal bands of activation for the Wave-RNN (left), corresponding to waves of activity propagating between hidden neurons over time. For the baseline simple RNN (iRNN) right, despite sorting the hidden state neurons by 'onset time of maximum activation to uncover any potential traveling wave activity, we for a duration of time and then fade. In the bottom row, following the analysis techniques of Davis al. 2.2T), we plot the corresponding 2D Fourier transform of the above activation time series. In this plot, the vertical axis corresponds to spatial frequencies while the horizontal axis corresponds to temporal frequencies. In such a 2D frequency space, constant speed traveling wave (or general moving object (Mahmoud et al. 1988)) will appear as linear correlation between space and time frequencies where the slope corresponds the speed. Indeed, for the Wave-RNN, we see strong band of energy along the diagonal corresponding to our traveling waves with velocity Ro 1 unit
such bands exist, but instead stationary bumps of activity exist
timestep as
expected, for the iRNN we see no such diagonal band in frequency space. In Appendix C Figure T0
we show additional visualizations of waves propagating at multiple different speeds simultancously within the same network, demonstrating flexibility of learned wave dynamics. Further, in Appendix
Figure [8] we show how the wave dynamics change through training for a variety of different initalizations and architectures. We see that even in wRNN models with random recurrent kernel initalizations, although waves are not present at initialization, they are learned through training in the service of sequence modeling, indicating they are valuable solution for memory tasks.
EXPERIMENTS
In this section we aim to leverage the model introduced in Section 2 to test the computational hypothesis that traveling waves may serve as mechanism to encode the recent past in a wave- field short-term memory. To do this, we first leverage a suite of frequently used synthetic memory tasks designed to precisely measure the ability of sequence models to store information and learn dependencies over variable length timescales. Following this, we use a suite of standard sequence modeling benchmarks to measure if the demonstrated short-term memory benefits of wRNNs persist in more complex regime. For each task we perform a grid search over learning rates, learning rate schedules, and gradient clip magnitudes, presenting the best performing models from each category on a held-out validation set in the figures and tables. In Appendix [B] we include the full ranges of each grid search well exact hyperparameters for the best performing models in each category.
4

---

Published as a conference paper at ICLR 2024
Copy Task. As first analysis of the impact of traveling waves memory encoding, measure the performance of the wRNN on the standard "copy task', as frequently employed in prior work (rrvee et al. 2014 Arjovsky al. 2016; Gu et al. 2020b). The task is constructed of sequences of categorical inputs of length 7 + 20 where the first 10 elements are randomly chosen one-hot vectors representing category in {1, (1....8 88.. The following 7 tokens
category 0, and form the time duration where the network must hold the information in memory. The next token is set to 9, representing a delimiter, signaling the RNN to begin reproducing the stored memory output, and the final 9 tokens are again set to category 0. The target for this task is another categorical sequence of length T + 20 with all elements set to category except for the last 10 elements containing the initial random sequence of the input to be reproduced. At high level, this task tests the ability for a network to encode categorical information and maintain in memory for T timesteps before eventually reproducing it. Given the hypothesis that traveling waves may serve to encode information in an effective 'register'. we hypothesize that wave-RNNs should perform significantly better on this task than the standard RNN. For each sequence length compare wRNNs with 100 hidden units per channel and 6 chamnels (n = 100, c 6) with two baselines: iRNNS of comparable parameter counts (n 100 = 12k params.), and iRNNS with comparable numbers of activations/neeuross (n 625) but a significantly greater parameter count (=> 403k params.). In Figure] we show the performance of the best iRNNS and wRNNs, obtained from our grid search, for T {0, 30, 80). We that the wRNNs achieve more than 5 orders of magnitude lower loss and learn exponentially faster for all sequence lengths. Furthermore, we see that the iRNN with (n 625) still fails to perform even close to the wRNN despite having an equivalent number of activations and nearly 40 times more parameters. From the visualization the model outputs in Figure we see that the iRNN has trouble holding items in memory for longer than 10 timesteps, while the comparable wRNN has no problem copying data for up to 500 timesteps. Ultimately, this experiment demonstrates exactly the distinction between traveling wave and static systems illustrated in Figure - the iRNN (static) system is unable to accurately maintain the relative order of sequence elements that have the same input encoding, while the wRNN wave field has problem encoding both timing and position, facilitating decoding.
T=0
T=30
T=80
10
10-1.
10
10-2
10-1
12'
10
10
10*
10
WRNN 1=000. 0
30
#RNN ni-100, cE
WRNN ni:100 CEO
RNN 025
10-1
IRNN n2:255
10-9
IRN n=:225
10
RVN n00
IRNN 1=000
IRNN n=100
10
100
200
300 100
500
600
100
200
300
400
500
600
100
200
300
400
500
600
Stops (x100)
(X100)
Steps (4100)
Figure 3: Copy task with lengths T=(0, 30, 80J. wRNNs achieve > 5 orders of magnitude lower loss
than iRNNs with approximately equal number of parameters (n 100) and activations (n 625).
T=0
T=10
T=80
4B0
WRNN
Input
Output
IRNN (n=625)
Input
Outpot
Figure 4: Examples from the copy task for wRNN (n=100, c= and iRNN (n=625). We see the iRNN loses significant accuracy after T=10 while the wRNN remains perfect at T=480 (MSE 10-9) 9).
Adding Task. To bolster our findings from the copy task, we employ the long-sequence addition task originally introduced by Hochreiter & Schmidhuber (1997). The task consists of a two dimensional input sequence of length T, where the first dimension is random sample from (((), 1)), and the second dimension contains only two non-zer zon-zero elements (set to 1) in the first and second halves of the sequence respectively. The target is the sum of the two elements in the first dimension which
5

---

Published as a conference paper at ICLR 2024
correspond to the non-zero indicators in the second dimension. Similar to the copy task, this task allows us to vary the sequence length and measure the limits of each model's ability. The original iRNN paper (Le et al. 2015) demonstrated that standard RNNs without identity initialization struggle to solve sequences with but begins to struggle with sequences of length greater than 400 (a result which reconfirm here). In our experiments depicted in Figure and [ we find that the WRNN not only solves the task much more quickly than the iRNN, but it is also able solve significantly longer sequences than the iRNN (up to 1000 steps). In these experiments we use an iRNN with n 100 hidden units (10.3k parameters) and a wRNN with 100 hidden units and 27 channels (10.29k parameters).
150, while the iRNN is able to perform equally as well as an LSTM,
T-100
T-400
T-1000
10"1
10-1
10-1
IRNN n=-00
JANN 200
10:2
10-2.
IRNN 1=100
10-1
MEE
10-3
10-1
10-3.
10-
10
0
20
30
40
50
60
10
20
30
40
50
10
15
20
25
30
Steps (x100)
Steps (x100)
Steps (x100)
Figure 5: wRNN and iRNN Training curves on the addition task for three different sequence lengths (100, 400, 1000). We see that the wRNN converges significantly faster than the iRNN on all lengths, achieves lower error, and solve tasks which are significantly longer.
Seq. Length (T)
100
200
400
700
1000
Test MSE
x 10
4 x 10
1 x 10 10-4
0.16
0.16
iRNN
Solved Iter
14k
22k
30k
X
X
Test MSE Solved Iter.
4 x 10
2 x 10
4 X 10 1k
8 x
10 6 x
10 5
WRNN
300
1k
3k
2k
Table I: Long sequence addition task for different sequence lengths. The wRNN finds the task solution (defined as MSE < 5 x 10 multiple orders of magnitude quicker and is able to solve much longer tasks than the iRNN. The indicates the model never solved the task after 60k iterations.
Sequential Image Classification. Given the dramatic benefits that traveling waves appear to afford in the synthetic memory-specife tasks, in this section we additionally strive to measure if waves will have any similar benefits for more complex sequence tasks relevant to the machine learning community. One common task for evaluating sequence models is sequential pixel-by- pixel image classification. In this work we specifically experiment with three sequential image tasks: sequential MNIST (sMNIST), permuted scquential MNIST (psMNIST), and noisy sequential CIFAR (nsCIFARIO). The MNSIT tasks are constructed by feeding the 784 pixels of each image of the MNIST dataset at time to the RNN, and attempting to classify the digit from the hidden
SMNIST
PSMNIST
100
90-
Test Accuracy
80
Test Accuracy
80
70
70
60
50-
wRNN n-256, c-16
wRNN n-256, c-16
40
IRNN n- 256
iRNN n-256
40
0
20
40
60 Epochs
80
100
120
50
100
150
200
250
Epochs
Figure 6: sMNIST (left) and PSMNIST (right) training curves for the iRNN & wRNN. The wRNN trains much faster and is virtually unaffected by the sequence permutation, while the iRNN suffers.

---

Published as a conference paper at ICLR 2024
state after the final timestep. The permuted variant applies random fixed permutation to the order of the pixels before training, thereby increasing the task difficulty by preventing the model from leveraging statistical correlations between nearby pixels. The nsCIFARIO task is constructed by feeding each row of the image (32 x 3 pixels) flattened as vector input to the network at each timestep. This presents significantly higher iiput-dimensionaliyy than the MNIST tasks, and additionally contains more complicated sequence dependencies due to the more complex images. To further increase the difficulty of the task, the sequence length is padded from the original length (32), to a length of 1000 with random noise. Therefore, the task of the model is not only to integrate the information from the original 32 sequence elements, but additionally ignore the remaining noise elements. As in the synthetic tasks, we again perform grid search over learning rates, learning rate schedules, and gradient clip magnitudes. Because of our significant tuning efforts, we find that our baseline iRNN results are significantly higher than those presented in the original work (98.5% vs. 97% on SMNIST, 91% vs. 81% on PsMNIST), and additionally sometimes higher than many "state of the art' methods published after the original iRNN. In the tables below we indicate results from the original work by a citation next to the model name, and lightly shade the rows of our results.
show results in comparison with existing work on the sMNIST and PSMNIST. Despite the simplicity of our proposed approach, we see that it performs favorably with many carefully crafted RNN and convolutional architectures. We additionally include wRNN + MLP, which is the same as the existing WRNN, but replaces the output map W with 2-layer MLP. We see this increases performance significantly, suggesting the linear decoder of the basic wRNN may be performance bottleneck. In Figure [6] (left), we plot the training accuracy of the best performing WRNN compared with the best performing iRNN over training iterations on the sMNIST dataset. We see that while the iRNN reaches slightly higher final accuracy (+0.9%), the wRNN trains remarkably faster the beginning of training. taking the iRNN roughly 50 epochs to catch up. On the right of the figure, we plot the models' performance on the permuted variant of the task (psMNIST) and see the performance of the Wave-RNN is virtually unaffected, while the simple RNN baseline suffers dramatically. Intuitively, this performance difference on the permuted task may be seen to come from the fact that by directly encoding the input into the wave-field 'buffer', the wRNN is able to learn to classify sequences invari- ant of the ordering permutations of the input (through the fully-connected readout matrix W), while the iRNN has no such buffer and thus struggles to encode sequences with less temporal structure.
In Table 2]
Model
sMNIST
PSMNIST
n/ #0
uRNN
95.1
91.4
512 9k
iRNN
98.5
92.5
256 / 68k
LSTM
98.8
92.9
256 / 267k
GRUP2
99.1
94.1
256 / 201k
NWM
98.6
94.8
128 / 50k
IndRNN (61)
99.0
96.0
128 / 83k
Lip. RNN(
99.4
96.3
128 / 34k
coRNN(7
99.3
96.6
128 34k
LEM
99.5
96.6
128 68k
wRNN (16c)
97.6
96.7
256/478
URLSTM
99.2
97.6
1024 / 4.5M
wRNN + MLP
97.5
97.6
256 / 420k
FlexTCN(5
99.6
98.6
-/ 3/5k
Table 2: sMNIST & PsMNIST test accuracy, sorted by psM-
NIST score. Baseline values from: [1] Arjovsky et al. 2016
[2]
Rusch et al. 2022)_ [3] Liet al. (0018). Gu et al. (2020b)
Romero et al. (002)). Tol Erichson et (2021).
M1 Rusch
& Mishra (202 1st Keller & Welling 2023)
We note that in addition to faster training and higher accuracy, the wRNN model additionally exhibits sub- stantially greater parameter efficiency than the iRNN due to its convolutional recurrent connections in place of fully connected layers. To exemplify this, in Figure 7 we show the accuracy (y-axis) of suite of wRNN models plotted as a function of the number of param- eters (X-axis). We see that compared with the iRNN, the wRNN reaches near maximal performance with significantly fewer parameters, and retains perfor- mance gap over the iRNN with increased parameter counts.
PSMNIST Accuracy # Parameters
95
0y-04, c-32)
15w256 c-1d)
freidt co4)
30-
m-33. cod)
Test Accuracy cocuracyyyyo
a5
Ira cill
80
anr330 c-11
75
WRNN
IRNN
70
1r-360
0
10k
20k
30k
40k
50k
60k
70k
Parameters
Figure 7: Num. parameters vs. accuracy for wRNNs & iRNNS on PSMNIST.
7

---

Published as a conference paper at ICLR 2024
Finally, to see if the benefits of the wRNN extend to more complicated images. we explore the noisy sequential CIFARIO task. In Figure we plot the training curves of the best performing models on this dataset, and see that the Wave-RN RNN still maintains a significant advantage over the iRNN in this setting. In Table 3 we see the performance of the wRNN is ahead of standard gated architectures such GRUs and LSTMs, but also ahead of more recent complex gated architectures such as the Gated mmme-tyiettiic RNN (Chang et al. 2019). We believe that these results therefore serve as strong evidence in support of the hypothesis that traveling waves may be valuable inductive bias for encoding the recent past and thereby facilitate long-sequence learning.
nsCIFAR10
Model
Acc.
units / params
LSTM
11.6
128 / 116k
50-
GRU
43.8
128 / 88k
Test Accuracy
40-
anti-sym. RNN [2]
48.3
256 36k
iRNN
51.3
529 336k
Incremental RNN [3]
54.5
128 / 12k
30
Gated anti-sym. RNN
54.7
256 / 37k
wRNN (16c)
55.0
256 435k
WRNN n n-256, c-16
Lipschits RNN [4]
57.4
128 / 46k
IRNN n=9
CORNN [5]
59.0
128 46k
50
100
Eoochs
150
200
250
LEM [1]
60.5
128 116k
Figure 8: Training curves noisy se- quential CIFARIO. We again see the wRNN trains faster and reaches higher accuracy than the wave-free model.
Table 3: Test set accuracy the noisy sequential CI-
FAR dataset sorted by performance. Baseline values from:
[1] Rusch et al. (2022). [2] Chang et al. (2019). [3] Kag et al.
2020 [2] Brichson et al. (2021) [5 Rusch & Mishra (021a))
Ablation Experiments. Finally, we include ablation experiments to validate the architecture choices for the Wave-RNN. For each of the results reported below, we again grid search over learning rates, activation functions, initalizations, and gradient clipping values. In Table we show the performance of the wRNN the copy task initialization, and V initialization (as described in Section 2). At a high level, we see that the wRNN as proposed performs best, with u-shiff initialization having the biggest impact performance, allowing the model to successfully solve tasks greater than length 7 = 10. In addition to ablating the WRNN, we additionally explore initializing the iRNN with a shift initialization (U Σ) and sparse identity initialization for V to disassociate these effects from the effect of the convolution operation. We see that the addition of Σ initialization the iRNN improves its performance dramatically, but it never reaches the same level of performance of the wRNN indicating that the sparsity and tied weights of the convolution operation are critical to memory storage and retrieval on this task.
ablate various proposed components such as convolution, u- u-shift
Sequence Length (T) 10
Model
80
wRNN
9 X 10 12
1 1 10 10
8 x 10 -1
1 x 8 -11
V-init
1 x 11
2 10 -1
4 x 10 10
4 x 10 11
u-shift-init
5 x 10-1T
3 x 10
7 10-
6 10
V-iini u-shift-init
8 x 10
1 10
3 x 10
7 10
IRNN (n=100)
1 1x 10
3 10
2 10-3
1x 10-3
Σ-init
1 10-8
1 10-7
2 10-7
2 10-5
+ E-init + V-init
1 10-7
1x 10-7
1 10-6
8 x 10-6
Table Ablation test results (MSE) on the copy task, Best results are bold, second best underlined.
4 DISCUSSION
In this work we have discussed theories from neuroscience relating traveling waves to working memory, and have provided one of the first points of computational evidence in support of these
8

---

Published as a conference paper at ICLR 2024
theories through experiments with our Wave-RNN - a novel minimal recurrent neural network capable of exhibiting traveling waves in its hidden state. In the discussion below we include a brief summary of related work, limitations, and future work, with an extended survey of related work in Appendix A
Related Work. As mentioned in the introduction, most related to this work in practice, Keller & Welling (2023) showed that an RNN parameterized as a network of locally coupled oscillators (NWM) similarly exhibits traveling waves its hidden state, and that these waves served as a bias towards learning structured representations. This present paper can be seen as reduction of the
Keller& & Welling (2023) model to its essential elements, allowing for controlled study of hypotheses relating traveling waves and memory. Additionally, we see in practice that this reduction sometimes improves performance, as on the PSMNIST task in Table 2 where the wRNN outperforms the NWM.
Another related line of work proposes to solve long-sequence memory tasks by taking inspiration for biological 'time-cells' (Jacques et al. [021] 2022). The authors use temporal convolution with set of exponentially scaled fixed kernels to extract set of time-scaleiinver invariant features at each layer. Compared with the wRNN, the primary conceptual difference is the logarithmic scaling of time features which is purported to increase efficiency for extremely long time dependencies. As we demonstrate in Appendix Fig. I0 however, the wRNN is also capable of learning waves of multiple speeds, and therefore with proper initialization may also exhibit logarithmic scaling of time-scales.
Similar in results to our work, Chen et al. (2022) demonstrated that in predictive RNN autoencoder learning sequences, Toeplitz connectivity emerges spontaneously, replicating multiple canonical neuroscientific measurements such as one-shot learning and place cell phase precession. Our results in Figure 18 further support these findings that with proper training and connectivity constraints, recurrent neural networks can learn to exhibit traveling wave activity in service of solving task.
Finally, we note that the long- long-sequence modeling goal of this work can be seen as similar to that of the recently popularized Structured State Space Models (S4, HIPPO, LSSL) (Gu et al.] [022] 2020a 202 In practice the recurrent connectivity matrix of the wRNN is quite different than the HiPPO initialization pioneered by these works, however we do believe that linear recurrence may be an equivalently useful property for wRNNs. Importantly, however, the H3 mechanism of Fu et al. (2023) does indeed include a 'shift SSM' with toeplitz structure very similar to that of the wRNN. However, this shift initialization was subsequently dropped in following works such as Heyena and Mamba (Poli et al.) 2023) Gu & Dao 2023). In future work we intend to explore this middle ground between the wRNN and state space models as we find it to likely be fruitful.
I Limitations& & Future Work. The experiments and results in this paper are limited by the relatively small scale of the models studied. For example, nearly all models in this work rely on linear encoders and decoders, and consist of single RNN layer. Compared with state of the art models (such as S4) consisting of more complex deep RNNs with skip connections and regularization, our work is therefore potentially leaving significant performance on the table. However, as described above, beginning with small scale experiments on standard architectures yields alternative benefits including more accurate hyperparameter tuning (due to smaller search space) and potentially greater generality of conclusions. As the primary goal of this paper was to demonstrate the computational advantage of traveling waves over wave-free counterparts both necessary and beneficial. In future work, we plan to test the full implications of our results for the machine learning community and integrate the core concepts from the wRNN into more modern sequence learning algorithms, such as those used for language modeling.
memory tasks, the tradeoff for small scale was
We note that the parameter count for the wRNN on the CIFARIO task is significantly higher than the other models listed in the table. This is primarily due to the linear encoder mapping from the high dimensionality of the input (96) to the large hidden state. In fact, for this model, the encoder V alone accounts for > 90% of the parameters of the full model (393k/435k). If one were to replace this encoder with a more parameter efficient encoder, such as a convolutional neural network or a sparse matrix (inspired by the initialization for V), the model would thus have significantly fewer parameters, making it again comparable to state of the art. We leave this addition to future work, but believe it to be one of the most promising approaches to improving the wRNN's general competitiveness.
Finally, we believe this work opens the door for significant future work in the domains of theoretical and computational neuroscience. For example, direct comparison of the wave properties of our model with neurological recordings (such as those from (King & Wyarl (02T)) may provide novel insights into the mechanisms and role of traveling waves in the brain, analogous to how comparison

---

Published as a conference paper at ICLR 2024
of visual stream recordings with convolutional neural network activations has yielded insights into the biological visual system (Cadieu et al. 2014).
ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
We would like to thank the creators of Weight & Biases (Biewald 2020 and Py PyTorch (Paszke et al. 2019). Without these tools our work would not have been possible. We thank the Bosch Center for Artificial Intelligence for funding T. Anderson Keller for the initial stages of this project, and The Kempner Institute for funding him through the final stages of the project.
REFERENCES
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48. ICML' 16, pp. 1120-1128 JMLR.org. 2016.
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past, 2016.
Adolf Beck. Die bestimmung der gehirn-und cecken-markfuucciiones vermittelst der electrischen ersheimungen. bl. Physiol, 4:473-476, 1890.
Gabriel Benigno, Roberto Budzinski, Zachary Davis, John Reynolds, and Lyle Muller. Waves traveling over a map of visual space can ignite short-term predictions of sensory input. Research Square Platform LLC, August 2022. doi: 00.21203/r..3.ss19003 144vl. URL https://dor
orq/10. 21203 3. rs 903114/v
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL /twtps://ww... wandb com/ Software available from wandb.com.
György Buzsáki, Nikos Logothetis, and Wolf Singer. Scaling brain size, keeping timing: Evolutionary
preservation of brain rhythms. Neuron, 80(3):751-764, 2013. ISSN 0896-6273. doi: https://doi.
0gg/100.1016/,..errn..00000 10.002. URL httttttt//www..et.. http:///wwwwwwiieeeeeeee.... com/ science/ Aaticleep11/S00922231111
Charles F. Cadieu, Ha Hong, Daniel L. K. Yamins, Nicolas Pinto, Diego Ardila, Ethan A. Solomon, Najib J. Majaj, and James J. DiCarlo. Deep neural networks rival the representation of primate it cortex for core visual object recognition. PLOS Computational Biology, 10(12):1-18, 12 2014. doi: 10. (0.133// pcbi. 1003963. URL /ttp::://6....rg/11.111111 1003963
pcbi
Richard Caton. The electric currents of the brain. British Medical Journal, 2:278, 1875.
Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. AntisymmetricRNN: A dynamical system view recurrent neural networks. In International Conference on Learning Representations, 2019. http:ttts:
net for fttt:/oooommm--dyeedeeeeee
Yusi Chen, Huanqiu Zhang, and Terrence J. Sej.owsk........ Hippocampus as generative circuit for pre- dictive coding future sequences. bioRxiv, 2022. doi: 10. 111111////2222...999.2.222222000 URL https
www. wwwrxiv..iii.............i..iwv.
05. 4927311
Zachary W. Davis, Lyle Muller, Julio Martinez-Trujillo, Terrence Sejnowski, and John H. Reynolds. Spontaneous travelling cortical waves gate perception in behaving primates. Nature, 587(7834): 432 436, October 2020. doi: 10. 10.1038/s41586-000 2802-y.
URL lattps://coi.og///
1038/441586-020-1
Zachary W. Davis, Gabriel B. Benigno, Charlee Fletterman, Theo Desbordes, Christopher Steward, Terrence Sejnowski, John H. Reynolds, and Lyle Muller. Spontaneous traveling waves naturally emerge from horizontal fiber time delays and travel through locally waancchrooou--irrggl... states. Nature Communications, 12(1), October 2021. doi: 10. 167108//4444--021-2220
URL https:
10

---

Published as a conference paper at ICLR 2024
Johanne Gertrude de Mooij-van Malsen, Niels Röhrdanz, Anna-Sophia Buschhoff, Thomas Schif- felholz, Torfi Sigurdsson. and Peer Wulff. Task-specific oscillatory synchronization of pre- frontal cortex, nucleus reuniens, and hippocampus during working memory. iScience, 26(9): 107532, 2023. ISSN 2589-0042. doi:
httts
www. www.sccience.c.... rect com htt.a//cccccc.c.....o.aaaaoaaaaa
Andreas Engel, Pascal Fries, and Wolf Singer. Dynamic predictions: Oscillations and synchrony in top-down processing, Nature Reviews Neuroscience, 2(10):704-715 716, October 2001. doi: 10. 10.1038/35094555.
N. Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W. Mahoney. Lipschitz recurrent neural networks. In International Conference Learning Representations, 2021.
Pascal Fries, Rhythmic attentional scanning. Neuron, 111(7)954--70, 2023, ISSN 0896-6273, doi:
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinsska Sergio Gómez Colmenarejo, Edward Grefenstette. Tiago Ramalho, John Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski. Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis. Hybrid computing using neural network with dynamic external memory. Nature, 538(7626): 471-476, October 2016. doi: 10. 10100
URL https:://// http://://ooo.oo//...11..2
nature20101
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Re. Hippo: Recurrent memory with optimal polynomial projections, 2020a.
Albert Gu, Caglar Gulcehre, Tom Le Paine, Matt Hoffman, and Razvan Pascanu. Improving the gating mechanism of recurrent neural networks, 2020b.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state-space layers, 2021a.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022. URL https //openreview.o hetp:/iwwwwweeeeeeeeeeeeeeeeeeee wttp::/eoooummidd-eeeoooooo//
Yameng Gu, Lueas E Sainburg, Sizhe Kuang, Feng Han, Jack W Williams, Yikang Liu, Nanyin Zhang, Xiang Zhang, David A Leopold, and Xiao Liu. Brain activity fluctuations propagate as traversing the cortical hierarchy. Cerebral cortex, 11,,,,110 3109)3988-4405, 2021b.
Mikacl Henaff, Arthur Szlam, and Yann LeCun. Recurrent orthogonal networks and long- memory tasks. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2034-2042, New York, New York, USA, 20-22 Jun 2016. PMLR. URL mlr www...peessvviei http//////hh/afffeeeeeeame html
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
11

---

Published as a conference paper at ICLR 2024
Brandon Jacques. Zoran Tiganj, Mare Howard, and Per B Sederberg. Deepsith: Effi- cient learning via decomposition of what and when across time scales. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 27530-27541. Curran Associates, Inc., neur www.eurrrp...pppppppppeppppppppp cc/paper firr_fllee fflles/paper/22111 Paper pdf
2021. URL http:::
Brandon G Jacques, Zoran Tiganj, Aakash Sarkar, Marc Howard, and Per Sederberg. A deep convolutional neural network that is invariant to time rescaling. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedines of Machine Learning Research, pp. 9729-9738. PMLR, 17-23 Jul 2022. URL httttpss://rrcceeddngg html
mlr. wtw.ppees/v12///ccuueeees
Anil Kag, Ziming Zhang, and Venkatesh Saligrama. Rnns incrementally evolving on equi- librium manifold: A panacea for vanishing and exploding gradients? In International Conference Learning Representations, 2020. httlhtttss//pppneeeet httpenrrwwwwweeeeeeeeeeeeeeeeeee forum?
id (d-HylpqA4Fwa
Anderson Keller and Max Welling, Predictive coding with topographic variational autoencoders. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV) Workshops, pp. 1086 1091, October 202
Anderson Keller and Max Welling. Topographic vaes learn equivariant capsules. In Advances in Neural Information Processing Systems, volume 34, pp. 28585-28597, Curran Asso- ciates, Inc., 2021b. URL hsttps: .ew.eeerddiinnen.nip.p...... www.ipuurp.... httplppscccppper/22//22/
Paper..appppppp. pdf
T. Anderson Keller and Max Welling. Neural wave machines: Learning spatiotemporally structured representations with locally coupled oscillatory recurrent neural networks. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023.
T. Anderson Keller, Qinghe Gao, and Max Welling. Modeling category- selective cortical regions with topographic variational autoencoders. In SVRHM 2021 Workshop @ NeurIPS, 2021. URL htttp:://ppenree
net wwwffforrmmmooommmm wttp://eoooommmmGGGFmm
IW54bI
Réan-Rémi King and Valentin Wyart. The human brain encodes chronicle of visual events at each instant of time through the multiplexing of traveling waves. Journal of Neuroscience, 41
(34):7224-7233 2021. ISSN 0270-6474. doi: 10.152//NEUROSCI..000 20.2021. URL https: www. jneurosci. .rgg http:/oonenn///3//2220 /ttp://nnnn////////222//
Jonas Kubilius, Martin Schrimpf, Ha Hong, Najib J. Majaj. Rishi Rajalingham, Elias B. Issa, Kohitij Kar, Pouya Bashivan, Jonathan Prescott-Roy, Kailyn Schmidt, Aran Nayebi, Daniel Bear, Daniel L. Yamins, and James J. DiCarlo. Brain-Like Object Recognition with High-Performing Shallow Recurrent ANNs. H. Wallach, H. Larochelle, Beygelzimer, D' Alché- Buce Fox, and R. Garnett (eds.), Neural Information Processing Systems (NeurIPS), pp. 12785 12796. Curran Associates, Inc., 2019.
Quoc V. Le, Navdeep Jaitly, and Geoffrey E. Hinton. A simple way initialize recurrent networks of rectified linear units, 2015.
Hyodong Lee, Eshed Margalit, Kamila M. Jozwik, Michael A. Cohen, Nancy Kanwisher, Daniel L. K. Yamins, and James J. DiCarlo. 'Topographic deep artificial neural networks reproduce the hallmarks of the primate inferior temporal cortex face processing network. bioRxiy, 2020, doi: 10.1 10.1101/2020.0..000 09. 18. 1881
httppp/swwww.biooooooooo//eeeeee 09. 07..1.18111111161
Mario Lezcano-Cais Casado and David Martínez- Rubio. Cheap orthogonal constraints in neural networks:
A simple parametrization of the orthogonal and unitary group, 2019.
12

---

Published as a conference paper at ICLR 2024
S. Li. W. Li, C. Cook, C. Zhu, and Y. Gao. Independently recurrent neural network (indran): Building a longer and deeper rnn In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5457-5466, Los Alamitos, CA, USA, jun 2018. IEEE Computer URL hhttp:://doo. eeecomuutersoc sessociicty iety
Society. doi: 10.1109/CVPR.20 201. 0.0722..
otg/10.1109//c 1109/CV1 201.PP........ 005722
Stefanic Liebe, Gregor M Hoerzer, Nikos K Logothetis, and Gregor Rainer. Theta coupling between v4 and prefrontal cortex predicts visual short-term memory performance. Nature neuroscience, 15 (3):4564462, 2012.
Diego
and Rufin VanRullen. The hidden spatial dimension of alpha: 10-hz perceptual echoes propagate as periodic traveling waves in the human brain. Cell reports, 26(2): 374-380, 2019.
Evgueniy Lubenov and Athanassios G. Siapas. Hippocampal theta oscillations are travelling waves. Nature, 459(7246):534-539, May 2009. doi: 10.10388mature00000.. URL https 100o.ooorg/...1033//aaaeee111111
S.A. Mahmoud, M.S. Afifi, and R.J. Green. Recognition and velocity computation of large moving objects images. IEEE Transactions on Acoustics, Speech, and Signal Processing, 36(11):1190 1791, 1988.
doi: 10.1109/29.900 9020.
Peter Milner. model for visual shape recognition. Psychological Review (16):52) 1974. doi: 10.1037/h0037149 URL http::://bb..or//11...000
Lyle Muller, Alexandre Reynaud, Frédéric Chavane, and Alain Destexhe. The stimulus-evoked popu- lation response in visual cortex of awake monkey is propagating wave. Nature Communications, 5(1), April 2014.
doi: 10. 10.10088/ccmm URL htttps://doi.orgg////0
ncomms4677
Lyle Muller, Giovanni Piantoni, Dominik Koller, Sydney S Cash, Eric Halgren, and Terrence J Sejnowski. Rotating waves during human sleep spindles organize global patterns of activity that repeat precisely through the night. eLife. 5:e17267, nov 2016. ISSN 2050-084X. doi: 7554/eLife 17.
10.7554/eLife. 17267 URL http:////bbb.orr........//..eee/
Lyle Muller, Frédéric Chavane, John Reynolds, and Terrence Sejnowski. Cortical travelling waves: mechanisms and computational principles. Nature Reviews Neuroscience, 19(5):255 268, March 2018. doi: 10.1038/mm 2018.............. URL
2018.20
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023.
James C. Pang, Kevin M. Aquino, Marianne Oldehinkel, Peter A. Robinson, Ben D. Fulcher, Michael Breakspear, and Alex Fornito. Geometric constraints on human brain function. Nature, 618 (7965):566-577) 574, May 2023. doi: 10. 10.1038/641586-023-000000 1038/ 02
Il Memming Park, Ábel Ságodi, and Piotr Aleksander Sokół. Persistent learning signals and working memory without continuous attractors, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. 2019,
Walter Pitts and Warren S McCulloch. How know universals the perception of auditory and visual forms. The Bulletin mathematical biophysics, 9:127-147, 1947.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao. Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023.
13

---

Published as a conference paper at ICLR 2024
Sridhar Raghavachari, Michael J Kahana, Daniel S Rizzuto, Jeremy B Caplan, Matthew P Kirschen, Blaise Bourgeois, Joseph R Madsen, and John Lisman. Gating of human theta oscillations by a working memory task. Journal of Neuroscience, 21(9):3175-313 2001.
David W. Romero, Robert-Jan Bruintjes, Jakub Mikolaj Tomezak, Erik Bekkers, Mark Hoogendoorn, and Jan van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. In International Conference on Learning Representations, 2022. URL htttp::/
net/ /tt///oormmi=i333 fooo
Doug Rubino, Kay A Robbins, and Nicholas G Hatsopoulos. Propagating waves mediate information transfer in the motor cortex. Nature Neuroscience, 912):1549-1557, November 2006. doi: 10.1038/nn 1802.
T. Konstantin Rusch and Siddhartha Mishra. Coupled oscillatory recurrent neural network (cornn): An accurate and (gradient) stable architecture for learning long time dependencies. In International Conference on Learning Representations, 2021a.
Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies, 2021b.
05487
Konstantin Rusch, Siddhartha Mishra, N Benjamin Erichson, and Michael W Mahoney. Long ex- pressive memory for sequence modeling. In International Conference on Learning Representations, 2022.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. One-shot learning with mugmet-a neural networks, 2016.
Paul Sauseng, Wolfgang Klimesch, W Gruber, Michael Doppelmayr, Waltraud Stadler, and Manuel Schabus. The interplay between theta and alpha oscillations in the human electroenceehallgram reflects the transfer of information between memory systems. Neuroscience letters, 321 124, 2002.
Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai kin Wong, and Wang chun Woo. Convolutional 1stm network: machine learning approach for precipitation nowcasting. 2015.
Kazutaka Takahashi, Maryam Saleh, Richard Penn, and Nicholas Hatsopoulos. Propagating waves in human motor cortex. Frontiers in Human Neuroscience, 5, 2011. ISSN 1662-5161.
doi: 10.3389/fnhum.201 201 1.00040. URL httpt::///// h/wp::::/wwww..// wtt.://froriieseooro/o/rrrrreee.
10. 3389/3339855 ffng/m....mmm 2011. 00040
Corentin Tallec and Yann Ollivier: Can recurrent neural networks warp time?, 2018.
Rory G Townsend, Selina S Solomon, Spencer C Chen, Alexander NJ Pietersen, Paul R Martin, Samuel Solomon, and Pulin Gong. Emergence of complex wave patterns in primate cerebral cortex. Journal of Neuroscience, 35(1 1):4557-4622. 2015.
James C.R. Whittington, Timothy H. Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy E.J. Behrens. The tolman-eichenbaum machine: Unifying space and relational memory through gencralization in the hippocampal formation. Cell, 183(5):1249- 1263.e23, 2020. ISSN 0092-8674. doi: http:///oi.ong//00.06//..c.....p 10.024. URL https w/www. www.sccienceiicccmm rectt. com/ cttp://ccciece.rrr.aeaaaa//////a
Yiben Xu, Xian Long, Jianfeng Feng, and Pulin Gong. Interacting spiral wave patterns underlie complex brain dynamics and are related to cognitive processing. Nature Human Behaviour, 7 (7):1196 1215, June 2023. doi: 10.1038/s41562 021-0-066--5.55
URL htttt://///o.... /ttp:///oo..../////////o
1038 008//4156222 023 626-5
Erfan Zabeh, Nicholas C. Foley, Joshua Jacobs, and Jacqueline P. Gottlieb. Beta traveling waves in monkey frontal and parietal areas encode recent reward history. Nature Communications, 14(1), September 2023. doi: 10. 40.10338/44447--23444111
URL ht https: http:::://oo.ooooeeeee /ttp:::///ol..o0/01.1.13...
241467-023-4112--
14

---

Published as a conference paper at ICLR 2024
Theodoros P Zanos, Patrick J Mineault, Konstantinos T Nasiotis, Daniel Guitton, and Christopher C Pack. A sensorimotor role for traveling waves in primate visual cortex. Neuron, 85(3):615-677, 2015.
Honghui Zhang, Andrew J. Watrous, Ansh Patel, and Joshua Jacobs. Theta and alpha oscillations traveling waves in the human neocortex. Neuron, 98()): 169-11281 .e4, June 2018. doi: 10.1016/j. ccuron.0018.05.019 URL https://o//..ooooo otg.ooooorgggg/ 1011/101111/ neuroo. 2018.
Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms, 2016.
15

---

Published as a conference paper at ICLR 2024
Appendices
Table of Contents
Related Work
16
B Experiment Details
Additional Results
21
A RELATED WORK
Deep neural network architectures that exhibit some brain-like properties related to ours: CORnet
(Kubilius et al. 2019), a shallow convolutional network with added layer-wise recurrence shown to be better match primate neural responses; models with topographic organization, such the TVAE
(Keller& & Welling 2021 Keller et al. 2021 Keller & Welling] 202 and TDANN (Lee et al. 2020D; and models of hippocampal-corte the TEM (Whittintoon et al., (020). Our work is unique in this space in that it is specifically focused on generating spatio-temporally synchronous activity, unlike prior work. Furthermore, we believe that our findings and approach may be complimentary to existing models, increasing their ability to model neural dynamics by inclusion of the Wave-RNN fundamentals such as locally recurrent connections and shift initalizations.
interactions such as the PredRAE (Chen et al. 2022 and
In the machine learning literature, there are a number of works which have experimented with local connectivity in recurrent neural networks. Some of the earliest examples include Neural GPUs (Eukasz Kaiser& & Sutskever, 2016) and Convolutional LSTM Networks (Shi et al, 2015). These works found that using convolutional recurrent connections could be beneficial for learning algorithms and spatial sequence modeling respectively. Unlike these works however, the present paper explicitly focuses the emergence of wave-like dynamics in the hidden state, and further studies how these dynamics impact computation. To accomplish this, we also focus on simple recurrent neural networks as opposed to the gated architectures of prior work providing a less obfuscated signal to the computational role of wave-like dynamics.
One line of work that is highly related in terms of application is the suite of models developed to increase the ability of recurrent neural networks to learn long time dependencies. This includes models such as Unitary RNNs (Arjovsky et al.] (0166). Orthogonal RNNs (Henaff' et al. 2016).
expRNNs (Lezcano-Casado & Martinez-Rubio, 2019). the chronoL-TM Tallec & Ollivier 2018),
anti. mymierymmet RNNs (Chang et al. 2019)) Lipschitz RNNs (Erichson et al. 2021), coRNNs (Rusch
& Mishra 2021a). unicoRNNs (Rusch & Mishra 2021b). LEMss (Rusch et 2022), periodic &
quasi-periodic continous attractors (Park ct al. 20033). Recurrent Linear Units (Orvieto et al., 2023),
and Structured State Space Models (S4) (Gu al. (022). Additional models with external memory may also be considered in this category such as Neural Turing Machines (Graves et al. 2014), the DNC (Graves al. 2016), memory augmented neural networks (Santoro et 2016)) and Fast-weight RNNs (Ba et al.) 2016). Although we leverage many of the benchmarks and synthetic tasks from these works in order to test our model, we note that our work is not intended to compete with state of the art on the tasks and thus we do not compare directly with all of the above models. Instead, the present papers intends to perform a rigorous empirical study of the computational implications of traveling waves in RNNs. To best perform this analysis, we find it most beneficial to compare directly with as similar of a model as possible which does not exhibit traveling waves, namely the iRNN. We do highlight, however, that despite being distinct from aforementioned models algorithmically, and arguably significantly simpler in terms of concept and implementation, the wRNN achieves highly competitive results. Finally, distinct from much of this prior work (except for the coRNN (Rusch & Mishra 2021a) and DeepSITH JeepSITTJJcqques Jacques et al. (2021)). ), our work uniquely leverages neuroscientific memory problem, thereby potentially offering insights into neuroscience observations in return. Of the above mentioned models, our model is perhaps most intimately related in theory to the periodic and quasi-periodic attractor networks of Park et al. (2023)
inspiration solve the
16

---

Published as a conference paper at ICLR 2024
and therefore encourage readers to consult that work for greater theoretical understanding of the wave-RNN.
In final related paper. Benigno al. (2022) studied complex-valued recurrent neural network with distance-dependee connectivity and signal-pus propagation time-delay, in order to understand the potential computational role for traveling waves that have been observed in visual cortex. They showed that when recurrent strengths are set optimally, the network is able to perform long-term closed-loop video forecasting significantly better than networks lacking this spatially-organized recurrence. Our model is complimentary to this approach, focusing instead on the sequence integration capabilities of waves, rather than on forecasting, and leveraging a more traditional deep learning architecture.
B EXPERIMENT DETAILS
In this section, include all experiment details including the grid search ranges and the best performing parameters. All code for reproducing the results can be found at the following repository: htt.gghhhb..ommaaanyyy..ece.
RNNs, The code was based on the original code base from the coRNN paper (Rusch & Mishra 2021a) found at http://iitubb.coom/ttuucchcchhee
Pseudocode. Below we include an example implementation of the wRNN cell in Pytorch (Paszke
et al. 2019):
import torch. as
class
wRNN Cell (nn. Module)
def
_init_ (self, ninn, n, k=3): super (RNN_Cell) self) init_()) ()
self.n
self.c
self.V Lineir (n n * c)
self.U nn. nn.Convld (c, C, k, 1, k//2, padding_mode- ci circular') gelf.act = nn. ReIJ))
12
Sparse identity initialization for nn zendt..erro_ssseeet nn.init. zeros zerooe with torch corch.no_grad(): self. www.ww...t...m vieeigh..iiwwe,, n, nin))
13
16
17
w(:, 0] - 1.0
18
Shift initialization for U
wts zeros (c, c, ii.in (Wtss (wts)
21
22
wts = torch. roll (wts, 1, -1)
23
with torch.no_grant self.U. copy_ (wts)
26
def fcrward x, hy)
27
hy self.act (self. f. (x) view self.e, self.n) + sel f.wy (hy)
return by
Figure The results displayed in Figure are from the best performing models of the sMNIST experiments, precisely the same as those reported in Figure[[[[eft) and Table The hyperparamters of these models are described in the sMNIST section below. To compute the 2D Fourier transform. we follow the procedure of Davis et al. 2021): we compute the real valued magnitude of the 2-dimensional Fourier transform of the hidden state activations over time (using torch. fft.. fft2 seq) abs () ). To account for the significant autocorrelation in the data, we normalize the output by the power spectrum of the spatially and temporally shuffled sequence of activations. We note that although this makes the diagonal bands more prominent, they are still clearly visible in the un-normalized spectrum. Finally, we plot the logarithm of the power for clarity.
Copy Task. We construct each batch for the copy task as follows:
17

---

Published as a conference paper at ICLR 2024
def get copy.ppppppp. task batch (osz, T)
np. zeros (bs2, T+10))
3
data np. rap.random.roo randint (low=1, high-9, size- (bs2, 10)
x[:, :10] data
x[:, -[10+1]
y ne.zeros (bsz, T+10))
y[:, -10:] x|:, :10]
X = F.one_het (torch. tensor (x, dtype=torch. 6nt permute (1, 0), 10)
Y torch.tensor (Y) intt () permute (1,
12
return x,
For the copy task, we train each model with batch size of 128 for 60,000 batches. The models are trained with cross entropy loss, and optimized with the Adam optimizer. We grid search over the following hyperparameters for each model type:
iRRN
Gradient Clip Magnitude: [0, 1.0, 10.0]
U Initialization: [I, U v )
Learning Rate: [0.01, 0.001, 0.0001]
Activation: [ReLU, Tanh]
wRNN
Gradient Clip Magnitude: [0, 1.0, 10.0] Learning Rate [0.01, 0.001, 0.0001]
We find the following hyperparameters then resulted in the lowest test MSE for each model:
Sequence Length (T)
Model
Parameter
0
10
30
80
480
wRNN
Learning Rate Gradient Magnitude Clip
le-3
le-3
le-3
le-3
le-4
(n=100, c=6, k=3)
0
0
Learning Rate
1e-3
1e-3
le-4
le-4
le-4
IRNN
Gradient Magnitude Clip
10
1
1
(n=100)
U-initialization
I
I
I ReLU
Activation
ReLU
ReLU
ReLU
ReLU
Learning Rate
1e-3
le-4
le-4
le-4
1e4
iRNN (n=625)
Gradient Magnitude Clip
10
1
U-initialization Activation
u
I
I
I ReLU
ReLU
ReLU
ReLU
ReLU
Table Best performing hyperparamters for each model on the Copy Task. Gradient clipping means not applied, U initialization 2I means Kaming Uniform Initialization.
The total training time for these sweeps was roughly 1,900 GPU hours, with models being trained on individual NVIDIA 1080Ti GPUs. The iRNN models took between to 10 hours to train depending on and n, while the wRNN models took between to 15 hours to train.
MNIIS. For the sequential MNIST task we use iRNNs and wRNNs with 256 hidden units. To have a similar number parameters, we use 16 channels with the wRNN. The models are trained with a batch size of 128 for 120 epochs. The learning rate schedule is defined such that the learning rate is divided by a factor of Ir drop rate every Ir drop_epoch epochs. We grid search over the following hyperparameters for each model type. We find that the iRNN does not need gradient clipping this task and achieves smooth loss curves without it. We highight in yellow the hyperparameters which achieve the maximal performance, and were thus reported in the main text:
18

---

Published as a conference paper at ICLR 2024
iRRN
Learning Rate: [0.001, 0.0001, 0.00001]
"__dro__rate [3.33, 10.0]
Ir Ir_drop_epoch: [40, 100]
wRNN
Gradient Clip Magnitude: [0, 1. 10, 100]
Learning Rate: [0.001, 0.0001, 0.00001]
Ir droopprrtee: [3.33, 10.0]
Ir ir_drop_epoch: [40, [00]
The total training time for these sweeps was roughly 1000 GPU hours, with models being trained on individual NVIDIA 1080Ti GPUs, iRNN models taking roughly 12 hours each, and wRNN models taking roughly 18 hours each.
PSMNIST. For the permuted sequential MNIST task, we use the same architecture and training setup for the sMNIST task. We find that this task the iRNN requires gradient clipping to perform well and thus include it in the search as follows:
iRRNN
Gradient Clip Magnitude: [0, 1, 10, 100, 1000]
- Learning Rate: [0.001, 0.0001, 0.00001]
Ir rrtp [3.33, 10.0]
Ir_drop_epoch: [40, 100]
wRNN
Gradient Clip Magnitude: [0, 1, 10, 100, 1000]
Learning Rate: [0.001, 0,0001, 0.00001]
drop_rrrree pate: [3.33, [0.0]
Ir arop_epoch: [40, 100]
In an effort to improve the baseline iRNN model performance, we performed additional hyperpa- rameter searching. Specifically, tested with larger batch sizes (120, 320, 512). different numbers of hidden units (64, 144, 256, 529, 1024), additional learning rates (1c-6, 5c-6), a larger number of epochs (250), and more complex learning rate schedules (exponential, cosine, one-cycle, and reduction on validation plateau). Ultimately we found the parameters highlighted above to achieve the best performance, with the only improvement coming from training for 250 instead of 120 epochs. Regardless, in Figure 6 we see the iRNN performance is still significantly below the wRNN performance, strengthening the confidence in our result. The total compute time for these sweeps was roughly 1,900 GPU hours with models being trained on individual NVIDIA 1080 Ti GPUs. The iRNN models took rougly 12 hours each, with wRNN models taking roughly 18 hours each.
For each of the wRNN models in Figure 7 the same hyperparameters are used as highlighted above and found to perform well. The wRNN is tested over combinations of = (16, 36, 64, 144) x (1, 4, 16, 32), and the best performing models for each parameter count range are displayed. For the iRNN, we sweep over larger batch sizes (128, 512), learning rates (Ie-3, le-4, le-5) and gradient clipping magnitudes (0, 1, 10, 100) in order to stabalize training, displaying the best models.
CICIARI,.. For the noisy sequential CIFARIO task, models were trained with a batch size 256 for 250 epochs with the Adam optimizer, Ir rr_drop_epoor gradient clipping was not necessary for the wRNN model on this task, and thus perform a grid search follows, with the best performing settings highlighted:
100, and lr ddop rate 10. We found
iRNN
Learning Rate: [0.001, 0.00010 0.00001]
-
Number hidden units (n): [144, 256, 529, 1024]
19

---

Published as a conference paper at ICLR 2024
WWRNN
Learning Rate: [0.001, 0.0001, 0.00001]
Number hidden units (n): [144, 256
The total compute time for these sweeps was roughly 1,600 GPU hours with models being trained on individual NVIDIA 1080 Ti GPUs. The iRNN models took roughly 15 hours each, with wRNN models taking roughly 22 hours each.
Adding Task. For the adding task we again train the model with batch sizes of 128 for 60,000 batches using the Adam optimizer. For this task models are trained with mean squared error loss. For both the iRNN and wRNN, we then grid-search over the following hyperparameters:
Gradient Clip Magnitude: [0, 1, 10, 100, 1000]
Learning Rate: [0.01, 0.001, 0.0001]
In Table we report the best performing parameters for each task setting. The total training time for these sweeps was roughly 1 1,200 GPU hours. The iRNN models took between to 7 hours, while the wRNN models took 12 hours each
Sequence Length (T)
Model
Parameter
100
200
400
700
1000
wRNN
Learning Rate Gradient Magnitude Clip
le-3
le-4
le-4
le-4
le-4
(n=100, c=27, k=3)
100
100
1
100
10
iRNN
Learning Rate
1e-3
1e-3
le-3
le-4
le-3
(n=100)
Gradient Magnitude Clip
1000
100
10
100
Table 6: Best performing settings on the Adding Task. Gradient clipping 0 means not applied. We note that the iRNN failed to solve the task meaningfully for lengths T = 700 & 1000, thus the hyperparameters found here are not significantly better than any other combination for those settings.
Ablation. For the ablation results in Table we first report the test MSE of the best performing wRNN and iRNN models from the original Copy Task grid search (identical to those in Figure 3). We then added the ablation settings to the grid search, and trained the models identically to those reported above the Copy Task section. This resulted the final complete grid search:
iRNN
Gradient Clip Magnitude: [0, 1.0, 10.0]
U Initialization: [1, u - )
V Initialization: [A[(0, 0.001). Sparse-Identity)
Learning Rate: [0.01, 0.001, 0.0001]
Activation: [ReLU, Tanh]
wRNN
Gradient Clip Magnitude: [0, 1.0, 10.0]
Initialization: [Dirac, ud- -
रα +a). u--uisiitt u-siitt]
V Initialization: (O, 0.001), Sparse-Identity)
Learning Rate [0.01, 0.001, 0.0001]
where Sparse-Identity refers to the V initialization described in section 2 In the case of the iRNN, the matrix is defined to have channel for the purpose of this initialization. The Dirac initialization is equivalent an identity initialization for convolutional layer and is implemented using the Pytorch function: torch. nn. init. dirac_. We report the best performing models from this search Table
4
For the wRNN (-u-shift-init), we note that for the sequence lengths where the wRNN model does solve the task (MSE 10-10), i.e. T=0 & 10, the best performing models always the
20

---

Published as a conference paper at ICLR 2024
initialization (00 0.001) and appear to learn to exhibit traveling waves in their hidden state, while the dirac initialization always performs worse and does not exhibit waves.
C ADDITIONAL RESULTS
Performance Means & Standard Deviations. In the main text, in order to make a fair comparison with prior work, we follow standard practice and present the test performance of each model with the corresponding best validation performance. In addition however, distributional properties of the model performance after multiple random initalizations. In Table 7 we include the means and standard deviations of performance from 3 reruns of each of the models.
find it beneficial to report the
Task
Metric
iRNN
wRNN
MSE
1.40x10-5 x.21x10-6
6.12x10-5 + 7.95x10
Adding T=100
Solved iter.
11,500.00 2,910.33
233.33 + 57.74
MSE
1.1x10-9+ 11,000.00
8.54x10-5 7.91x10-
Adding T=200
Solved iter.
1,000.00 = 0.00
MSE
7.70x10- 8.49x00 8.49x10-2
1.59x10-4 1,07x00
Adding T=400
Solved iter.
30,000.00 -
1,333.33 577.35
MSE Solved iter.
0.163 + 2.08 10-3
5.29x + 3.19x10
Adding T=700
3,000.00 = 0.00
MSE
0.160 +
4.36x 1.91
Adding T=1000
Solved iter.
1,666.67 577.35
SMNIST
Test Acc.
98.20 0.32
97.30 = 0.34
PSMNIST
Test Acc.
90.85 + 1.47
96.60 + 0.10
IsCIFAR1O IO
Test Acc.
51.80 +: 0.54
54.70 + 0.42
Table 7: Mean and standard deviation of model performance over random initializations for the best performing models in each category. We model performance is consistent with the best performing models reported in the main text. The x means the models never solved the task after 60,000 iterations, and (#: -) means that the other 2/3 random initalizations also did not solve the task after 60,000 iterations or crashed.
Visualizing Multiple Channels. To visualize how activations might flow between channels, we include visualization of the channels for the wRNN on the sMNIST task in Figure
Figure 9: Visualization of 15 separate hidden state channels for wRNN trained the sMNIST task. We see wave activity differs between channels has the potential to flow between channels although it is difficult to precisely quantify due to their lack of predefined ordering.
21

---

Published as a conference paper at ICLR 2024
Waves of Multiple Velocities. To compliment visualization of waves in the main text, show in Figure I0 multiple distinct channels of a single wRNN trained on sMNIST. For such a network, randomly initlize the velocity of each channel.
deeseeee Steoe
150
150
250
500
200
Space
:50
Thee Deguere (te03
100
50
800
Space (Htiden Unit0)
6.5
200
7500 Delyeeee Stpr1
250
500
620
700
Space IHibden (N415)
6.6
Figure 10: Depiction of traveling waves of multiple speeds within the same model. Each row shows a separate channel of the hidden state for a single model. The speed of the waves are given by the slopes of the lines of high activation.
Unsorted Visualization of iRNN Hidden State. To ensure that our sorting procedure of iRNN neurons in Figure did did not artificially destroy localized activity which may have otherwise been visible, in the following Figure [T] we plot the same of neural activations, with neurons unsorted for the iRNN. We see that there is less structure as expected, and still no wave-like activity.
Sorted Hidden State (units)
Simple RNN (IRNN)
Activation Magnitude
Space Freq, (cyc/unit)
Comallzzd pee Power
100
208
30
258
Time Freq, (oyo step)
Time (steps)
Figure 11: Visualization of hidden state (left) and associated 2D Fourier transform (right) for an iRNN after training the sMNIST task. This identical to Figure without the sorting of iRNN neurons by onset time.
22

---

Published as a conference paper at ICLR 2024
Fourier Analysis Calibration To demonstrate the relationship between the speed of traveling waves and the resulting slope of peak power in the 2D Fourier transform, in Figure 12 we compute the 2D Fourier transform of synthetic data where have induced known velocities. As expected, we see the slope matches the observed wave velocities. We note that we only plot the magnitude of positive frequency components in this work since all signals real values and thus the remaining components are simply the complex conjugates of these components.
Velocity 0.5 unit / step
206
Tex: (Saguone Steo)
Velocity unit / step
200
Time (Seguimer Stepo
Velocity units step
100
Time (Seouimee Stepe)
Figure 12: Examples of synthetic waves with known velocities (left) and the resulting 2D Fourier Transforms (right). We see indeed the slope of the peak magnitude in Fourier space corresponds to the wave velocities.
Additional Ablation: Locally Connected RNN. In order to test if the emergence and maintenance of waves requires the weight sharing of the convolution operation, or is simply due to local connectiv- ity, we perform an additional experiment where we use the exact same model as the default wRNN, however we remove weight sharing across locations of the hidden state. This amounts to replacing the convolutional layer with an untied 'locally connected' layer. In practice, when initialized with same the ahift-initialization we find that such model does indeed exhibit traveling waves its hidden state depicted in Figure 1 Although notice that the locally connected network does train to comparable accuracy with the convolutional wRNN on sMNIST, we present this preliminary result as simple ablation study and leave further tuning of the performance of this model on sequence tasks
for future work.
Hildan State zartto)
Iteration 10,000
Activation Magnitude
Iteration 600,000
Aetivatin Seggritus
8
Time (stspa)
Time 1steps)
Figure 13: Hidden state of locally connected RNN showing the existence of traveling waves. These results imply that it is simply local connectivity which is important for the emergence of traveling waves rather than shared weights.
23

---

Published as a conference paper at ICLR 2024
Additional Baseline: Linear Layer MNIST. One intuitive explanation for the performance of the wRNN is that the hidden state "wave-fiele fieee- acts like a register or "tape' where the input is essentially copied by the encoder, and then subsequently processed simultaneously by the decoder at the end of the sequence. To investigate how similar the wRNN is to such a solution, we experiment with training single linear layer on flattened MNIST images, equivalent to what the decoder of the wRNN would process if this "tape" hypothesis were correct. In Figure [4] we plot the results of this experiment (again showing the best model from grid search over learning rates and learning rate schedules), and we see that the fully connected layer achieves maximum performance of 92.4% accuracy compared with the 97.6% accuracy of the wRNN model.
FC vs. WRNN
100
80
Test Accurocy
60
40
wRNN
20
0
20 40
60
80
100 120
Epochs
Figure 14: Training curves for fully connected layer trained flattened MNIST images (FC) versus the standard wRNN presented in the main text. We see the wRNN still performs significantly better, however the FC model does mimic the rapid learning capability of the wRNN, suggesting that the wRNN's learning speed may be partially attributable to the wave-field's memory-tape-lin quality.
Additional Ablation: Frozen Encoder & Recurrent Weights. To further investigate the difference between the wRNN model and model which simply copies input to a register, we propose to study the relative importance of the encoder weights V and recurrent connections U for the wRNN. We hypothesized that the wRNN may preserve greater information in its hidden state by default, and thus may not need learn flexible encoding, or perform recurrent processing. To test this, we froze the encoder and recurrent connections, leaving only the decoder (from hidden state to class label) to be trained. In Figure 15 we plot the training curves for a wRNN and iRNN with frozen U and V. We see that the wRNN performs remarkably better then the iRNN in this setting (89.0% vs. 44.3%). indicating that the wave dynamics do indeed preserve input information by default far better than standard (identity) recurrent connections.
SMNIST Frozen Encoders and Recurrence V
90
Test Accuracy
70-
50-
40-
30
wRNN frozen
IRNN frozen
20
50
150
250
Epochs
Figure 15: Training curves for wRNN and iRNN models with frozen encoder and recurrent connec- tions (U& & V) on the sMNIST task. We see that the wRNN performs drastically better, indicating that the wRNN requires significantly less ficxibility in its encoder and recurrent connections in order to achieve high accuracy.
24

---

Published as a conference paper at ICLR 2024
Additional Visualizations of Copy Task. Here include additional visualizations of the smaller iRNN (m=100) for the copy task. We see that it performs significantly worse than the larger iRNN (n=625) displayed in the main text, unable to recall sequences for any length of time.
T=0
T=10
T=30
T8480
WRNN
Inout
Output
Input
IRNN
Output
Figure 16: Examples from the copy task for wRNN (n=100, c=6) and iRNN (n=100). We see the iRNN loses significant accuracy after T=10. These results clearly show that the iRNN does not have the appropriate machinery for storing memories over long sequences while the wRNN does.
On the Emergence of Traveling Waves. In this section we expand on the emergence of traveling waves in recurrent neural networks with different connectivity and initialization schemes. Specifically, Figure [8 we show the hidden state visualization for models with varying initalizations. We see that models with shift-initialization exhibit waves directly from initialization, while randomly initialized convolutional models do not initially exhibit waves but learn to exhibit them during training, and identity initialized models never learn to exhibit waves. Furthermore, in Figure 17 we show the respective training curves for randomly initialized and identity initialized wRNNs. We see that the randomly initialized wRNN achieves higher final accuracy in correspondence with the emergence of traveling waves, reinforcing the conclusion that traveling waves are ultimately beneficial for sequence modeling performance.
WRNN u-Initialization Comparison
100
95
90
Test Accuracy
85
75
70
wRNN u- u-shift-init
WRNN rand-init
wRNN identity-init
60
20
40
60
80
100
120
Epochs
Figure 17: Training curve on the sMNIST task for wRNN with three different initialization schemes: u-shift (default), random (u1(- vn Ja)), and identity (dirac). We see that the random initalization does not have the same rapid learning speed as the u-shift initialization, however, it does still achieve significantly higher beneficial to performance.
final accuracy than the identity initialization, implying traveling waves
25

---

Published as a conference paper at ICLR 2024
Iteration 10,000
Iterrtion 600,000
WRNN u-shift
50
155
100
200
300
480
902
600
WRNN Random
139
200
250
225
WRNN Identity
230
702
300
206
200
400
355
666
100
128
350
500
600
X0
IRNN Identity
280
-00
iRNN Random
50
200
228
45
690
101
130
200
300
400
506
900
700
IRNN U-shift
300
200
255
Figure 18: Visualization of hidden state (y-axis is hidden units) over timesteps (X-axis) for a variety of different models and initalizations for U. We see that the wRNN with u-shift initialization achieves the most consistent waves throughout training. Interestingly, some other models learn to achieve traveling waves despite not having them at initalization (wRNN with random (kaming uniform) initialization): while other models, (iRNN with U-shift) initially have stronger traveling waves, and slowly lose them throughout training. We see the wRNN with identity (dirac) initialization never learns waves despite using convolutional recurrent connections.
26

---

Published as a conference paper at ICLR 2024
Waves with Variable Velocity. As described in the related work, there is relation between the ideas presented in this paper and the notion of 'time cells' from the neuroscience literature. One interpretation of time cells is that they could be represented by progresses over the cortical surface. To explore the impact of such modification, we implemented locally-connected wRNN architecture which is initialized such that the velocity systematically decreases for each t'th row of the 2 matrix as In such a model the wave speed could indeed be interpreted to decrease it progresses. In Figure T9 we plot the hidden state of such model after being trained on the sMNIST task. We see that waves indeed travel significantly more slowly, and appear to have curved trajectories, implying 'slowing'. In terms of performance, we see that the model reaches roughly the same test accuracy the default wRNN presented in the main text (%67.76), suggesting that such a modification could performant alternative to study in the future.
which slows down as it
Figure 19: Visualization of the hidden state of a locally connected wRNN model, trained on sMNIST; initialized with velocities that decrease down as propagates. We see that the waves do indeed appear to have velocities which change as function of position in the hidden state, and the model still learns to perform well on the task.
function of position, emulating wave which slows
Impact of Hidden State Size. Due to cyclic boundary conditions, the wRNN should in theory
be able to handle arbitrary time delays with constant computational complexity (i.e. waves simply propagate in circles until decoded). Wave-RNNs should therefore be able solve sequence tasks with number of steps significantly greater than the number of hidden units. For example, we have found that models with as little as 49 hidden units (6 channels) can solve the copy task with a delay of 1000 time steps. Despite this theory, practice, we have found models with larger hidden states still typically perform better on longer sequence tasks. We demonstrate this in Table T We believe this divergence from theory be primarily due to optimization challenges. our experiments we found runtime to correlate significantly sub-linearly with hidden state size, with the majority of computational complexity arising from iterating over sequence length.
Seq. Length (T)
50
100
200
500
hid 9
8c-3
le-2
le-2
2c-4
n hid 16
Ic-3
2e-2
6e-3
4c-4
n_hid 25
le-4
2e-3
2e-3
le-5
n_ 1_hid= 49
1e-8
8e-5
9e-6
2e-4
hid 100
3e-9
7c-7
4e-7
1e-6
Table 8: Test loss on the copy task for the wRNN for different Hidden State sizes (n (n_ vs. Sequence Lengths. We see the larger hidden state sizes almost consistently achieve lower enror.
27


Reviews:
[{"review_id": "5bwSrmILsA", "reviewer": "Reviewer_h1zv", "summary": "In this paper, the authors explore how patterns of wavelike activity, observed in brains, might help artificial neural networks learn and recall sequences of inputs. They discretize the 1-dimensional wave equation, finding a common structure between the discretization matrix and convolutions that they exploit to set up a simple RNN that naturally supports wavelike activity patterns.  After demonstrating that the network indeed exhibits waves, they test how such waves may facilitate sequence learning and recall in a number of different tasks by comparing the network to a very similar one initialized in a manner that does not naturally support waves.  Overall, the authors’ RNN performed impressively well across a number of sequence tasks, including a more complex permuted MNIST sequence task.", "questions": "Please see **Weaknesses**", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors overall provide a very good job in providing motivation, relevant biological background, and mathematical intuition for their network setup.  The results of their simple setup are impressive, and are generally carefully presented and analyzed, including with ablation studies.  Moreover, the authors provide a fuller characterization of the wave activity and performance of their network in the Appendix, including, transparently, their network result distributions, which is generally far too lacking in the field. Similarly, they provide their code in an anonymous repo and specify the parameters of the network, increasing the reproducibility profile of the work.", "weaknesses": "Note, the below concerns have resulted in a lower score, which I would be happy to increase pending the authors’ responses. \n\n**A. Wave fields**\n\nThe wave-field comparisons, claims, and references seem a bit strained and unnecessary.  Presumably, by “wave-field,” the authors simply mean a vector field that supports wave solutions.  In any case, since this term is not oft-used in neuroscience or ML that I am aware of, a brief definition should be provided if the term is kept.  However, I am unsure that it is necessary or helpful.  That the brain supports wavelike activity is well-established, and some evidence for this is appropriately outlined by the authors.  Many computational neuroscience models support waves in a way that has been mathematically analyzed (e.g., Wilson-Cowan neural fields equations).  The authors’ discretization methodology suggests a similar connection to such analyses.  However, appealing to “physical wave fields” to relate waves and memory seems to be overly speculative and unnecessary for the simple system under study in this manuscript.  The brain is a dissipative rather than a conservative system, so that many aspects of physical wave fields may well not apply.  Moreover, the single reference the authors do make to the concept does not apply either to the brain or to their wave-RNN.  Instead, Perrard et al. 2016 describe a specific study that demonstrates that a particle-pilot wave system can still maintain memory in a very specific way that does not at all clearly apply to brains or the authors’ RNN, despite that study studying a dissipative (and chaotic) system. Instead, the readers would benefit much more from gaining an intuition as to why such wavelike activity might benefit learning and recalling sequential inputs.  Unfortunately, Fig. 1 does little to help in this vein. \n\nHowever, the concept certainly is simple enough, and the authors provide a few intuitions in the manuscript that help.  I believe the manuscript would improve by removing the discussion of wave fields and instead providing / moving the intuitive explanations (e.g., the “register or ‘tape’” description on p. 20) as to how waves may help with sequential tasks to the same portion of the Introduction.  \n\n**B. Fourier analysis**\n\nOverall, I found the wave and Fourier analysis a bit inconsistent and potentially problematic.  While I agree that the wRNNs clearly display waves when plotted directly, the mapping and analysis within the spatiotemporal Fourier domain (FD below) does not always match patterns in the regular spatiotemporal plots (RSP below).  Moreover, it’s unclear how much substance they add to the analysis results.  In more detail: \n\n1. Constant-velocity, 1-D waves don’t need to be transformed to the FD to infer their speeds.  The slopes in the RSP correspond to their speeds.  For example, in Fig. 2 (top left), there is a wave that begins at unit 250, step ~400, that continues through to unit 0, step ~650, corresponding to a wave speed of ~1.7 units/step, far larger than the diagonal peak shown in the FD below it that would correspond to a speed of ~0.3 units/step, as indicated by the authors.  \n\n2. Similar, seemingly speed mismatches can be observed in the Appendix.  E.g., in Fig. 9 (2nd column, top), the slopes of the waves are around 0.35-0.42 units/step (close enough to likely be considered the same speed, especially as they converge in time to form a more clustered wave pulse) from what I can tell, whereas the slopes in the FD below it are ~0.3 for the diagonal (perhaps this is close enough to my rough estimate) and ~0.9, well above any observable wave speed. Perhaps there is a much faster wave that is unobservable in the RSP due to the min/max values set for the image intensity in the plot, but in that case the authors should demonstrate this.  Given (a) the potential mismatch in the speeds for the waves that can be observed, (b) the mismatch in the speeds discussed above in Fig. 2, and (c) the fact that some waves may be missed in FD (see below), I would worry about assuming this without checking.\n\n3. As alluded to in the point above, iRNN in Fig. 2 appears to have some fast pulse bursts easily observed in the RSP that don’t show in the FD. For example, there is a very fast wave observable in the RSP in units ~175-180, time steps 0-350.  Note, the resolution is poor, but zooming in and scrolling to where the wave begins around unit 175, step 0 makes it clear.  If one scrolls vertically such that the bottom of the wave at step 0 is just barely unobservable, then one can see the wave rapidly come into view and continue downwards.  Similarly some short-lasting, slower pulses in units near 190, steps 0-350 are observable in the RSP.  None of these appear in the FD.  Note, this would not take away from the claim that wRNNs facilitate wave activity much more than iRNNs do, but rather that some small amounts—likely insufficient amounts for facilitating sequence learning—of wave activity might still arise in iRNNs.  If the authors believe these wavelike activities are aberrant, it would be helpful for them to explain why so.\n\n4. I looked over the original algorithm the authors used (in Section III of “Recognition and Velocity Computation of Large Moving Objects in Images”—RVC paper below—which I would recommend for the authors to cite), and I wonder if an error in the initial calibration steps (steps 1 & 2) occurred that might explain the speed disparities observed between the RSPs and FDs.\n\n5. There do seem to be some different wave speeds—e.g., in Fig. 9, there appear to be fast and narrow excitatory waves overlapping with slow and broad inhibitory waves. But given that each channel has its own wave speed parameter $\\nu$, it isn’t clear why a single channel would support multiple wave speeds.  This should be explored in greater depth, and if obvious examples of sufficiently different speeds of excitatory waves are known (putatively Fig. 9, 2nd column), these should be clearly shown and carefully described and analyzed.\n\n6. Is there cross-talk across the channels?  If so, have the authors examined the images of the hidden units (with dimensions __hidden units__ x __channels__) for evidence of cross-channel waves?  If so, perhaps this is one reason for multiple wave speeds to exist per channel?\n\n7. Overall, it is unclear overall what FT adds to the detection of 1-D waves.  If there are such waves, we should be able to observe them directly in the RSPs.  In skimming over the RVC paper, it seems like it would be most useful in determining velocities of 2-D objects and perhaps wave pulses.  That suggests that one place the FD analysis might be useful is if there are cross-channel waves as I mention above.  If so, the waves should still be observable in the images (and I would encourage such images be shown), but might be more easily characterized following the marginalization decomposition procedure described in the original algorithm in Section III of the RVC paper.  Note, the FDs might also facilitate the detection of multiple wave speeds in the network, as potentially shown in Fig. 9.  However, in that case it would seem they should only appear in Fig. 9, and if the speeds are otherwise verified.\n\n8. The authors mention they re-sorted the iRNN units to look for otherwise hidden waves.  This seems highly problematic.  If there are waves, then re-sorting can destroy them, and if there is only random activity then re-sorting can cause them to look like waves.  \n\n**C. Mechanisms**\nFinally, while the results are overall impressive, and hypotheses made regarding the underlying mechanisms for the performance levels of the network, there is too little analysis of the these mechanisms.  While the ablation study is important and helpful, much more could be done to characterize the relationship between wavelike activity and network performance.\n\n**D. Minor**\n1. Fig. 2: Both plots on the right have the leftmost y-axis digits obscured\n2. Fig. 9, top, plots appear to have their x- and y- labels transposed (or else the lower FD plots and those in Fig. 2 have theirs transposed.\n3. Fig. 15 needs axis labels", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GcFjXpmNdP", "reviewer": "Reviewer_ovPq", "summary": "This paper presents Wave-RNN, in which the hidden states are organized in time to resemble a traveling wave.  Because the position of the wave enables one to reconstruct the time of the event that triggered it, the network can maintain precise information about what happened when in the past.   A one-layer wRNN is evaluated on some artificial tasks (e.g., psMNIST) to evaluate its ability to solve problems with long-range dependencies.  It is compared to other models, especially Identity RNN which the authors argue provides a fair comparison because it has long memory but does not exhibit waves.", "questions": "Suppose that the velocity v in Eq. 2 changed systematically as a function of the row of the matrix.  In particular, what if v went down like 1/x?  How would that wave behave?  How would this change how the longest time horizon of the memory scales with number of weights/units?  Would this model behave better or worse on the problems in this paper than wRNN does?   Presumably worse because the temporal resolution at long delays would be poor and problems like psMNIST require precise timing information.  Could those problems be mitigated with a deep network?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This model can remember not only retain information about input to the network but also the time at which it was experienced as long as the waves persist.  Information about time is very useful.", "weaknesses": "The connection between this computational model and traveling waves in the cortex is extremely tenuous.  The problem is that the goal of this model is to allow information to be remembered for a long time whereas the traveling waves over the cortex (and hippocampus and striatum) are much faster.  For instance, the Siapas & Lubenov paper shows that theta travels the length of the hippocampus in on the order of 200 ms.  \n\nThere is some evidence for very slow oscillations in MEC, but there is no evidence these are traveling waves.\nhttps://doi.org/10.1101/2022.05.02.490273 \nhttps://doi.org/10.1016/j.celrep.2023.113271\n\nOn the other hand, there is extremely robust evidence for reliable sequences of firing in the brain over time scales relevant for memory.   \nhttps://doi.org/10.1038/s41467-018-07020-4\nhttps://doi.org/10.1016/j.cub.2021.01.032 \nThis phenomenon is often (but not always) referred to as ``time cells.'' There is not evidence that these sequence are anatomically organized, but that doesn't seem to be important for the model.  \n\nThere is a critical difference between the sequences observed with time cells and the traveling waves here.  In particular, the sequences in the brain slow down as they unfold.  This is as if the wave was traveling through a fluid whose properties change systematically from one end of a channel to the other.  This does not seem to be a property of this model.  Although Figure 9 in Appendix C shows waves that travel at different speeds in different channels, the waves within each channel proceed at a constant velocity (the right panel is more complicated but certainly doesn't slow down the way the neural data do).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yoCPtrhnqK", "reviewer": "Reviewer_2PWQ", "summary": "The present study uses a neurally inspired convolutional recurrent network model with traveling wave states to accomplish a set of tasks. The authors claim the traveling wave recurrent network outperforms the one without waves.", "questions": "- Do you need to retrain the RNN to copy sequences with different lengths, or produce waves with different speeds?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The present study is a good and novel example of brain-inspired computation in that it uses the wave RNN (inspired by the brain) to implement a set of benchmark tasks. It also considers a set of experiments that support the wave RNN indeed outperforms its counterparts without traveling waves. It provides us insight that structured spatiotemporal dynamics can have its advantages in real applications.", "weaknesses": "### Major\n- Although the concept of this study is novel, I feel the present paper can be strengthened by conducting a deeper analysis to show why internally generated traveling waves in RNNs can improve those computational tasks. For example, since the network model is small (I see the limitation discussed in the end), the authors could perform a dimensionality analysis to demonstrate the network's evolution in the low-dimensional manifold. I do see the author providing an intuitive explanation of the benefit of traveling waves in the introduction, and showing neurons' response in Fig. 2, but they are not enough from my point of view. \n\n- The illustrative example in Fig. 1 (middle) is probably not a perfectly matching example in explaining the benefit of traveling waves. The Fig. 1 illustrates an example of a two-way wave equation, however, the traveling wave in the RNN is only a one-way wave equation that the network state moves in a single direction. I am not clear about why the traveling wave could maintain the spatiotemporal information.\n\nI'd like to raise my rating if the two concerns can be properly addressed.\n\n### Minor\n- Eq. 2: it seems that the matrix $\\Sigma$ misses the time step $\\Delta t$ in discretizing Eq. 1. Otherwise the speed $v$ cannot exceed 1 (the term $1-v$ in $\\Sigma$). I mean you can absorb the $Delta t$ into a new parameter in Eq. 2 without explicitly expressing it, but in this case the $v$ in Eq. 2 is not the same $v$ in Eq. 1.\n\n- Below Eq. 2: I am confused about what the input channel means.\n\n- Below Eq. 3: it is not clear how the convolution between $u$ with dimension [c,c,f]  and $h$ with dimension [c,n'] is calculated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5JLpCcxrLv", "reviewer": "Reviewer_6SWP", "summary": "Over the past decade, recurrent neural networks, trained by backpropagation through time, have been used to infer mechanisms employed by networks of biological neurons to perform cognitive tasks. Yet, most existing literature on biological RNNs focus on attaining fixed or slow points at steady-state, and neglects the other well-known solution -- oscillations. On the other hand, traveling waves have been observed in the brain and studied for multiple decades. Waves in general are well-understood mechanistically, but their computational role in the brain remains a hot debate. Here, the authors bridge this very obvious gap by introducing an RNN architecture that produces traveling waves. Wave-RNNs store waves in multiple rings (which they call \"channels\") and seemingly perform reasonably well across a variety of sequence-based tasks.\n\nI am supportive of the general idea and the rigor of this work. However, the model introduction was very hard to parse and requires several rereads and returning to after reading future sections in order to reconcile any confusion.", "questions": "(1) Did the authors really initialize $\\mathbf{V}$ to be zero (as claimed on page 3)? I think it is a typo and they actually meant initializing $\\mathbf{b}$ to be zero instead?\n\n(2) There seems to be no additional point in training wRNN + MLP? It provides a small improvement and a single sentence explaining that linear decoding is a bottleneck (which is in fact an important bottleneck to prevent overfitting in neural data) -- am I missing something?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This work is timely and important for the biologically-plausible RNN community to acknowledge the possibility of oscillatory/wave-like solutions that can arise from gradient-based training. The work is rigorous and well-motivated.", "weaknesses": "While the general narrative makes sense, I find the introductory narrative very hard to understand. \n\n(1) Since $\\mathbf{u}$ has dimensions $c \\times c \\times f$, this means that there are convolutional kernels between rings, but the intuition provided does not seem to acknowledge this.\n\n(2) The usage of identity-initialized RNNs was also not clear. From the parameter count of such models, it seems like every element of the weight matrix is being optimized. That is not mentioned anywhere in the text (or if it is, it should be highlighted more clearly). Indeed, that is the case for Le et al 2015, but this would detract from the bump model which the authors originially intended to be the main baseline model to compare with wave-RNNs.\n\n(3) The way wave-RNNs are named is confusing and inefficient. At some times, it is labeled as $n = 100, c=6$ to represent $n$ as the total number of neurons and $c$ as the number of channels. This means that there are 16 neurons per channel, and where does the last 4 neurons go (how does the floor function in page 3 work)? At other times it is called $16c$ which I assume refers to the same thing. In both ways of naming, there is no mention about the dimensionality of the kernel -- I suspect it may not be constant because of the remaining 4 neurons, but there has to be a better way to label everything.\n\n(4) In Figure 2, neurons in the iRNN are sorted according by time of maximum activation. How are the neurons in the wRNN sorted and where is the channel separation?\n\n(5) The training curves in Figures 3,5,6 are extremely transient and subject to random \"resets\" in performance. This means that in a single training session, the efficiency of training depends (by luck) on number of resets that happen, which is very apparent in Figure 3. This makes any conclusions drawn about training efficiency unconvincing. More models should be trained and the loss curves averaged.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)