Paper: Parallelizing non-linear sequential models over the sequence length (ICLR-2024)

Reviewer comment (Experiments):
It would have been interesting to see the DEER method applied to some of the other methods such as LEM or UnICORNN since they could presumably be trained with the DEER method also. Would they maintain the same performance? Could better performance be achieved with the DEER method since they could potentially be trained faster and thus be hyperparameter tuned more? An exploration of this would strengthen the paper and broaden the impact.

Author response:
We have ran experiments using DEER on LEM. It is more challenging to reproduce their results than we thought as running the code from the official repository does not really reproduce the results on the paper. So we used DEER method to tune the hyperparameters efficiently and we can obtain similar results to their paper. With DEER, one experiment can run in less than 30 minutes, while with the code from the official repository, it took us 6 hours to run an experiment.  We will put the results in the paper.

Reviewer comment (Experiments):
A greater variety of tasks would increase the impact of the paper and potentially broaden the audience. For example, there are many other common sequential "long sequence" tasks ranging from sequential MNIST, sequential CIFAR, to the LRA benchmarks. It would have been interesting to see how a nonlinear RNN such as GRU performed on LRA.

Author response:
Yes, we have added a new experiment result on sequential CIFAR-10 (sCIFAR-10) using multiple head GRU with different strides. We can get around the cubic complexity of DEER by splitting the GRU into multiple heads. For example, instead of using correlated 256 states, we split them into 16 heads with 16 states each. We ran an experiment with this idea on sequential CIFAR-10 where we elaborate in the "Response to all reviewers" section 1 above. We will put the results in the paper.

Reviewer comment (Experiments):
Can you somehow quantify the claim that the DEER method can improve the performance of the GRU since it allows more hyperparameter tuning? Perhaps report the results of a single sequential GRU run, and then show the improvement that can be achieved by hyperparameter tuning the GRU for multiple runs (allowed by the parallelization of the DEER method).

Author response:
We ran hyperparameters tuning during the discussion period as the reviewer suggested and we can achieve the improvement for GRU from 82.1% to 88.0%, beating a more modern technique Neural RDE. We will update the results in the paper.

Reviewer comment (Experiments):
Can you apply the DEER method to LEM and/or UniCORNN and report the results and compare the runtimes?

Author response:
We have ran experiments using DEER on LEM. It is more challenging to reproduce their results than we thought as running the code from the official repository does not really reproduce the results on the paper. So we used DEER method to tune the hyperparameters efficiently and we can obtain similar results to their paper. With DEER, one experiment can run in less than 30 minutes, while with the code from the official repository, it took us 6 hours to run an experiment.  We will put the results in the paper.

Reviewer comment (Experiments):
Can we add more tasks? Is it possible to run a GRU or other nonlinear RNN on Long Range Arena? Or does the cubic complexity of DEER prevent running large enough models for some of these tasks? Can we at least add other common RNN benchmarks such as sequential CIFAR?

Author response:
Yes, we have added a new experiment result on sequential CIFAR-10 (sCIFAR-10) using multiple head GRU with different strides. We can get around the cubic complexity of DEER by splitting the GRU into multiple heads. For example, instead of using correlated 256 states, we split them into 16 heads with 16 states each. We ran an experiment with this idea on sequential CIFAR-10 where we elaborate in the "Response to all reviewers" section 1 above. We will put the results in the paper.

Reviewer comment (Experiments):
Can you add a PDE example (discussed in the appendix)?

Author response:
Although DEER can be used for PDE, the focus of our paper is for one dimensional differential equation (i.e. ODE or discrete difference equation). Applying it for a PDE would require writing the solver of the linear equation which would take a significant amount of effort. Exploring this direction is a possible future work. For now, we put an example on how to apply DEER framework to a simple PDE case in the Appendix.

Reviewer comment (Evaluation):
The practical importance of this method is somewhat unclear. From a practical point of view, the DEER method is a mechanism to use more memory in order to speed up the forward and backwards pass. Therefore, many of the experiments, particularly figure 2, are not really a fair comparison, as it's very common to increase the batch size until the memory is fully utilized. In other words, the fairer comparison would be to fix the throughput (i.e. FLOP/s or memory usage) of DEER and the sequential method the same and to compare the two methods. I think it's quite plausible that for a given throughput, using DEER to do more minibatch steps with smaller batch size can result in faster convergence. However, I don't believe this argument is made in the paper. Figure 4 gets close, but I don't think this is normalized in terms of throughput.

Author response:
We ran an experiment using LEM on EigenWorms to compare DEER and sequential methods. To make the memory consumption similar, we increase the batch size when running with the sequential method. As the result, although DEER method took more steps to achieve the similar results as sequential method, it was 3 times faster in terms of wall-clock time. This confirms Reviewer EMpw's hypothesis that *"for a given throughput, using DEER to do more minibatch steps with smaller batch size can result in faster convergence"*. We will put the results in the paper.

