Traveling Waves Encode The Recent Past and Enhance Sequence Learning

Summary

We thank the reviewer for their acknowledgment of the biological relevance of our work and for the contribution of a model which exhibits wave-like solutions from gradient-based training.

Below we address the reviewer’s proposed weaknesses and questions.

(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.

Indeed the convolutional kernels are fully parameterized and therefore support waves between channels. We believe the intuition provided in our introduction still supports this. Specifically, since the mapping between channels and the cortical surface is not firmly fixed, it would fit within our model that channels may be laid out parallel on the cortical surface, and such cross-channel communication could then again simply be implemented by local lateral connectivity. We appreciate the reviewer bringing this up, and will strive to make the multi-channel nature of the convolution more clear.

(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 originally intended to be the main baseline model to compare with wave-RNNs.

In both settings all parameters are being optimized as this is necessary for both models to achieve optimal performance. We appreciate the reviewer highlighting that this is not clear, and we have modified the text to reflect this. As we found and demonstrated empirically in Figure 2, despite this full parameterization and optimization, the iRNN still learns a solution which resembles a static ‘bump’ solution, and therefore we believe it to still fit within our motivation.

(3) The way wave-RNNs are named is confusing and inefficient. At some times, it is labeled as $n = 100, c=6$ to represent 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.

We thank the reviewer for highlighting potential confusion with respect to the total number of neurons and our notation. In our plots, $n$ refers to the number of neurons per channel, so the total number of neurons is $n \times c$. We have modified the text to make this more clear.

(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?

There is no sorting of the wRNN neurons, they inherently have spatial locality due to the convolution operation. We therefore simply plot the hidden state vector in the order of the array over which the convolution is applied.

(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.

While we agree that the dramatic increases in training loss are less than ideal, we would like to make two notes which may at least temporarily ease the reviewers concerns. First, we note that the y-axis is on a logarithmic scale, therefore although it may look like the loss is 'resetting’, in practice these changes in performance are not as significant as they may first appear. Secondly, we note that all of the believably harmful 'resets’ which happen occur for the wRNN model, which still learns orders of magnitude faster than the iRNN.

Although we unfortunately do not have time to rerun all of the loss curves and average them for this rebuttal we would certainly be happy to do so for the camera ready if accepted, and we can assure the reviewer that for the reruns we did during development, we saw the same results consistently — the wRNN trained orders of magnitude faster than the iRNN.

Summary

We thank Reviewer 2PWQ for their time spent with our paper, for acknowledgement of the novelty of our submission, and for their helpful comments which have helped make the presentation of our work more clear. In the following we explain how we have attempted to address their major (and minor) concerns in our updated submission.

Major Point 1: Deeper Understanding

We appreciate the reviewer noting how the clarity of our presentation could be improved. In response to this, we have updated the intuition provided in our introduction to give a better understanding of the working mechanism of our model which we believe may partially address the reviewer's concern regarding a lack of analysis. Specifically, as pointed our by reviewer h1zv, it is potentially easiest to understand the success of the wRNN model as a having a hidden state in the form of a 'register' or 'stack' where information is sequentially written or pushed. The propagation of waves then serves to prevent overwriting of past input, allowing for superior invertible memory storage.

Pursuant to this, we greatly appreciate the reviewer's suggestion to perform low dimensional analysis of the hidden state dynamics, however we believe that with an intuitive understanding of the dynamics as described above this may not be necessary. Specifically, since we can interpret the hidden state as a set of different `stacks' (one for each channel) where information largely flows within each individual channel separately, the low-dimensional manifold over which the hidden state will evolve is explicitly biased towards these 1-D circular topologies that we have plotted in Figure 2, and in the appendix. We encourage the reviewer to look at Figure 9 of the updated appendix which includes visualization of all channels, as this reinforces the understanding that in general inter-channel communication is relatively minor, and the subspaces of interest for hidden state evolution are exaclty those plotted, as this is the intended design of the network.

Major Point 2: Figure 1

We thank the reviewer for highlighting the potential confusion caused by this figure. Indeed in response to these concerns and those of Reviewer h1zv, we have updated Figure 1 to better match the 1-dimensional structure of the wRNN hidden state. Additionally, we believe this new figure better reflects the new intuition provided in the introduction, suggesting to interpret the hidden state as a 'register' or 'stack' where information is written/pushed. We believe this may help to better convey the understanding of the model that the reviewer felt was lacking and thereby additionally contribute to addressing the reviewer's first major point.

Minor Point 1: Missing $\Delta t$

We thank the reviewer for highlighting the skipped steps in our derivation of the wRNN. Indeed we had assumed that the timestep was absorbed into $\nu$ in equation 2. The full derivation is as follows: $h_{t+1}(x) = h_t(x) + \Delta t \left[\nu \frac{\partial h(x, t)}{\partial x}\right]$ $h_{t+1}(x) \approx h_t(x) + \Delta t \nu \left[h_t(x+1) - h_t(x)\right]$ $h_{t+1}(x) = (1 - \Delta t \nu) h_t(x) + \Delta t \nu h_t(x+1)$ In the updated submission we have changed the velocity in equation 2 to $\nu'$ and noted that this incorporates the timestep.

Minor Point 2: Input Channels

We use the term 'input channels' in the standard sense employed when describing convolutional neural network layers. The only difference is that in our setting the convolution operates on an RNN hidden state vector, rather than an input image or feature map. By channels of the hidden state, we then refer to subsets of the hidden state obtained by reshaping the traditional $n$ dimensions into $c$ separate $n' = \lfloor\frac{n}{c}\rfloor$ dimensional subsets (as described after Equation 3). In traditional convolutional neural networks, these subsets of the input signal are typically called the input channels, and the resulting subsets of the output vector (feature maps) are often called the output channels. This naming extends to the convolutional kernel as well, which has weights for each input-output pair of channels respectively. This then extends into the reviewer's minor point 3 below.

Minor Point 3:

The convolution between the kernel $\mathbf{u}$ with dimensionality [$c_{out} \times c_{in} \times f$] and the input with dimensionality [$c_{in} \times n'$], can be understood as computing $c_{out}$ separate convolutions of shape [$c_{in} \times f$] with signal [$c_{in} \times n'$], where each convolution has a different set of weights -- this is identical to convolution in a standard convolutional neural network layer. In our setting, since we are employing this in a recurrent neural network, the input and output dimensionality must be the same and thus $c_{out} = c_{in} = c$.

Summary

We thank the reviewer for their careful review and for their acknowledgement of the value of our contribution. To summarize our understanding of the reviewer’s comments, they are impressed with our experimental results and appear to find our contribution valuable. The reviewer appears to be primarily concerned with the biological connection of our model and specifically the ability for traveling waves in the brain to be a viable mechanism for encoding memory. In response we have provided a few references which demonstrate the potential for traveling waves over appropriately longer time scales relevant for memory. Furthermore, we have attempted to describe our interpretation of the relation between our work and the notion of 'time cells'. Finally, we have performed the additional experiment proposed by the reviewer in their questions and included the result in our appendix with extended discussion of these points.

Waves as a memory mechanism

In response to the reviewer’s concern about the biological relevance of our model, and the uncertainty regarding the reality of waves as a memory mechanism in the brain, we would like to highlight two potential answers.

The first is that there is indeed preliminary evidence that some traveling waves do exist over significantly longer timescales than the millisecond scale referenced by the reviewer, and that these do indeed have a relation to memory. Specifically (as noted in our introduction), "King & Wyart (2021) provide evidence that traveling waves propagating from posterior to anterior cortex serve to encode sequential stimuli in an overlapping `multiplexed' manner". In detail, they use a task of carefully constructed sequential stimuli presented over 2 seconds and show that such stimuli are decodable from EEG measurements at sequentially propagated regions of the cortex, implying a wave of activity encoding this information. Others such as Zabeh et al. (2023) also show traveling waves in frontal and parietal lobes encode memory of recent rewards. Therefore, while we certainly agree with the reviewer that certain waves (such as those which exist in the millisecond timescale of the hippocampus) are inappropriate for encoding longer term memories, we believe that there are likely a diversity of types of traveling waves in the brain and therefore this mechanism should not be discounted so quickly.

The second potential answer to the reviewers concern is that the current studies of traveling wave dynamics in the brain are inherently very limited. It is simply very challenging to measure traveling waves in the brain due to the necessity for both high spatial and temporal resolution simultaneously. Therefore, while we agree that there is currently limited evidence that traveling waves may support the type of memory implied by our model, we again do not believe that this should be a reason to discount such a potential mechanism. Specifically, we believe that our model may indeed encourage further targeted study of waves. Furthermore, we note that we are not the first to suggest such a relation between waves and memory as outlined in our introduction and in the review paper (Muller et al. 2018).

I read the section entitled ``Waves with Variable Velocity''. I am not certain how one can see that the velocity is variable from Fig 19, but I appreciate the inclusion of this. Bumping up a notch.

Summary

We thank Reviewer h1zv for their thorough review and constructive commentary, we believe it has truly helped us improve our manuscript and we appreciate the reviewer's time. Below we summarize the reviewer's comments and our responses:

First, (A) it appears the reviewer is primarily concerned by the direct description of neural activity in terms of physical wave dynamics and would prefer an alternative description of the hidden state as a 'register'.

• In response we have altered our introduction to include the mechanistic 'register' intuition, and altered our phrasing to highlight the physical wave comparison is an imperfect analogy.

Secondly, (B) they were concerned with apparent inconsistency between the wave velocity estimates obtained from the regular spatio-temporal activity plots and the Fourier domain plots, leading them to question the value of the Fourier analysis.

• In response we have identified the issue (an initial incorrect calibration based on sampling rate) and corrected the axes of our Fourier analysis plots. We have additionally included a new figure in the appendix (Figure 12) validating our calibration with synthetic data.

In the sections below we adress these concerns in more detail. We note that in our updated submission we have highlighted all text changes in blue to allow for quicker evaluation.

Weakness A - Wave fields

The reviewer is correct in understanding the name 'wave-field' is simply a reference to the hidden state of the network which supports and encourages wave solutions. This is a term derived from physics to describe "the extended area or space taken up by a wave" (wiktionary). We decided to use this name in order to distinguish this type of hidden state from the baseline RNN hidden state which has no notion of spatial organization and no bias towards traveling wave dynamics. Our initial motivation for appealing to physical wave dynamics was in reference to the review paper (Muller et al. 2018) in which they make an analogy with the pilot-wave work of (Perrard et al. 2016). We include a quote below from that review which we believe may directly address some of the reviewer's concerns related to its inclusion:

A recent study examined the dynamics of a silicone droplet bouncing on a vibrating oil bath (Perrard et al. 2016). The vibrations of this medium propel the droplet upward and cause it to bounce continually, creating a pattern of successive impacts across the surface (Fig. 6a). With each bounce, the droplet creates a set of waves in the oil bath. In our analogy, these droplet waves may be similar to those created in the sensory cortex by the impact of afferent, stimulus-driven action potentials. The waves evoked by each impact of the droplet then create a pattern across the oil bath that leaves a trace of the path of the droplet. With sufficient knowledge of the dynamic properties of the medium, one could decode the history of the path of the droplet simply from the observed spatial pattern at each point in time. Furthermore, temporal reversibility is preserved for waves even in chaotic systems; this property, along with a clever manipulation of the experimental apparatus, can be used to cause the droplet to exactly reverse its trajectory (Fig. 6a). It is argued that this last manipulation shows that the global wave field stores information about the trajectory of the droplet with sufficient precision that it can be read back out. While the authors did not have neuroscience in mind and, indeed, note that they have not yet considered practical implementations, they posit that this system implements the basic elements for general computations that could find application in future work.

We appreciate that the reviewer has highlighted that our initial emphasis on this analogy may have been distracting to readers. We have therefore followed the reviewer's suggestion and rephrased our introduction to highlight this is an imperfect analogy while including the more intuitive `register' description earlier on. We have additionally updated Figure 1 to match the wRNN dynamics more closely and resemble physical wave dynamics less.