By: Random blog posts about machine learning in Julia

Re-posted from: https://rkube.github.io/jekyll/update/2021/04/09/nn-precond-2.html

This post is a follow-up from the last one. But now we are focusing on solving
a non-linear system of the form

\[F(x) = 0\]

using a quasi-Newton method. You find the definition on
Wikipedia. Basically it
is Newton’s method but you replace the Jacobian with an approximation.
In this blog post we are considering a simple two-dimensional system.
Using Newton’s method for this example is trivial and leads to a converged
solution after only two or three steps, since it converges quadratically.
A quasi-Newton method on the other hand converges only linearly. And we
aim to improve on that convergence behaviour by using a machine-learned
preconditioner. Now, preconditioners are usually used to accelerate convergence of
linear solvers, like GMRES. And I don’t claim that they are generally useful for
non-linear systems. So we do it just to see if this works in principle.

Quasi-Newton iteration

We consider the simple two-dimensional system

\[\begin{align}
F_1(x_1, x_2) = x_1 + \alpha x_1 x_2 – s_1 & = 0 \\
F_2(x_1, x_2) = \beta x_1 x_2 + x_2 – s_2 & = 0
\end{align}\]

where \(\alpha\) and \(\beta\) are of the order 0.01 and the source terms are
\(s_{1,2} \sim \mathcal{N}(1, 10^{-2})\). Picking the source terms from a
Normal distribution allows us to consider a large number of instances of this
problem. Later we will use this to generate training and test sets.

Newton’s method can be used to find the root of $F(x)$. Starting out at an
initial guess $x_0$, it updates the guess by guiding it along its gradient:
\(F(x_0 + \delta x) \approx F(x_0) + J(x) \delta x \approx 0\). Here
\(J^{-1}(x)\) is the inverse of an approximation on the Jacobian Matrix whose
entries are is defined as \((J(x))_{ij} = \partial F_i / \partial x_j\).
Solving for \(\delta x\) gives the update \(\delta x = J^{-1}(x) F(x)\). Now the
Jacobian matrix of the system can be easily calculated. For the quasi-Newton
method we consider the approximation

\[\widetilde{J}(x) = \begin{pmatrix}
1 + \alpha x_2 & 1 \\
1 & 1 + \beta x_1
\end{pmatrix}\]

where the off-diagonal elements of the original Jacobian have been modified
as \(\partial F_1 / \partial x_2 = \alpha x_1 \approx 1\) and
\(\partial F_2 / \partial x_1 = \beta x_2 \approx 1\).

A quasi-Newton aims to find the root of \(F(x)=0\) using the same iteration scheme
but using an approximation to the Jacobian instead
\(\begin{align}
x_{k+1} & = x_k – \widetilde{J}^{-1}(x) f(x).
\end{align}\)

If you use the true Jacobian you have Newton’s method and the rate of
convergence is quadratic. If you don’t, you have a quasi-newton method and the
rate of convergence is linear. How can we accelerate convergence in this
case?

To accelerate convergence, we can use a Preconditioner \(P^{-1}\). This is
a invertible matrix that we squeeze into the update by inserting a one:

\[\begin{align}
J(x_k) \delta x_k &= -F(x_k) \Leftrightarrow \\
\left( J(x_k)P^{-1} \right) \left( P \delta x_k \right) &= -F(x_k)
\end{align}\]

The update from \(x_k \rightarrow x_{k+1}\) then needs to calculated in
two steps.

• Calculate \(\delta w = \left( \widetilde{J} \right)^{-1} P^{-1} \left(F(x_k) \right)\)
• Calculate \(\delta x_{k} = P^{-1} \delta w\).
• Update \(x_{k+1} = x_{k} + \delta_x\).

That is, we apply \(P^{-1}\) twice. A closed form of \(P\) is not required.
Next we will see how to implement \(P^{-1}\) as a neural network and
how to train it using differentiable code.

Neural-Network preconditioner

We parameterize the preconditioner with a Multilayer perceptron as
\(P^{-1}_{\theta}: \mathbb{R}^{m} \rightarrow \mathbb{R}^2\):