Reviewer comment (Experiments):
How do RNNs with smaller hidden sizes compare to those with larger hidden sizes? Is it the case that the increased speed of DEER is able to offset for a reduction in hidden size that might be needed to use DEER?

Author response:
The unfavourable scaling of DEER with respect to the hidden size can be avoided by having multiple “heads” where each head has a smaller hidden size. We ran a new experiment with this idea on sequential CIFAR-10 where we elaborate in the "Response to all reviewers" section 1 above.

Reviewer comment (Experiments):
How many iterations are typically required to converge? How does this compare against zero-order methods that don't use a Jacobian?

Author response:
With a typical untrained GRU with 16 hidden sizes, 10000 sequence length, float32, it usually takes about 4 iterations to converge (i.e. max absolute error < 1e-4). With zero-order methods, it typically takes about 23-26 for an untrained GRU. Despite it not looking very different for untrained GRU, the number of zero-order iteration could significantly increase with different settings. For example, multiplying all the parameters of untrained GRU by 5 would get the numbers to about 7 for our method and 330 - 360 for the zero order method.

Reviewer comment (Experiments):
It is not clear whether this approach improves over direct multiple shooting (which is also applicable to GRUs). Is there some drawback to the linearization you introduce? Do you require more steps to converge?

Author response:
In regards to the linearization we refer to in the appendix, we realise that we made a mistake: the direct multiple shooting layer (MSL) equation (eq. 17 in our paper) is actually already a linear equation ($b_{i+1}^{(k+1)}$ is linear w.r.t. $b_i^{(k+1)}$), so we didn’t actually do linearization over MSL because it’s already linear. We apologise for the mistake. As they are both linear and can use Newton’s method, the number of iterations to converge should be the same.

To show the advantage of our method over the multiple shooting layer (MSL), we ran the runtime comparisons between DEER, multiple shooting layers from torchdyn (maintained by the authors of the MSL paper: S. Massaroli, *et al.*, 2021), and our JAX implementation of MSL from torchdyn. In the comparison, we used a simple ODE case with 10,000 time points.

| Method | Device | Runtime | Notes |
|-|-|-|-|
| **DEER** | **GPU** | **1.62 ms** | |
| torchdyn's MSL | GPU | 12.47 s | 20 sections |
| torchdyn's MSL | GPU | 11.80 s | 200 sections |
| JAX MSL | GPU | 202 ms | 20 sections |
| JAX MSL | GPU | 117 ms | 200 sections |
| DEER | CPU | 105 ms | |
| torchdyn's MSL | CPU | 4.59 s | 20 sections |
| torchdyn's MSL | CPU | 3.60 s | 200 sections |
| JAX MSL | CPU | 8.39 ms | 20 sections |
| **JAX MSL** | **CPU** | **7.32 ms** | **200 sections** |

Although our method is slower in CPU, DEER can take more parallelization advantage with GPU. We will put this result in the paper.

Reviewer comment (Reproducibility):
How is the recurrent step baseline implemented here? Since you use JAX, have you tried to optimize the recurrence (jit-compile, or fuse)?

Author response:
Yes, we used jax.lax.scan and jit-compile (which I believe includes operator fusion) for the sequential method.

Reviewer comment (Experiments):
Although the scaling of the proposed method in with state dimension is quite unfavourable, many existing state-of-the-art sequence models (S4, H3, Hyena, RWKV) take a specific "multi-SISO" (single-input, single-output) form, where each channel of the model (in width) is assigned its own state. Although these models typically use linear recurrences that can simply apply the convolution theorem or a parallel-scan, the same structure can be used in the nonlinear case. Since these models are usually trained with small state dimensions (8 - 64), this would correspond to integrating many more systems in parallel (model width times batch size, rather than batch size as in the GRU case). Can the authors comment on whether DEER would be applicable in this case (and ideally show some preliminary result on a simple task?)

Author response:
We agree with the reviewer’s suggestion on having multiple independent states. We can group the channel dimension into several “heads”, thus avoiding the cubic scaling for DEER. We ran an experiment with this idea on sequential CIFAR-10 where we elaborate in the "Response to all reviewers" section 1 above.

Reviewer comment (Evaluation):
How many steps do you need to converge with DEER? How much higher is the latency of one DEER evaluation vs one step of a corresponding GRU?

Author response:
With a typical untrained GRU, it usually takes about 4 iterations, but during the training it can take more steps to converge. For a GRU of input and hidden sizes of 16, batch size 16, sequence length 10K, a DEER iteration typically takes 5.4 ms, while one step of GRU (from one sequence point to another sequence point) takes 22 μs. Although one DEER iteration takes much longer than one sequential step, the number of required iterations is usually much less than the number of sequence steps.

Reviewer comment (Theory):
Can you comment on the different rates of approximation of this linearized version of multiple shooting, and regular multiple shooting?

Author response:
This is related to our earlier comment regarding linearization. We made a mistake by saying that we do linearization over direct multiple shooting Newton’s method update step, while the direct multiple shooting update step is actually already linear.