Smoothness-dependent reconstruction with Heun’s method.
The bounds in Prop. 3 and 4 do not depend on the smoothness with respect to the weights of the f(.,\theta_{n}^{N}) . Only the magnitude of the residuals plays a role in the correct recovery of the activations and estimation of the gradient. Hence, there is no apparent benefit of having such a network behave like a Neural ODE. We now turn to Heun’s method, a second order integration scheme, and show that in this case smoothness in depth of the network improves activation recovery. A HeunNet (Maleki et al., 2021 ) of depth N with parameters \theta_{1}^{N},\dots,\theta_{N}^{N} iterates for n=0,\dots,N-1 :

y_{n}=x_{n}+\frac{1}{N}f(x_{n},\theta_{n}^{N})\quadand\quad x_{n+1}=x_{n}+\frac{1}{2N}(f(x_{n},\theta_{n}^{N})+f(y_{n},\theta_{n+1}^{N}))

These forward iterations can once again be approximately reversed by doing for n=N-1,\dots,0 :

\tilde{y}_{n}=\tilde{x}_{n+1}-\frac{1}{N}f(\tilde{x}_{n+1},\theta_{n+1}^{N})\quadand\quad\tilde{x}_{n}=\tilde{x}_{n+1}-\frac{1}{2N}(f(\tilde{x}_{n+1},\theta_{n+1}^{N})+f(\tilde{y}_{n},\theta_{n}^{N})), (9)

which also enables approximated backpropagation without storing activations. When discretizing an ODE, Heun’s method has a better O(\frac{1}{N^{2}}) error, hence we expect a better recovery than in Prop. 3 . Indeed, we have:

Assume assumption 2 . Denote by L_{f}^{\prime} the Lipschitz constant of x\mapsto\frac{1}{2}(f(x,\theta_{n+1}^{N})+f(x-\frac{1}{N}f(x,\theta^{N}_{n+1}),\theta_{n}^{N})) , by L_{\theta} the Lispchitz constant of \theta\mapsto f(\cdot,\theta) and by L^{\prime}_{\theta} that of \theta\mapsto\partial_{x}f(.,\theta) . Let C^{\prime}_{f}=\frac{1}{4}L_{\theta}^{\prime}L_{\theta} . Finally, define \Delta_{\theta}^{N}\coloneqq\max_{n}\|\theta_{n+1}^{N}-\theta_{n}^{N}\|^{2} .
Using Heun’s method, we have: \|x_{n}-\tilde{x}_{n}\|\leq\frac{(e^{L_{f}^{\prime}}-1)C^{\prime}_{f}}{L_{f}^{\prime}N}\times\Delta_{\theta}^{N}+O(\frac{1}{N^{2}}).

This bound is very similar to that in proposition 3 , with an additional factor \Delta_{\theta}^{N} .
Hence, we see that under the condition that \Delta_{\theta}^{N}=O(\frac{1}{N}) , the reconstruction error \|x_{n}-\tilde{x}_{n}\| is in O(\frac{1}{N^{2}}) . In the linear case, we have proven under some hypothesis in Prop. 2 that such a condition on \Delta_{\theta}^{N} holds during training. Consequently, the smoothness of the weights of a HeunNet in turns helps it recover the activations, while it is not true for a ResNet.
This provides better guarantees on the error on gradients:

Suppose assumption 2 . Suppose in addition that \partial_{x}f(.,\theta) admits Lipschitz constant, \partial_{\theta}f(.,\theta) admits a Lipschitz constant and an upper bound, all of which are independent of \theta . Then one has \|\tilde{\nabla}_{\theta^{N}_{n}}L-\nabla_{\theta^{N}_{n}}L\|=O(\frac{\Delta_{\theta}^{N}}{N^{2}}+\frac{1}{N^{3}}).

Just like with activation, we see that Heun’s method allows for a better gradient estimation when the weights are smooth with depth. Equivalently, for a fixed depth, this proposition indicates that HeunNets have a better estimation of the gradient with the adjoint method than ResNets which ultimately leads to better training and overall better performances by such memory-free model.