Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This method proposes a second-order optimization network for neural ODEs. Inspired by OC, this paper adopts differential programming to derive backward ODE for higher-order derivatives at constant memory cost. This paper also explores using low-rank representation for the second-order optimization to make more efficient updates. The paper shows their method achieves faster convergence than first-order methods across three different tasks. This framework also enables the optimization of the architecture with the example of integration time of neural ODE.
main_review: Originality
- This work proposes a novel second-order optimization network specifically for neural ODEs. It adopts ideas from OCP to neural ODE to achieve this goal. It is clear how this work differs from previous work, and the related works are appropriately cited. I like how it introduces the OCP perspective to neural ODE and it can treat integration time as an optimization variable to optimize the architecture directly.

Quality
- This work is technically sound. It makes sense that using a second-order optimization approach can make neural ODE converge faster and this claim is well-supported by three different tasks, which demonstrate faster convergence and higher accuracy using their method than using the first-order method.

- I would like to see the reasoning behind Figure 9, which claims that this method is less sensitive to hyper-parameter changes. It is not very obvious based on the theoretical ground and is worth exploring.

Clarity
- The overall structure is well-organized. 

- Different fonts of x are present in lines 19 and 21 before it is clarified in preliminary.

- On line 118: what are “These functions are known as...” Which are the functions here?

Significance
- Other researchers may use their method to run neural ODE, but I doubt the significance of developing even higher-order optimization based on this work. Also, I am not sure if it is worth implementing the method given its complexity and limited improvement on computational time. Please discuss these in Section 5.

Additional Questions
-What is the relationship between step sizes and integration time? Does increasing the time step size have the same effect as decreasing the integration time for image and CNF datasets?

-Why do you choose a Runge-Kutta 4 adaptive solver instead of an implicit solver such as the implicit Adams method? Do the benefits depend on the specific solver you choose?

-For Eq. 14, why do you call this specific formula a preconditioned update law? Why do you choose this to update theta?

-Is there any mathematical formulation that supports how much the second-order optimization improves the convergence rate of Neural ODE?

-Why using second-order optimization can also improve the accuracy of the experiments as demonstrated in figure 9?

limitations_and_societal_impact: The limitations were adequately discussed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper proposes a method to derive second-order derivatives for Neural ODEs, which requires O(1) memory in theory. For practical feasibility, the authors propose to use a low-rank approximation of the (dynamic) Hessian matrix, together with Kronecker-based factorization to further reduce the memory cost. The authors validated their methods in image classification, time-series analysis and continuous generative modeling.
main_review: Pros:

This method, to my knowledge, is novel in terms of an efficient method to estimate the second-order gradient information for Neural ODE models. The paper is in general well-written with solid theoretical analysis and experimental validations. It would be better if the authors could release code for such a complicated method.

Cons:

I have a few comments and concerns rather than critiques for the authors to clarify.

1. I quickly went through the proof, found the key result rely on two important assumptions (line 526 in appendix), a) $z(t) \otimes g(t)$ are uncorrelated in time $t$, b) $z$ and $g_i$ are pairwise independent. Assumption a) looks very strong even unpractical to me, since in the update it's always $z(t+\delta (t))=z(t)+\Delta(t)$, omit the details of $\Delta(t)$ here. It's clear that $z(t)$ and $z(t+\delta (t))$ are highly-correlated.

2. The author did not specify how to choose a proper rank $R$

3. In terms of memory consumption, If I get it correctly, Full rank second-order adjoint > Low-rank second-order adjoint > first-order adjoint, so the memory cost of the proposed method is in-between. Even though it's O(1) in time $t$, I'm concerned the memory cost has a large constant. \
Specifically, suppose the parameter number is $n$ and activation dimension is $m$, first-order adjoint needs $O(n+m)$ memory, while the proposed method needs $O(n+m+Rm+Rn)$, where $R$ is the rank. This memory cost is required to store the first and second-order gradients. \
In terms of update, the method requires a much higher peak memory of $O(n^2+m^2)$, though it's reusable, because Eq.13 requires the inversion of two matrices of size $n^2$ and $m^2$, this inversion could also be very expensive.  \
Table 7 in appendix shows that the proposed method requires 40% more memory than first-order methods. Note that this is for small-size image ($m=32\times 32$) on CIfar, on large images such as ImageNet ($m=224\times 224$), I expect the memory would be even much larger. \
So the scalability to large scale networks especially with high-resolution images is one of my concerns.

Suggestions

It's shown that numerical errors in gradient for Neural ODE is the cause of bad empirical performances, so I wonder how does the proposed method compare with first-order methods that tackle the numerical issues carefully [1,2,3].

