Parallelizing non-linear sequential models over the sequence length

Summary

This paper proposes parallelizing evaluation and training of nonlinear sequential models using fixed point iteration methods. The paper proposes a method that restates nonlinear differential equations as fixed-point iterations with quadratic convergence, as in Newton's root finding method. The result is the DEER method. Speed and performance results are presented for a NeuralODE and GRU that are parallelized using DEER. Theoretical equivalences and convergence rates are established in the appendix.

Strengths

- The proposed DEER method is well-motivated and presented clearly - The theoretical results seem sound, though I only skimmed the proofs to follow the arguments, rather than check every detail line-by-line - The method is general and can be applicable to a host of nonlinear differential equation methods such as neural ODEs and any nonlinear RNN (e.g. LSTM, GRU, etc) of broad interest to the sequence modeling community. - The method has the advantage of theoretically having quadratic convergence. Experimental results support the claim that this method should lead to nearly equivalent results (up to numerical precision) as evaluating the nonlinear ODE sequentially - For small hidden sizes, the empirical speed ups are substantial

Weaknesses

- The biggest weakness of the paper is the empirical results, particularly related to performance. Only two tasks are considered, a synthetic physics system where HNNs are trained and the EigenWorms task where a GRU is trained. - For the EigenWorms task, the GRU does not significantly perform that much better than several of the baselines and is outperformed by others. It is claimed that the DEER method enables faster training and thus experimentation to identify optimal GRU architectures. But is this the best the GRU could do? - 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. - A greater variety of tasks would increase the impact of the paper and potentially broaden the audience. E.g. 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. - There are many interesting theoretical results pushed to the Appendix. The authors might potentially consider packaging some of the formal statements as propositions in the main paper, to signpost the results and guide the reader to understand the rigor behind the method as well as point to the location of the proofs more explicitly in the Appendix. - A weakness of the proposed method is the cubic complexity in the hidden size that the DEER method incurs. This limits scalability and the applicability of the method to larger scale systems. This is mentioned, but perhaps a further discussion of potential paths to address this could be interesting and helpful.

Questions

1. 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). 2. Can you apply the DEER method to LEM and/or UniCORNN and report the results and compare the runtimes? 3. 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? 4. Can you add a PDE example (discussed in the appendix)?

Rating

6: marginally above the acceptance threshold

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper proposes a novel parallelization scheme for sequential models. The method works by recasting the model as a fixed point problem and using Newton's method to solve it in parallel. Involves computing the Jacobian and Hessian of the model explicitly to enable quadratic convergence. An adjoint method is developed to compute the gradient in parallel. Experimental results are shown in speeding up RNNs and neural ODEs, showing significant speedups especially for long sequence lengths. Extensive proofs are given.

Strengths

+ The method clearly has large speedups in certain training regimes, particularly for long sequences and small batch sizes + The theory of the method is clearly described, with proofs provided in the appendix. + The proofs in the appendix are clearly described. + The method is generally applicable to any sequential method, unlike previous methods which require specific architectures or structural assumptions. + Convergence is quicker than previous methods which did not incorporate Jacobians + Potentially, the method could also be applied at inference time to speed up sequence generation, not just at training time.

Weaknesses

+ 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. + The increased memory usage of DEER at higher hidden dimensions seems a big stumbling block. As far as I am aware, a hidden dimension of 8 would be considered very small and impose a limit on the expressiveness of a recurrent model. This issue is not really addressed in the work.

Questions

+ Can you provide comparisons to the sequential method when the throughput of the GPU is held constant (e.g. the memory is approximately fully utilized)? + 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? + Would you be able to use an approximation for the Jacobian (low-rank, etc) to reduce the memory usage? + How many iterations are typically required to converge? How does this compare against zero-order methods that don't use a Jacobian?

Rating

6: marginally above the acceptance threshold

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

2 fair

Summary

This paper introduces a method for parallel evaluation of non-linear sequence models (neural differential equations, GRUs) as a limit case of differentiable multiple shooting. The key idea is to linearize the multiple shooting update equation and use exact parallel methods (e.g., parallel scans) to evaluate each step of the root-finding procedure. The method is evaluated on time series classifiction, hamiltonian neural ODE training, and benchmarked for efficiency.

Strengths

* The method is clearly presented. The authors do a good job contextualizing the approach with respect to direct multiple shooting, including recent work on adapting it to Neural ODEs (Appendix A.1, A.2) * The efficiency evaluation is quite thorough, investigating the effect of batch size, state dimension and sequence length.

Weaknesses

