Learning via nonlinear conjugate gradients and depth-varying neural ODEs

TL;DR

This paper introduces a novel method for reconstructing depth-varying parameters in neural ODEs using nonlinear conjugate gradients, demonstrating improved stability and smoothness in deep learning models with infinite layers.

Contribution

It develops a new parameter reconstruction approach for neural ODEs based on nonlinear conjugate gradients, including mathematical analysis and a Sobolev gradient for stability.

Findings

Method performs well on synthetic datasets

Outperforms standard gradient approaches

Ensures stability and smoothness in deep networks

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.