Learning by solving differential equations

TL;DR

The paper addresses training instability in deep learning caused by highly variable gradient fields and investigates using higher-order Runge-Kutta methods to more accurately track the gradient-flow ODE $\dot \theta = -\nabla L(\theta)$. It shows vanilla RK4 can be competitive on small-batch DL tasks but faces stiffness and generalization gaps in large-batch settings; to counter this, it introduces three adaptations—preconditioning, adaptive learning rates, and momentum—that empirically bridge the gap or surpass Adam in several simple benchmarks. Through AdaGrad-like preconditioning, Drift-Adjusted Learning rate (DAL) variants, and RK4 with momentum, the approach demonstrates meaningful performance gains on MNIST and Fashion-MNIST and offers a promising optimization perspective at the intersection of numerical ODE methods and deep learning. The study also discusses limitations, notably computational cost and scalability to larger, more realistic workloads, and outlines directions for broader evaluation and integration with standard DL tricks. Overall, the work motivates further exploration of RK-based optimizers as a principled alternative to first-order methods in neural network training.

Abstract

Modern deep learning algorithms use variations of gradient descent as their main learning methods. Gradient descent can be understood as the simplest Ordinary Differential Equation (ODE) solver; namely, the Euler method applied to the gradient flow differential equation. Since Euler, many ODE solvers have been devised that follow the gradient flow equation more precisely and more stably. Runge-Kutta (RK) methods provide a family of very powerful explicit and implicit high-order ODE solvers. However, these higher-order solvers have not found wide application in deep learning so far. In this work, we evaluate the performance of higher-order RK solvers when applied in deep learning, study their limitations, and propose ways to overcome these drawbacks. In particular, we explore how to improve their performance by naturally incorporating key ingredients of modern neural network optimizers such as preconditioning, adaptive learning rates, and momentum.

Figures (18)

• Figure 1: For MNIST ( top row) and Fashion MNIST ( bottom row) trained on MLP, vanilla RK4 achieves better test accuracy compared to Adam for small batches ( left column) while we observe a generalization gap for large batches ( right column). See Appendix \ref{['appendix:figure_1']} for experiment details.
• Figure 2: RK4 with AdaGrad-like preconditioning helps bridge the gap between Adam and vanilla RK4 in the large batch regime for MLP trained on MNIST (left) and Fashion-MNIST (right). See Appendix \ref{['appendix:bridging_the_gap_preconditioning']} for experiment details. Additional experiments on CIFAR-10 and a ResNet-18 model yield similar results; see Fig. \ref{['figure:conditioners_cifar']} and Appendix \ref{['section:cifar10_rk4_modified_adagrad']}.
• Figure 3: RK4 combined with an adaptive learning rate bridges the generalization gap in the large batch setting and is competitive with Adam for MLP trained on MNIST (left) and Fashion-MNIST (right). See Appendix \ref{['appendix:bridging_the_gap_adaptive_learning_rate']} for experiment details. Additional experiments on CIFAR-10 and a ResNet-18 model yield similar results; see Fig. \ref{['figure:cifar_momentum_dalr']} and Appendix \ref{['section:cifar10_rk4_momentum_adaptive_lr']}.
• Figure 4: Combining momentum with RK4 bridges the generalization gap in the large batch regime and achieves better test accuracy than both Adam and vanilla RK4 on MNIST (left) and Fashion MNIST (right). See Appendix \ref{['appendix:bridging_the_gap_momentum']} for experiment details. Additional experiments on CIFAR-10 and a ResNet-18 model yield similar results; see Fig. \ref{['figure:cifar_momentum_dalr']} and Appendix \ref{['section:cifar10_rk4_momentum_adaptive_lr']}.
• Figure 5: RK4 is competitive with Adam on MNIST with a 3 hidden layer MLP of 500 neurons each when trained with small batches of size 16 for 10000 steps. Adam's learning rate was tuned to 0.001 while the decay parameters were set to Optax defaults optax2020github, and RK4 learning rate was tuned to 0.003; 5 seeds.

Theorems & Definitions (5)

• Definition 2.1
• Lemma 4.1: Preconditioning lemma
• proof
• Lemma B.1
• proof