[1] Gholami, Amir, Kurt Keutzer, and George Biros. "Anode: Unconditionally accurate memory-efficient gradients for neural odes." arXiv preprint arXiv:1902.10298 (2019).  \
[2] Zhuang, Juntang, et al. "Adaptive checkpoint adjoint method for gradient estimation in neural ode." International Conference on Machine Learning. PMLR, 2020.  \
[3] Zhuang, Juntang, et al. "MALI: A memory efficient and reverse accurate integrator for Neural ODEs." arXiv preprint arXiv:2102.04668 (2021).
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This work proposes an alternative to the recursive adjoint method for computing higher order gradients of neural ODEs based on optimal control theory that promises to be more efficient and more accurate. 
main_review: My main concern is that the authors do not take sufficient care in statements about the memory and time costs. The adjoint method is O(1) memory *in depth* i.e. in time, but it is O(n) in the number of parameters, i.e. $dim_R(\theta) = n$. Naively computing higher-order derivatives by recursively applying the adjoint method has O(n) memory footprint and O(n) time, but these calculations are independent and can be parallelised to give O(n^2) memory and O(1) time (as described by the authors L98). The proposed method maintains O(1) in depth, but has a footprint of O(n^2) ($Q_{uu}$ is matrix valued $\in R^{n\times n}$), and is O(1) in time.

So there isn’t actually any change in the memory or time costs as compared to recursively applying the adjoint method. There may be a difference in the error, as stated in L41-43, but there is no evidence presented in favour of this claim nor evidence that this is of practical consequence in the training of neural ODEs, which is the evaluation criteria used in this work.

