Learning via nonlinear conjugate gradients and depth-varying neural ODEs

Abstract

The inverse problem of supervised reconstruction of depth-variable (time-dependent) parameters in a neural ordinary differential equation (NODE) is considered, that means finding the weights of a residual network with time continuous layers. The NODE is treated as an isolated entity describing the full network as opposed to earlier research, which embedded it between pre-and post-appended layers trained by conventional methods. The proposed parameter reconstruction is done for a general first order differential equation by minimizing a cost functional covering a variety of loss functions and penalty terms. A nonlinear conjugate gradient method (NCG) is derived for the minimization. Mathematical properties are stated for the differential equation and the cost functional. The adjoint problem needed is derived together with a sensitivity problem. The sensitivity problem can estimate changes in the network output under perturbation of the trained parameters. To preserve smoothness during the iterations the Sobolev gradient is calculated and incorporated. As a proof-of-concept, numerical results are included for a NODE and two synthetic datasets, and compared with standard gradient approaches (not based on NODEs). The results show that the proposed method works well for deep learning with infinite numbers of layers, and has built-in stability and smoothness.

Standalone Neural ODEs with Sensitivity Analysis

“…As we will demonstrate, this enables a general sensitivity analysis of the mapping. There is a recent example where NODEs were used in isolation from conventional network layers (Baravdish et al, 2022) but only for small networks on 2D toy datasets. We substantially scale the size of the NODE network, and apply it to image input data.…”

“…This optimization method is further extended to include the Sobolev gradient for trainable weights. Such problem formulation has only been considered very recently in (Baravdish et al, 2022), and we significantly expand on this.…”

“…We solve the minimization problem using a conjugate gradient based iterative method proposed by Baravdish et al (2022), where the descent direction, involving the gradient of E(W, b), are calculated via the adjoint problem (see Lemma 4.1 below) and the learning rate is determined by the neural sensitivity problem (2) with ξ 0 = 0.…”

“…In this experiment, we compare the Euler numerical scheme against an adaptive ODE solver (ODE45) for the discretization of sNODE when NCG training is considered on (14×14) MNIST. We note that the ODE45 solver was previously tested by Baravdish et al (2022) in the same setting of training sNODE via NCG; however, experiments were limited to 2D toy datasets. We choose tanh activation functions due to their smoothness property, and set µ 1 = 1, µ 2 = µ 3 = µ 4 = 0 in (3) hence minimizing the 2 -loss without weight decay.…”