using Flux
dim = 2
Pinv = Chain(Dense(2*dim + 2, 50, celu), Dense(50, 50, celu), Dense(50, dim))

Here we allow for \(m > 2\) to pass additional inputs to the Network. We are
also using CeLU activation functions to avoid having a non-differentiable point
at \(x=0\) and we are choosing a simple 3-layer feed-forward architecture to start
out with.

While this gives us a parameterization we now have to find out how to train
the network. The task of the preconditioner is to minimize the residual
after \(n_\mathrm{i}\) iterations of Newtons method. The following loss
function tells us how well \(P^{-1}\) is doing over a number of examples
\(F(x) = 0\), where the RHS terms \(s_1, s_2 \sim \mathcal{N}\left(2, 0.01\right)\).

sqnorm(x) = sum(abs2, x) 


function loss2(batch)
# batch is a 2d-array. dim1: index within batch. dim2: indices the batch
    # Loss of the current sample. Defined as the norm of f(x) after nᵢ Newton iterations
    loss_local = 0.0
    # Batch-size
    batch_size = size(batch)[2]
    # Iterate over source vectors in the current mini-batch
    for sidx in 1:batch_size
        # Grab the current source terms 
        s = batch[:, sidx]
        # Define a RHS
        function f(x) 
            F = Zygote.Buffer(zeros(dim))
            F[1] = x[1] + α*x[1]*x[2] - s[1]
            F[2] = β * x[1] * x[2] + x[2] - s[2]
            return copy(F)
        end

        # run num_steps of Newton iteration for the current RHS
        sol = picard(x, x0, Pinv, s, 1)
        loss_local += norm(f(sol))
    end

    return log(loss_local) / batch_size
end

One has to be careful in evaluating the preconditioner within calls to the
Picard iteration. After each iteration, the input \(f(x)\) will decrease by
an order of magnitude. Neural Networks work best of the input is of order
one on the other hand. Thus we need to find a way to provide input of order
one to \(P^{-1}\) at each iteration. In the routine below the input is

• \(f(x) / \vert f(x) \vert_{2}\) – Input \(f(x)\), scaled to its L² norm
• \(s\) – The source terms in \(F(x) = 0\).
• \(\log \vert f(x) \vert\) – Gives the order of magnitude of the input
• \(i\) – The current iteration

The first term represents the desired input, re-scaled to be order unity, and
the logarithm \(\log \vert f(x) \vert\) estimates the order of magnitude of this
term. For example for \(\vert f(x) \vert = 10^{-5} \rightarrow \log \vert f(x) \vert \approx -11.5\). I also pass the source term \(s\) and the
current Picard iteration \(i\) into \(P^{-1}\).

function picard(f, x, P⁻¹, s, num_steps)
# Run num_Steps of Newton iteration
# P⁻¹: Preconditioner
# s: source terms in F(x) + s = 0

    for i = 1:num_steps
        # P⁻¹ receives the following inputs
        # f(x) / norm(f(x)) : Scaled input
        # s                 : source term
        # log(norm(f(x)))   : Order of magnitude of f(x)
        # i                 : Iteration number I
        δw = pinv(J(x)) * P⁻¹(vcat(f(x) / norm(f(x)), s, log(norm(f(x))), i))
        δx = P⁻¹(vcat(δw / norm(δw), s, log(norm(δw)), i))
        x = x - δx
    end
    return x
end

I am generating 100 training samples which I divide into 10 mini-batches
for training. Similarly, the test-set is 100 examples. Here I don’t use
a validation set but instead only consider the loss on the training set.
While training I am monitoring the loss on the training set. If it has not
decreased for 5 epochs, I’m halfing the learning rate.

In order to optimize the neural network as to give a smaller residual
after one iteration one needs to calculate the gradients of the loss
with respect to the parameters of the neural network. Zygote allows us
to do this using only a single call to gradient:

   
# Start training loop
loss_vec = zeros(num_epochs)
η_vec = zeros(num_epochs)
# Counts how many steps we haven't improved the loss function.
stall_thresh = 5