* 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? * The benchmarking is quite limited, both tasks showcased are small scale. The tasks do a good job at showing relative performance (end-to-end time) improvements, but they do not provide any insight on the method itself. Are there important hyperparameters, methods that could impact the final results and convergence? You mention mid-point in passing in Sec 3.3, could you elaborate on the choice of other interpolation methods that you have tried?

Questions

* How is the recurrent step baseline implemented here? Since you use JAX, have you tried to otimize the recurrence (jit-compile, or fuse)? * 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 commend on whether DEER would be applicable in this case (and ideally show some preliminary result on a simple task?) * 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? * Can you comment on the different rates of approximation of this linearized version of multiple shooting, and regular multiple shooting?

Rating

6: marginally above the acceptance threshold

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Response to all reviewers: new experiments and results

We thank all the reviewers for their comments and questions. There is a typical main concern from the 3 reviewers which is about lacking experimental variations. Therefore, we would like to use this space to inform the reviewers about the new experiments and new results that we obtained during the discussion period. Please note that no changes in the paper has been made yet. We need more time to add new results and arrange the content to fit within the 9-page limit. ### 1. Sequential CIFAR-10 with multihead GRU To address the concerns about the cubic scaling of DEER method, we show that this can be avoided using multihead GRU. For example, instead of using a GRU with 256 channels, we divided them into a GRU with 32 heads with 8 channels each. The channels from different heads are detached, so the Jacobian matrix for each head is only $(8\times 8)$ instead of $(256\times 256)$. Moreover, each head can have different strides and in our case, we use strides from the set $\\{2^0, 2^1, ..., 2^7\\}$ (so there are 4 heads that have the same stride value). The idea of detaching channels is not new, it has been shown to work really well in state-space models like S4, S4D, Hyena, *etc*. Our experiments with GRU + DEER can run about 3 - 5 times faster than with GRU + sequential method, allowing us to run more experiments during this discussion period. The best test set accuracy we obtained was 77.63%. Although this is not the state-of-the-art for sequential CIFAR-10, this is the best accuracy obtained by non-linear recurrent neural network (as shown in [PapersWithCode](https://paperswithcode.com/sota/sequential-image-classification-on-sequential-1)). **Please note that** the aim of our paper is not to introduce a new architecture, but to focus on a computational method that allows faster training for non-linear sequential models. We will put the results in the paper. [comment]: <> (<p style="color:blue">[TODO: put in the paper ???]</p>) ### 2. EigenWorms with LEM The state-of-the-art for EigenWorms dataset is currently held by LEM (T. Konstantin Rusch, *et al.*, 2022). Therefore, we ran experiments using DEER method with LEM architecture, trying to replicate their results. The replication attempt is not as easy as we expected because running the official code didn't give us the results presented in the paper, so we need to run with various sets of hyperparameters. Running the official code (written in Pytorch, using sequential method) took us about 6 hours to complete 1 experiment, while using DEER in JAX only took us less than 30 minutes. This allows us to find the hyperparameters that can achieve the similar results in a short amount of time. We will put the results in the paper. [comment]: <> (<p style="color:blue">[TODO: put in the paper ???]</p>) ### 3. Improving GRU on EigenWorms We use this discussion period to find the hyperparameters that can improve the performance of GRU in EigenWorms. With the new hyperparameters, we can improve the results from 82.1% to 88.0%. Although it is not the state-of-the-art, we show that classical architecture such as GRU can beat Neural RDE, a more modern technique on learning with long sequence. We will update the results in the paper. [comment]: <> (<p style="color:blue">[TODO: put in the paper ???]</p>)

Updates on the paper

We have updated our paper to follow the recommendations and our responses to the reviewers. Here are the list of changes we made in the paper: ### Main content 1. Added a new section on sequential CIFAR-10 experiment using multi-head GRU (section 4.4). Note that we obtained a new result that has higher accuracy than our previous response. 2. Updated Table 1 to reflect the latest results of our hyperparameter tuning of GRU and our reproducibility attempt at LEM. 3. Added a new paragraph (3rd paragraph in section 3.5) to elaborate the only 1 hyperparameter of DEER method. 4. Added a new paragraph (last paragraph in section 4.3) to briefly explain our reproducibility attempt at LEM for EigenWorms. ### Appendix 1. Added a new section on "Other kinds of interpolations" to list the other interpolations we considered in NeuralODE (section A.6) 2. Added a new section on experimental details of multi-head GRU on sequential CIFAR-10 (section B.4) 3. Added a new section on the effect of tolerance level on convergence, to explain the only 1 hyperparameter of DEER (section C.1) 4. Added a new section to present the results of DEER vs sequential method using the same amount of memory (section C.3)