7th June 2023, 4 min read

Neural Network Training using Stiff ODE Solvers

Original post is here eklausmeier.goip.de/blog/2023/06-07-neural-networking-training-using-stiff-ode-solvers.

This post recaps the paper from Aaron J. Owens and D.L. Filkin from 1989: Efficient Training of the Back Propagation Network by Solving a System of Stiff Ordinary Differential Equations.

1. Neural network. Below is a single "neuron":

Input to node $Q$ is given in $A_{i} \in R$ , weights $W_{i} \in R$ , then node $Q$ computes

Q_{in} = \sum_{i} W_{i} A_{i} .

Output from node $Q$ is

Q_{out} = f (Q_{in}) a .

The semilinear function $f (x)$ , sometimes called "squashing function", is an increasing function of $x$ that has a limited range of output values. A commonly used squashing function is the sigmoid

f (x) = \frac{1}{1 + \exp (- (x + ϕ))},

$ϕ$ is called the threshold of the node. The squashing function has limits $0 < f (x) < 1$ , which guarantees that the output of node $Q$ is limited to the range $0 < Q_{out} < 1$ , regardless of the value $Q_{in}$ . It looks like this, from Wolfram Alpha, for $ϕ = 0$ .

Derivative of $f$ is

\frac{d f (x)}{d x} = {(f (x))}^{2} \exp (- (x + ϕ)) = (1 - f (x)) .

Error from output $Q_{out}$ versus desired output $Y$ is

E := {(Y - Q_{out})}^{2} .

Multiple "neurons" from above can be stacked together. All nodes, which are not directly connected to input or output, are called "hidden layers". Below is a three-layer neural network.

2. Gradient descent. To change the weights $W_{i}$ via the method of steepest descent, find the gradient of $E$ with respect to $W_{i}$ and then change $W_{i}$ to move $E$ toward smaller values. Thus change according

\begin{matrix} (1) & Δ W_{i} = - \frac{d E}{d W_{i}} . \end{matrix}

Applying the chain rule and using above squashing function $f$ gives

\begin{matrix} (2) & Δ W_{i} = 2 (Y - Q_{out}) Q_{out} (1 - Q_{out}) A_{i} . \end{matrix}

The weights are initialized to small random values. The weights are now changed according below rule:

\begin{matrix} (3) & W^{ν + 1} = W^{ν} + η Δ W^{ν} + α (W^{ν} - W^{ν - 1}) . \end{matrix}

Both $η$ and $α$ are adjustable parameters and problem dependent. Typically with magnitudes smaller unity. Parameter $α$ is called "momentum", parameter $η$ is called "training rate".

3. Stiff training. The authors, Owens+Filkin, were impressed by the similarities between temporal history of the weight changes and those of ordinary differential equations that are stiff. In place of the discrete equation (1), the weights $W_{i}$ are now to be considered functions of time. To change the weights so that the squared predition error $E$ decreases monotonically with time, the weights are required to move in the direction of gradient descent. Thus the discrete equation (1) is replaced with the ordinary differential equation (ODE)

\frac{d W_{i}}{d t} = - \frac{d E}{d W_{i}} .

Using the schematic equation

\begin{matrix} (4) & \frac{d W_{i}}{d t} = Δ W_{i}, \end{matrix}

where $Δ W_{i}$ represents the weight changes given in (2). There is one ODE for each adjustable weight $W_{i}$ , and the right-hand sides are nonlinear functions of all weights. Comparing equations (3) and (4) shows that each unit of time $t$ in our differential equations corresponds to one presentation of the training set in the conventional algorithm, with a training rate $η = 1$ .

The Hessian matrix

H_{i j} = - \frac{\partial^{2} E}{\partial W_{i} \partial W_{j}}

is the Jacobian matrix for the differential equation. All explicit numerical solution schemes have a limiting step size for stiff stability, which is proportional to $1 / Λ$ , where $Λ$ is the largest eigenvalue of the Hessian matrix.

Modern stiff differential equation solvers are A-stable, so that the stability of the numerical solution is not limited by the computational stepsize taken.

The tolerance used for the numeric solver can be quite loose. Owens+Filkin chose $10^{- 2}$ , i.e., 1%. Owens+Filkin provide below performance table.

RMS error	Nonstiff / stiff	gradient descent / stiff
10%	1	1
5%	2	4
2%	5	>20

4. Classification and orientation. The results of Owens+Filkin (1989) have been confirmed by Alessio Tamburro (2019).

Fewer iterations or presentations of the data to the network are needed to reach optimal performance. [Though] The use of a ODE solver requires significantly greater computation time.

It remains remarkable that stiff ODE solvers are not used more often in the setting of neural networks. See Reply to: Neural Network Back-Propagation Revisited with Ordinary Differential Equations. The main points were:

Tolerances employed were very strict, too strict in my opinion, especially during initial integration
Completely incomprehensible that zvode, the "Complex-valued Variable-coefficient Ordinary Differential Equation solver" was used
A "switching" solver, i.e., one that can automatically switch between stiff and non-stiff would likely improve the results

Currently used neural networks have quite a number of weights.

Neural network	number of weights	reference
LLaMa	13 × 10⁹	Wikipedia
ChatGPT-3	175 × 10⁹	Wikipedia
ChatGPT-4	1000 × 10⁹	The decoder

Added 09-Jul-2023: Nice overview of some activation functions is here Activation Functions in Neural Networks. PyTorch's activation functions are listed here. Activation function in Twitter format as below. A separate discussion of the so called "rectifier function" is in Wikipedia. Also see Activation Functions: Sigmoid vs Tanh.