L97-101 make some arguments against the recursive method, stating that the method scales unfavourably due to recursion and errors. That k=2 is apparently the useful limit of this recursion (‘*second*-order neural ODE optimizer') diminishes the significance of the first complaint and it is diminished further when we note that the time cost is linear in the order of the optimization, which is a fixed hyperparameter. Evidence is not presented to support the second complaint, and this is a missed opportunity as whether the accumulation of errors is of any significance in training is an interesting scientific question. Recursively applying the adjoint method is the essential baseline to compare the method against, particularly given the proposals from section 3.2 onwards to approximate the optimisation and L194 which introduces a further approximation explicitly to trade off complexity for performance, such that both methods run with some error.

This is a shame as the work is otherwise interesting and highly creative, and seeing this comparison seems essential.

I’d also be interested in some discussion about how the number of terms in equation 9 grows in the order of the optimisation. I followed the discussion in the appendix to Almubarak 2019, though the appendix indicates that there is a further reference within Almubarak 2019 that is to be followed next, though doesn’t specify which. Combinatorially, we expect to get 2^(n+1) - 2 (i.e. 2 for first order, 6 for second order, 14 for third order…) This seems to suggest that there is a further term missing from the memory footprint. Including the optimisation order, k, the costs seem to be
Recursive adjoint: memory O(n^k) time O(k)
Proposal: memory O(2^k n^k) time O(1)
Given that k is generally a fixed hyperparameter and ≤ 2 I doubt this is a pressing problem, but some comment would be appreciated.

Section 3.3 would benefit from including the choice of initialisation in the discussion. We can freely vary the integration time without modifying the output by making corresponding changes to the magnitude of the weights in the derivative function i.e. if you want to double the integration time, half the weights in the final layer of the derivative network. So discussions about the integration interval are missing important context. Furthermore, given that the model uses a CNN feature extractor (Appendix A.4 gives this as a 3 layer conv net) prior to the NODE, and that output of the integration is passed through a linear layer to get the prediction, the results in Figure 5 are explained equally as the ODE being superfluous. The correct baselines for comparison here would be 0 integration, which shows 50% reduction in train time for what appears to be a 2-3 percentage point reduction in accuracy, and a model that has the same capacity as the NODE added. In this case the ODE function is a 2 layer conv net, so a 5 layer conv net in total would be appropriate. Figure 13 in the appendix, showing the same experiment for MNIST, presents similarly inconclusive results, in that case the change in accuracy is even less significant across the range of values tested.

The results in Figure 7 are mixed, which isn’t accurately reflected in the discussion in the **Convergence and computational efficiency** paragraph. On SpoAD, SGD trains as quickly and to the same performance. On CharT the three methods reach similar levels of performance. The absence of errors in Table 2 makes it difficult to compare the significance of the bolded results (and appropriate baselines discussed above are absent).

Table 3 does not select the best training time model from Figure 5. There is a higher-performing and faster-training model on that plot, something in the region of t = 0.4 or 0.5 looks to have performance at least as good as that of the t = 1.0 model at a train time percentage of 50-75%. 

Figure 9 is a useful plot, though the caption doesn’t seem to match the data: the model is sensitive to the choice of learning rate, with the training loss varying by approximately ~2x between high and low learning rates on CIFAR10. A higher learning rate is not investigated for Adam, which would appear to be on trend to match the performance of the proposal if a higher learning rate were used (perhaps it wouldn’t, the evidence isn’t provided.) Better would be to provide training plots as in Figure 7 for varying learning rates for the different methods.

Also absent is a discussion of the effect of batch size on the training plots. With the first order methods having a lower memory overhead, there should be more scope for increasing the batch size to accelerate training. Has this been investigated?

Overall, this is an exciting project that is let down by its evaluation. The absence of obvious baselines, most critically the recursive adjoint method, and seemingly cherry picked examples in the comparisons that are present, make it difficult to evaluate the impact of this work.

--- minor points
The presentation is clear and of high quality.
L91 should be ‘the backpropgation algorithm’
L122 seemly should be seeming
L160 ‘much more efficiently’


### Update
Following the rebuttal and back and forth with the authors, my concerns have been addressed and I'm happy to raise my score (4->7) to reflect this.
limitations_and_societal_impact: Impact is adequately addressed.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: This paper presents a novel second-order optimization technique for neural ordinary differential equations (NODEs), derived from the optimal control perspective of NODEs. In a nutshell, the authors consider the second-order Taylor expansion of the loss term, leading to matrix-valued ODEs that give second-order derivatives. A series of approximations are proposed to compute a low-rank representation of the second-order preconditioning matrices. The experiments demonstrate that the proposed optimization method converges faster than a first-order optimization routine in all neural ODE applications at the cost of slightly higher memory usage. The method also allows optimizing the integration time in continuous-time normalizing flows.
main_review: The integration operation in continuous-time models can be overwhelmingly time-consuming, which oftentimes makes it difficult to apply neural ODEs to even simple, low-dimensional problems. Therefore, the proposed approach addresses an overlooked but fundamental restriction of a rather recent model family. The paper is well-written. Despite the heavy technical details, the presentation is very clear and I had no difficulty grasping the motivation and ideas. The experiments and the discussion of the results are exhaustive. Experiments indeed show that the presented optimization scheme improves the NODE training in all three applications of NODEs. 

A few concrete suggestions and questions:
- Since the approximations are very involved, having a broader picture of them would be nice. For instance, a cartoon illustration of the series of approximations and what they aim to address could help the reader gain insights more easily. Also, additional (simple) ablation studies that compare approximations with the corresponding exact expressions might reveal how good/tight the approximations are.
- A direct comparison against a Quasi-Newton method (I suggest LBFGS) on a simple problem would be very interesting to see as such methods can scale to low/mid-size neural nets.
- Line 243: Which tolerance values (absolute/relative) does this refer to? Also, any justification for not using the default values?
- To support the claims in Line 272, it would be nice to test NODE with a regularization scheme, e.g., Finlay et al., 2020. Also, one may further check the number of function evaluations.
- Looking at Fig13 (train time), I wonder why it first drops and then goes up. Maybe because of a change in the stiffness of the vector field? Can this be experimentally demonstrated?

Minor details:
- I suggest writing "back-propagation" instead of "Back-propagation". 
- Line117: ',and etc' seems unnecessary. Also, what is w.l.o.g?
- An important reference to [1] is missing.

[1] Stapor, Paul, Fabian Fröhlich, and Jan Hasenauer. "Optimization and profile calculation of ODE models using second order adjoint sensitivity analysis." Bioinformatics 34.13 (2018): i151-i159.

**Post-rebuttal comment:** Thanks authors for the rebuttal. As noted in my original review, this is a very nice paper and should be definitely accepted.
limitations_and_societal_impact: The limitations are clearly described. No potentially negative societal impact I can think of.
review_title: null

Meta-review: This paper proposes a second-order optimization algorithm for neural ODEs. Inspired by optimal control theory, the paper derives backward ODEs for higher-order derivatives at constant memory cost. The authors shows their method achieves faster convergence than first-order methods across three different tasks.

The contribution is novel and relevant to NeurIPS, and is appreciated by the reviewers. The main criticism that surfaced during review is the absence of certain baselines, raised by reviewer ikHp: After back and forth discussion between the authors and reviewer this concern was addressed sufficiently for the reviewer to raise their review score. The authors also replied extensively to comments by the other reviewers, and no further serious concerns were identified. The reviewers formed a consensus recommendation to accept the paper.
Decision: Accept (Spotlight)