for epoch in 1:num_epochs
    # Randomly shuffly the training examples
    batch_idx = shuffle(1:num_train)
    for batch in 1:num_batch
        # Get random indices for the current batch
        this_batch = batch_idx[(batch - 1) * batch_size + 1:batch * batch_size]
        grads = Flux.gradient(Flux.params(P⁻¹)) do 
            l = loss2(source_train[:, this_batch])
        end
        for p in Flux.params(P⁻¹)
            Flux.Optimise.update!(p, -η * grads[p])
        end
        Zygote.ignore() do 
            loss_vec[epoch] = loss2(source_train[:, this_batch])
            η_vec[epoch] = η
            if mod(epoch, 10) == 0
                println("Epoch $(epoch) loss $(loss_vec[epoch])  η $(η)")
            end
        end
    end

    if (epoch > stall_thresh) && (loss_vec[epoch] > mean(loss_vec[epoch - stall_thresh:epoch]))
        global η *= 0.5
        println("    new η = $(η)   ")
    end
    if η < 1e-10
        break 
    end
end

Training the preconditioner is a bit tricky. I had to play quiet a bit with
the layout of the network, the learning rate, and the batch size to get useful
results. One setup that works is

x0 = [1.9; 2.2]
α = 0.03
β = -0.07

Pinv = Chain(Dense(2*dim + 2, 100, celu), Dense(100, 200, celu), Dense(200, 100, celu), Dense(100, dim))

num_train = 100
num_epochs = 100
num_batch = 10
num_test = 100

Let’s look at the loss on the training set:

For the first few epochs, there is no significant training. But around epochs
10 and then 20, there are significant drops. After 20, training proceeds uniform
and the loss function flattens out around epoch 60. Learning-rate scheduling
is essential here, it is halfed around the steep drops. Finally we end up with
an L² loss of about exp(-4)≈0.01 per sample.

We can evaluate the performance of the learned preconditioner by testing
it on unseen examples:

or idx in 1:num_test
    # Load the source vector from the test set
    s = source_test[:, idx]
    # Define the non-linear system with the current source vector
    function f(x)
        F = zeros(dim)
        F[1] = x[1] + α * x[1] * x[2] - s[1]
        F[2] = β * x[1] * x[2] + x[2] - s[2]
        return F
    end

    # Run 1 iteration with ML preconditioner, store solution after 1 iteration in x1_ml
    x1_ml = newton(f, x0, P⁻¹, s, 1)
    # Run 1 iteration without preconditioner, store solution after 1 iteration in x1_no
    x1_no = newton_no(f, x0, 1)

    # Run 5 more newton iterations, starting at x1_ml
    vsteps_ml = map(n -> norm(f(newton_no(f, x1_ml, n))), 0:5)
    # Run 5 more newton iterations, starting at x1_no
    vsteps_no = map(n -> norm(f(newton_no(f, x1_no, n))), 0:5)
end

The first iteration is performed using the preconditioner, followed by 5
iterations without preconditioning. This is compared to the norm of the
residuals in 6 iterations performed without preconditioner. The plot below
shows the convergence history of 100 test samples. Green line show ML-preconditioned iterations, red lines show the iteration history without
peconditioner.

We find that the Neural Network can act as a preconditioner in the first iteration.
The residual norm after one iteration shows large scatter, almost two orders of
maagnitude. After one accelerated initial step, the iteration continues with
linear rate of convergence for all test samples. That is at the same rate as the
un-preconditioned samples, shown in red here.

In summary, yes, we can use a machine-learned preconditioner to accelerate
Picard iteration to solve non-linear systems. It is a bit tricky to learn and
I only showed how to do the first iteration. So there are still many details that
have not been explored. I would also like to add that for the case here, Newton’s
method is the way to go. It has a quadratic rate of convergence – much faster than
the linear rate of Picard iteration. For the samples here, residuals are machine precision after about 3 to 4 iterations. But if that is not an option, using a
Neural Network can accelerate your Picard iteration.