Data Loaders are critical tools to work efficiently with deep learning. They are the glue that bring the data to your machine learning model. Doing this efficiently and without errors is important for successful training of any model.

Datasets and DataLoaders are concepts that encapsulate the logic for preparing the data to be input in your model. In their basic form, they prepare small batches of data for ingestion by reshaping individual samples into matrix form. More involved applications may load images from files, augment them through flipws, crops, rotations.

The good news is that in Julia, MLUtils has you covered for all your data loading needs. In this blog post we’ll dive into how to create both Datasets and DataLoaders. We start out with a simple-as-possible example and then walk through a more realistic example.

Data loaders in easy mode

Let’s start by writing a simple as possible data loader to illustrate the core concepts of data loading using the DataLoader. The task we are looking at is somewhat related to image classification, but simplified.

Our data may be a bunch of images, like 1,000 of them. Each of the images is assumed to be in one of ten classes. A mock up of that data is something like

X_train = randn(Float32, 10, 10, 1_000)
Y_train = randn(1:10, 1_000)

Here, X_train are 1_000 samples of 10x10 matrices. These are mocked up as normally distributed Float32 data with shape (10, 10, 1_000), corresponding to (width, height, sample). Julia stores arrays in column-major, so in our case the data for width is consecutive in memory. The last dimension, the sample dimension varies slowest. The targets are mocked as 10_000 samples of random integers between 1 and 10, indicated by the 1:10 syntax.

A nice thing about Julia is that you often get away without defining custom classes. In our simlpe case, we can get away with defining our dataset as

data = (X_train, Y_train)

No custom class definition needed. The dataset is just a tuple of two vectors.

Now we move on to the data loader. The job of the data loader is to sample from our dataset. It should do batching, i.e. loading multiple (x,y) samples in one go, and random shuffling. For this simple dataset, we just need to create a DataLoader:

loader = DataLoader(data, batch_size = 3)

DataLoaders are iterable, and we can look at the first item like so:

julia> (x_first, y_first) = first(loader);
julia> size(x_first)
(10, 10, 3)
julia> size(y_first)
(3,)

When we iterate over the dataloader, we can use individual samples. Here is the iteration skeleton:

julia> for (x, y) in loader
            @show size(x), size(y)
       end
[...]
(size(x), size(y)) = ((10, 10, 3), (3,))
(size(x), size(y)) = ((10, 10, 1), (1,))

Note that in the last iteration, the DataLoader only returns vectors with a single sample. This is the loader watching out for us and stopping at the end of the data. We have 1_000 samples, and using batchsize=3, the last iteration can only feature a single sample. This is because 1000 / 3 = 333.33333 and the dataset is exhausted after the first sample in the last iteration.

We can confirm this by querying the length of the dataloader:

julia> length(loader)
334

That’s it. In conclusion, for a simple in-memory dataset creating a tuple data = (X, Y) is all you need. The DataLoader does the rest.

A more realistic example

Now let’s look at a more involved case where we want to read text data from a file and tokenize it. This sample is inspired by Andrej Karpathy’s NanoGPT video, github collab.

For this application, it’s convenient to encapsulate information on the DataSet into a struct. This struct collects the length of the text, a block size for data loading, dictionaries that map characters to tokens and vice-versa, and a vector of all the tokens. The members important for the data loading logic are block_size and data. While data holds the data vector itself, block_size defines the length of an individual sample.

In addition, we give the struct a constructors that reads all lines in the text file, creates an array of unique characters, and dictionaries to map from character to token and a dictionary to map from token to character. While the dictionaries are not necessary for the data loading logic to work, they are included here for completeness.

struct NanoDataset
    block_size::Int64               # How long an observation is. This is important for the dataloader
    ch_to_int::Dict{Char, Int64}    # Maps chars to tokens
    int_to_ch::Dict{Int64, Char}    # Maps tokens to chars
    data::Vector{Int64}             # The tokens

    function NanoDataset(filename::String, block_size::Int64)
        lines = readlines(filename)             # Read all lines
        _out = [c for l ∈ lines for c ∈ l]      # Create char array
        chars = sort!(unique(_out))             # Sort all chars occuring in the dataset
        push!(chars, '\n')                      # Add the \n character that we stripped when only looking at lines
        ch_to_int = Dict(val => ix for (ix, val) in enumerate(chars))   # Mapping of chars to int
        int_to_ch = Dict(ix => val for (ix, val) in enumerate(chars))


        all_tokens = [ch_to_int[s] for s in join(lines, "\n")]

        new(block_size, ch_to_int, int_to_ch, all_tokens)
    end
end

To create this dataset we run

ds = NanoDataset(FILENAME, 16)

Great, now moving on to the DataLoader. To make the dataloader work, we have to implement the numobs and getobs interface as described in the Documentation. This basically means that we have to make DataLoader know how to get the length of the dataset and how to get a single observation.

Fortunately, the MLUtils documentation tells us how to implement both. numobs just requires to specialize Base.length for our type. This function returns how many observations (that is the number of samples) there are in the dataset.

We implement this like

Base.length(d::NanoDataset) = length(d.data) - d.block_size - 1

which makes numobs work for our dataset:

julia> numobs(ds)
1115375

Next, we’ll get getobs working.

For the easy case, an individual sample of X_train can be accessed by X_train[:,:,42]. That is, array indexing is the access. For NanoDataset we need to define how to get a single observation. We do this by specializing Base.getindex on our dataset:

function Base.getindex(d::NanoDataset, i::Int)
    1 <= i <= length(d) - d.block_size - 1|| throw(ArgumentError("Index is out of bounds"))
    return (d.data[i:i+d.block_size-1], d.data[i+1:i+d.block_size])

So when we access ds through [] braces, we get two vectors: The input X and the target Y. X is a sequence of tokens of length block_size and Y is this same sequence, shifted by a single index.

I’d like to note here, that Base.getindex is a good place to get fancy. For image loading, this function may load data from a file. Or apply augmentations through Augmentor.jl. The sky (or ChatGPT) is the limit!

Now we can access samples in our dataset like

julia> ds[1]
([18, 47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14], [47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43])

Note that Y is just X shifted by one index, just as implemented in getindex. Before we define a dataloader, let’s implement one more function that allows us to load minibatchs:

MLUtils.getobs(d::NanoDataset, i::AbstractArray{<:Integer}) = [getobs(d, ii) for ii in i]

This function defines to dispatch calls where requested samples are given by i::AbstractArray{<:Integer} to multiple calls of getobs. For example, loading the first two samples returns this:

julia> getobs(d, [1,2])
2-element Vector{Tuple{Vector{Int64}, Vector{Int64}}}:
 ([18, 47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14], [47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43])
 ([47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43], [56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43, 44])

Fantastic, now we have everything in place to create a DataLoader. We can create a dataloader that loads 2 batches at the same time:

julia> dl = DataLoader(d, batchsize=2)
557688-element DataLoader(::NanoDataset, batchsize=2)
  with first element:
  2-element Vector{Tuple{Vector{Int64}, Vector{Int64}}}

julia> (x,y) = first(dl)
2-element Vector{Tuple{Vector{Int64}, Vector{Int64}}}:
 ([18, 47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14], [47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43])
 ([47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43], [56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52, 10, 65, 14, 43, 44])

The data loader returns a vector of lenth 2, each element of the vector a tuple of vectors. Ideally for deep learning, we’d like a matrix of size (16,2) though. Remember the first dimension is the sequence length, dimensions 2 is the batch dimension. The DataLoader can do this by setting collate=true:

julia> dl = DataLoader(d, batchsize=2, collate=true)
557688-element DataLoader(::NanoDataset, batchsize=2, collate=Val{true}())
  with first element:
  (16×2 Matrix{Int64}, 16×2 Matrix{Int64},)

julia> (x,y) = first(dl)
([18 47; 47 56; … ; 65 14; 14 43], [47 56; 56 57; … ; 14 43; 43 44])

julia> x
16×2 Matrix{Int64}:
 18  47
 47  56
 56  57
 57  58
 58   1
  1  15
 15  47
 47  58
 58  47
 47  64
 64  43
 43  52
 52  10
 10  65
 65  14
 14  43

Now the DataLoader returns the desired matrix. In addition, DataLoader supports shuffling and multithreading, just add shuffle=true, parallel=true to the parameter list.

Where to go from here

In this tutorial we looked at how to use DataLoaders in Julia through two examples. In the first example, we had in-memory data that can be passed as a tuple into the DataLoader. That just worked. In the second example, we had to implement the numobs, getobs interface to make it work. This is a just a bit more work, but allows to work with arbitrary custom data. The getobs implementation is where to hook in lazy-loading from disk, augmentations, and so on.

Now that we looked at some examples on data loading, here are some things to try next

• Try building your own Dataloader and use it in a Lux or Flux training loop.
• Play around with image augmentation using Augmentor
• Try using threads and measure the speed-up compared to single-threaded loading.

To be more precise, let’s say we have a quantity \(\rho\) with a profile that has a profile with a peak whose amplitude is parameterized by the parameter \(t\):

The plot below shows how the peak of the amplitude varies with \(t\):

We see that for \(t = 0.2\ldots2.0\) the peak decreases in amplitude and increases as \(t\) increases from 2.0 to 4.0.

To complicate things a bit, we would like to know how a quantity that is derived from this profile depends on \(t\). In particular we want to

1. Calculate \(\rho_t = \frac{\partial^2 \phi_t}{\partial x^2}\)
2. Then maximize \(\int\limits_{0}^{L_x} \phi_t(x)^2 \, \mathrm{d}x\)

The ρ profile used here is after all a manufactured solution to the Poisson equation for

Then we only have to ask Wolfram Alpha to give us

But in this tutorial the focus is on using automatic differentiation for such tasks. In particular, we can use the approach presented here when we lack an analytic solution for ϕ. We can also use the approach presented here when we lack an analytic expression for ρ, as long as we can generate ρ given a value for t.

So let’s draft a game-plan. First we choose a discretization of the domain [0:Lx] using Nx grid points, equidistantially spaced with \(\triangle x\). Then we pick a value of t and

1. Calculate \(\rho_t(x)\).
2. Invert the poisson equation to find \(\phi_t(x)\).
3. Calculate the integral \(I(t) = \int \phi_t(x)^2 \mathrm{d}x\).

Starting it off, we import relevant libraries, in particular Zygote, and define parameters for the grid and the profile:

using Zygote
using FFTW: fft, ifft
using LinearAlgebra
using BenchmarkTools

Lx = 2π
Nx = 128
Δx = Lx / Nx
n0 = 0.0
μ = π
σ = 0.25
t = 1.0

Since we later ask for the \(\partial_t I(t)\) we write a function that calculates \(I(t)\):

function int_prof(t)
    # This calculates I(t)
    xrg = Δx * (0:1:(Nx-1))
    ρ = ρ_prof(xrg, 1.0, t, μ, σ)
    ϕ = invert_poisson(ρ, Nx, Lx)
    ϕ = ϕ .- mean(ϕ)
    sum(ϕ.^2)
end

That’s it. Now we can call Zygote’s gradient function on int_prof(t) to take calculate \(\partial_t I(t)\). Now there are some more details to explore. In the next section we digress into Poisson solvers. Feel free to skip this section and head straight to the results to see how well AD performs on this relatively simple forward model.

Poisson Solvers

Given \(\rho_t\), we want to solve the Poisson equation

for \(\phi_t\) on the domain \(x \in [0:L]\). We don’t even pretend that we have any class and use periodic boundary conditions. We want the solution on the grid points \(x_i = i \triangle x\), where \(\triangle x = L_x / N_x\) and \(N_x\) denotes the number of grid points.

For this task wwe can either follow this excellent tutorial or use a spectral solver, described for example here.

A simple spectral solver for can be implemented like this:

# Simple spectral solver for Laplace equation
function invert_laplace_spectral(y::Array{Float64}, Nz, Lz)
    k = [1e100; 1:(Nz ÷ 2); -(Nz ÷ 2 - 1):-1] .* 2π/Lz
	y_ft = -fft(y) ./ k ./ k
    y_ft = y_ft .* [0; ones(Nz - 1)]
    d2y = real(ifft(y_ft))
end

The first line assembles a vector of wave numbers. The solution to Poisson’s allows an arbitrary offset. Setting the element of the wave numbers vector, that corresponds to the zero frequency, to a large number , k[1] = 1^100, we fix the mean of the solution to zero since we later divide by k.
We then take the Fourier transformation, divide by \(k^{2}\) and explicitly zero out the zero-mode. The syntax is a bit cumbersome, we perform a point-wise multiplication with an array, since Zygote does not support array mutation. Finally we take the real part of the inverse transformation.

A particlar issue with Zygote is that it can’t handle code that mutates arrays. The following text snippet shows a way of how to assemble an array and a way that will fail:

[0; ones(Nz - 1)]   # This will work with Zygote

tt = ones(Nz)       
tt[1] = 0           # This will not work with Zygote

Alternatively, we can implement a finite difference solver for Poisson’s equation. Below is the implementation for periodic boundary conditions adapted from herel

function invert_laplace(y,  Nz, Lz)
    Δz = Lz / Nz
    invΔz² = 1.0 / Δz / Δz

    A0 = Zygote.Buffer(zeros(Nz, Nz))
    for n ∈ 1:Nz
        for m ∈ 1:Nz
            A0[n,m] = 0.0
        end
    end
    for n ∈ 2:Nz
        A0[n, n] = -2.0 * invΔz²
        A0[n, n-1] = invΔz²
        A0[n-1, n] = invΔz²
    end
    A0[1, 1] = 1.0
    A0[1, 2] = 0.0
    A0[Nz,1] = invΔz²

    A = copy(A0)
    yvec = vcat([0.0], y[2:Nz])
    ϕ_num = A \ yvec
end

After calculating sample spacing, the code constructs a Matrix for the finite difference Laplace operator. Here we use Zygote.Buffer which allows us to use array mutation syntax. Again, code that is to be differentiated with Zygote can’t mutate arrays. After assembling the matrix we copy it into a matrix, augment the input vector and solve the linear system.

Differentiation through Poisson solvers

The code below shows how I set up the problem. I choose \(\phi\) and \(\rho\) as a manufactured solution. This way I can compare the numerical solution for \(\phi_t\) to ϕ_prof. A Gaussian solution is also a craving test case for a spectral solver. Since the Fourier transformation of a Gaussian is a Gaussian, all Fourier modes will be non-negative. This way one can pick up errors where a mode is at a wrong place in an array.

The functions int_prof_fd(t) and int_prof_sp(t) implement the three prong workflow described above. In particular their only argument is t. This way we can pass the whole workflow to Zygote.gradient and this way get the derivative of I with respect to t, $$\partial_t I$.

Now let’s set up the code:

ϕ_prof(xrg, n0, t, μ, σ) = n0 .+ sin(t).^2 .* exp.(-(xrg .- μ).^2 ./ 2.0 ./ σ ./ σ)
ρ_prof(xrg, n0, t, μ, σ) = -sin(t)^2 * exp.(-(xrg .- μ).^2 / 2 / σ / σ) .* (σ^2 .- (xrg .- μ).^2) / (σ^4)

Lx = 2π
Nx = 128
Δx = Lx / Nx
n0 = 0.0
μ = π
σ = 0.25
t = 1.0

function int_prof_fd(t)
    xrg = Δx * (0:1:(Nx-1))
    ρ = ρ_prof(xrg, 1.0, t, μ, σ)
    ϕ = invert_laplace(ρ, Nx, Lx)
    ϕ = ϕ .- mean(ϕ)
    sum(ϕ.^2)
end

function int_prof_sp(t)
    xrg = Δx * (0:1:(Nx-1))
    ρ = ρ_prof(xrg, 1.0, t, μ, σ)
    ϕ = invert_laplace_spectral(ρ, Nx, Lx)
    sum(ϕ.^2)
end

trg = 0.0:0.01:2π
phi2_fd = zeros(length(trg))
phi2_sp = zeros(length(trg))

phi2_grad_fd = zeros(length(trg))
phi2_grad_sp = zeros(length(trg))

for idx ∈ 1:length(trg)
    t = trg[idx]
    phi2_fd[idx] = int_prof_fd(t)
    phi2_grad_fd[idx] = gradient(int_prof_fd, t)[1]

    phi2_sp[idx] = int_prof_sp(t)
    phi2_grad_sp[idx] = real(gradient(int_prof_sp, t)[1])
end

As a sanity check we first compare the output of the forward model when using the finite difference and spectral solver.

For the resolution I chose, \(\triangle x\) is approximately 0.5 in this example, both solvers perform well.

A final thing to look at is speed. Using Benchmarktool we can evaluate the performance of backpropagating through either the finite difference and the spectral solver:

using BenchmarkTools
julia> @benchmark gradient(int_prof_sp, t)[1]
BenchmarkTools.Trial: 
  memory estimate:  141.64 KiB
  allocs estimate:  752
  --------------
  minimum time:     110.582 μs (0.00% GC)
  median time:      118.785 μs (0.00% GC)
  mean time:        139.315 μs (9.20% GC)
  maximum time:     12.583 ms (68.76% GC)
  --------------
  samples:          10000
  evals/sample:     1

julia> @benchmark gradient(int_prof_fd, t)[1]
BenchmarkTools.Trial: 
  memory estimate:  6.07 MiB
  allocs estimate:  170758
  --------------
  minimum time:     28.093 ms (0.00% GC)
  median time:      32.338 ms (0.00% GC)
  mean time:        35.562 ms (3.07% GC)
  maximum time:     104.050 ms (0.00% GC)
  --------------
  samples:          141
  evals/sample:     1

The mean time it takes to backprop through the finite difference solver is about 250 times larger than for the spectral solver. Given that the spectral solver is cheaper to evaluate, it doesn’t require to solve a linear system, this is not completely unexpected.

To get started with differentiable programming we are looking at a very simple example. We are given a density profile n on a finite domain:

The background density is \(n_0\) on which a Gaussian density peak, centered around \(z_0\) modulated. The amplitude of the peak is given by \(\sin(t)\). And we are interested on how the integral of the profile over the domain varies with t. Of course this can be calculated analytically:

To study the dependence of this integral on the amplitude parameter t we evaluate the derivative:

To find the amplitude t that maximizes N we can now just set \(\partial N/\partial t = 0\) and solve for t. But in a situation where we are not so lucky and have an analytic expression we may want to use automatic differntiation. Let’s do this in Julia.

Automatic differentiation in Julia

First things first - import Zygote to do automatic differentiation and the plots package as well as the error function erf:

using Zygote
using Plots
using SpecialFunctions: erf

Now we define a function that gives our density profile:

gen_prof(xrg, n0, t, z0, σ) = n0 .+ sin(t).^2 .* exp.(-(zrg .- z0).^2 ./ 2.0 ./ σ ./ σ)

For automatic differentiation to do its magic we need to numerically integrate the profile.

function int_prof(t)
    my_prof = gen_prof(z0:Δz:z1, n0, t, μ, σ)
    my_sum = sum(my_prof) * Δz
end

This function generates a vector which holds values of the profile at the points starting at \(z_0\) to \(z_1\), spaced regularly with \(\triangle z\). The parameters \(\mu\) and \(\sigma\) are captured from the context where the function is called from later. After generating the profile vector we numerically integrate the profile using the rectangle rule.

Now we have all the ingredients to calculate \(\partial N/\partial t\) using automatic differentiation. In code, we first define the necessary parameters and allocate vectors for the t’s where we want the derivative:

n0 = 1.0
σ = 0.2
x0 = -1.0
x1 = 1.0
Δx = 0.01
n0 = 1.0
μ = x0 + (x1 - x0) / 2.0
σ = (x1 - x0) / 10.0

# Perform numerical integration of the profile for a range of t's:
t_vals = 0.0:0.05:6.5
sum_vals = similar(t_vals)
sum_vals_grad = similar(t_vals)

To find the gradient, we first evaluate int_prof for a given value of t. Then we call gradient(int_prof, t) on this call. The return value is just the gradient \(\partial N / \partial t\). Thist functionality is just Zygote’s gradient doing it’s work but on a more complicated example than in the documentation. Note that gradient returns a tuple where the actual gradient is the first element.

for tidx ∈ tvals
    t = t_vals[tidx]
    sum_vals[tidx] = int_prof(t)
    sum_vals_grad[tidx] = gradient(int_prof, t)[1]
end

The gradient calculated using automatic differentiation can be compared to the analytic formula:

dNdt_anl = 2. * sqrt(2π) * σ .* sin.(t_vals) .* cos.(t_vals) .* erf(1 / √(2) / σ )

p = plot(t_vals, sum_vals_grad, label="∂ ∫n(z,t) dz / ∂t (autodiff).")
plot!(p, t_vals, dNdt_anl, label="∂ ∫n(z,t) dz / ∂t(analytical).", xlabel="t")

As shown in the plot, the analytical derivative and the one calculated by automatic differentiation are almost identical.

And plotting the error, we find differences of the order of \(10^{-8}\) to \(10^{-9}\).

Summary

In this blog post we looked at a very simple use-case for automatic differentiation. We have a mathematical expression that we know depends on a paramter and we wish to find the optimum value of this expression. Using automatic differentiation we now can calculate the gradient of that expression with respect to the parameter. This approach allows one to find the optimize the expression with resepct to the value.

A common procedure for this optimization is gradient descent, used in deep learning. Instead of a optimizing a neural network, automatic differentiation can be applied to arbitrary code. So we can optimize a much broader class of items than just neural networks.

Project structure

I am taking a starting point in the vanilla GAN implementation on the FluxML website. This implementation works and the trained generator indeed generates images that look indistinguishable from images belonging to the MNIST dataset. But how do we arrive there? Why are the learning rates chosen as \(\eta = 2 \times 10^{-4}\)? IS the leakyrelu the optimal activation function or does it perform on-par with relu in some regime? To answer these questions we need a code that quickly allows us to change these parameters.

And while we are at it, lets bundle the code together with its dependencies in a Julia package. This allows us to conveniently a package dependencies to the code. Taken together, the code and well defined dependencies make the behaviour reproducible. The Julia documentation gives a comprehensive introduction on packages here.

In order to run the code in the project I first checkout the code from github, then enter the repository and then execute the runme script:

$ git checkout https://github.com/rkube/mnist_gan.git
$ cd mnist_gan
$ julia --project=. src/runme.jl --activation=ADAM --train_k=8 ...

All packages installed in the project are local to this project and don’t interfere with packages installed in the general environment. This allows for example to specify for certain version numbers and will give us producibility of our results.

Code structure

The code is structued as a standard Julia project. The root folder layout looks like this

├── Manifest.toml
├── Project.toml
├── README.md
└── src
    ├── Manifest.toml
    ├── mnist_gan.jl
    ├── models.jl
    ├── Project.toml
    ├── runme.jl
    └── training.jl

The root folder contains Manifest.toml and Project.toml which include information about dependencies, versions, package names. More information is given in the Pkg.jl documentation.

The src folder contains all source codes files. In particular it contains a mnist_gan.jl file. This is named after the package name and in the simple case here only twofines the package as a module, includes all other modules and my two source files

module mnist_gan

using NNlib
using Flux
using Zygote
using Base:Fix2

# All GAN models
include("models.jl")
# Functions used for training
include("training.jl")
end #module

As additional structure I put the models in models.jl and training functions in training.jl.

Command line arguments

To quickly train the GAN with specific hyperparameters one can either read the hyperparameters from a configuration file or pass them through the command line. Here we do the second approach. To comfortably parse command line arguments I’m using (ArgParse.jl)[https://argparsejl.readthedocs.io/en/latest/argparse.html].

Condensing to only single argument, my code looks like this:

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    "--lr_dscr"
        help = "Learning rate for the discriminator. Default = 0.0002"
        arg_type = Float64
        default = 0.0002

args = parse_args(s)

That’s it. Now I can access the single command line arguments via args[lr_dscr].

Logging

Keeping track of the model performance while training is crucial when performing parameter scans. For the vanilla GAN alone I defined 10 parameters that can be varied. Letting each parameter assume only two distinct values this allows for 1024 combinations. Julia’s [logging facilities(https://github.com/JuliaLogging) provide means to systematicallylog model training for a large hyperparameter scan.

In particular, we can use TensorBoardLogger.jl. TensorBoard provides a visualization of training and includes numerous useful features, such as visualization of loss curves, displaying of model output images and more. To use TensorBoardLogger.jl in my code I have to include the module, instantiate a logger. Then I can easily log my experiments:

# Import the modules
using TensorBoardLogger
...
# Instantiate TensorBoardLogger
# Let's log the hyperparamters of the current run. 
dir_name = join(["$(k)_$(v)" for (k, v) in a])
tb_logger = TBLogger("logs/" * dir_name)
with_logger(tb_logger) do
    @info "hyperparameters" args
end

# Wrap the main training loop in a with clause to enable logging
lossvec_gen = zeros(Float32, args["num_iterations"])
lossvec_dscr = zeros(Float32, ["num_iterations"])

with_logger(tb_logger) do
    for n ∈ args["num_iterations"]e
        # Do machine learning ...
        ...
        # Code to log PNG images to tensorboard, inside the main training loop
        if n % args["output_period"] == 0
            noise = randn(args["latent_dim"], 4) |> gpu;
            fake_img = reshape(generator(noise), 28, 4*28) |> cpu;
            # I need to clip pixel values to [0.0; 1.0]
            fake_img[fake_img .> 1.0] .= 1.0
            fake_img[fake_img .< -1.0] .= -1.0
            fake_img = (fake_img .+ 1.0) .* 0.5
            # 
            log_image(tb_logger, "generatedimage", fake_img, ImageFormat(202))
        end
        # Log the generato and discriminator loss
        @info "test" loss_generator=lossvec_gen[n] loss_discriminator=lossvec_dscr[n]
    end # for
end #  Logger

First, I’m generating a string from all keys and values defined in the command line argument dictionary. Later this will allow me to filter these arguments. Then I’m logging the args dictionary, which contains the hyperparameters of the current experiment. Then I’m generating a fake image using the generator and log it as well. Here I need to clip the pixel values to [0.0; 1.0]. Since the Generator is trained on images with pixel values between -1.0 and 1.0 I need to transform the pixel space. Note that he last argument to the call in log_image encodes the layout of the fake_img array. I had to look up the available encodings via

?

Loss functions on-the-fly

To resolve the correct loss function from command line arguments I’m using the getfield method. To make it a little more convoluted, we also need to distinguish between loss functions that take an additional, tunable parameterr like celu, elu, leakyrelu and trelu, and loss functions who do not. The following code block shows how to map a string that encodes the function name to the actual function using getfield. To create a closure over an optional parameter I’m using Fix2. The code below is from models.jl

function get_vanilla_discriminator(args)
    ...
    if args["activation"] in ["celu", "elu", "leakyrelu", "trelu"]
        # Now continue: We want to use Base.Fix2
        act = Fix2(getfield(NNlib, Symbol(args["activation"])), Float32(args["activation_alpha"]))
    else
        act = getfield(NNlib, Symbol(args["activation"]));
    end

    return Chain(Dense(28 * 28, 1024, act), 
        ...);

I found out that can have an impact on performance how I pass the activation function as an argument to the dense layer. By passing only the function, the implementation of Dense handles how the activation function is applied to the linear transformation. This is how it should be. If I manually prescribe how to apply the broadcast I find slower performance:

julia> d1 = Dense(100, 100, act)
Dense(100, 100, relu)  # 10_100 parameters

julia> @btime d1(randn(Float32, 100, 100));
  163.863 μs (6 allocations: 117.33 KiB)

julia> d2 = Dense(100, 100, x -> act(x))
Dense(100, 100, #5)  # 10_100 parameters

julia> @btime d2(randn(Float32, 100, 100));
  3.041 ms (20016 allocations: 430.30 KiB)

So manually prescribing how to perform the broadcast is about 20 times slower. Instead, I let the code above return a function that Flux knows how to apply a broadcast on.

Running a parameter scan

Now we are set up to run a parameter scan. For this I generate runscripts where I vary my command line arguments. The resulting scripts look like this

#SBATCH things

cd /location/of/the/repo
julia --project=. --lr_dscr=0.0002 --lr_gen=0.0002 --batch_size=8 --num_iterations=10000 --latent_dim=100 --optimizer=ADAM --activation=leakyrelu --activation_alpha=0.2 --train_k=8 --prob_dropout=0.3 --output_period=250

Of course the arguments vary across the scripts. After crunching all the numbers, the log file directory is populated with the tensorboard log files. The next blog post will discuss how the results look like and how to pick the best hyperparameters.

References

[1] I. Goodfellow et al. Generative Adversarial Networks

Where Q is an orthogonal m-by-n matrix, i.e. the columns of Q are orthogonal \(q_{i} \cdot q_{j} = \delta_{i,j}\). R is a square n-by-n matrix. Using the decomposition, we can reconstruct the columns of A as

Internally, Julia treats QR-factorized matrices through a packed format [1]. This format does not store the matrices Q and R explicitly, but using a packed format. All matrix algebra and arithmetic is implemented through methods that expect the packed format as input. And the QR decomposition itself is calling the LAPACK method geqrt which returns just this packed format.

A computational graph where one may want to backpropagate through a QR factorization may look similar to this one:

While the incoming gradients, \(\bar{f} \partial f / \partial Q\) and \(\bar{f} \partial f / \partial R\) depend on \(f\), the gradients for the QR decomposition \(\bar{Q} \partial Q / \partial A\) and \(\bar{R} \partial R / \partial A\) are defined through the QR factorization. While these can be in principle computed through an automatic differentiation framework, it can be beneficial to implement a pullback. For one, this saves compilation time before first execution. An additional benefit is less memory use, as the gradients will be propagated through fewer functions. A pullback for the QR factorization may thus also aid numerical stability, as fewer accumulations and propagations are performed.

Formulas for the pullback of the QR factorization are given in [2], [3] and [4]. Pytorch for example implements the method described in [3], see here.

In this blog post, we implement the pullback for the QR factorization in Julia. Specifically, we implement a so-called rrule for ChainRules.jl. While you are here, please take a moment and read the documentation of this package. It is very well written and helped me tremendously to understand how automatic differentiation works and is implemented n Julia. Also, if you want to skip ahead, here is my example implementation of the QR pullback for ChainRules.

Now, let’s take a look at the code and how to implement a custom rrule. The first thing we need to look at is the definition of the QR struct

struct QR{T,S<:AbstractMatrix{T}} <: Factorization{T}
    factors::S
    τ::Vector{T}

    function QR{T,S}(factors, τ) where {T,S<:AbstractMatrix{T}}
        require_one_based_indexing(factors)
        new{T,S}(factors, τ)
    end
end

There are two fields in the structure, factors and τ. The matrices Q and R are returned through an accompanying getproperty function:

function getproperty(F::QR, d::Symbol)
    m, n = size(F)
    if d === :R
        return triu!(getfield(F, :factors)[1:min(m,n), 1:n])
    elseif d === :Q
        return QRPackedQ(getfield(F, :factors), F.τ)
    else
        getfield(F, d)
    end
end

For the QR pullback we first need to implement a pullback for the getproperty function. The pullback for this function only propagates incoming gradients backwards. Incoming gradients are described using a Tangent. This struct can only have fields that are present in the parameter type, in our case LinearAlgebra.QRCompactWY. This struct has fields factors and τ as discussed above. Now the incoming gradients would be \(\bar{Q}\) and \(\bar{R}\). Thus the pullback needs to map between these two. Thus the pullback can be implemented like this:

function ChainRulesCore.rrule(::typeof(getproperty), F::LinearAlgebra.QRCompactWY, d::Symbol) 
    function getproperty_qr_pullback(Ȳ)
        ∂factors = if d === :Q
            Ȳ
        else
            nothing
        end

        ∂T = if d === :R
            Ȳ
        else
            nothing
        end

        ∂F = Tangent{LinearAlgebra.QRCompactWY}(; factors=∂factors, T=∂T)
        return (NoTangent(), ∂F)
    end

    return getproperty(F, d), getproperty_qr_pullback
end

Notice the call signature of the rrule. The first argument is always ::typeof(functionname), where functionname is the name of the function that we want a custom rrule for. Following this argument are the actual arguments that one normally passes to functionname.

Finally we need to implement the pullback that actually performs the relevant calculations. The typical way of implementing custom pullbacks with ChainRules is write a function that calculates that returns a tuple containing the result of the forward pass, as well as the pullback \(\mathcal{B}^{x}_{f}(\bar{y})\). Doing it this way allows the pullback to re-use results from the forward pass. In other words, by defining the pullback as a function in the forward pass allows to easily use cached results. There is more discussion on the reason behind this design choice in the ChainRules documentation.

Returning to the pullback for the QR factorization, here is a possible implementation:

function ChainRules.rrule(::typeof(qr), A::AbstractMatrix{T}) where {T} 
    QR = qr(A)
    m, n = size(A)
    function qr_pullback(Ȳ::Tangent)
        function qr_pullback_square_deep(Q̄, R̄, A, Q, R)
            M = R̄*R' - Q'*Q̄
            # M <- copyltu(M)
            M = triu(M) + transpose(triu(M,1))
            Ā = (Q̄ + Q * M) / R'
        end 
        Q̄ = Ȳ.factors
        R̄ = Ȳ.T 
        Q = QR.Q
        R = QR.R
        if m ≥ n 
            Q̄ = Q̄ isa ChainRules.AbstractZero ? Q̄ : @view Q̄[:, axes(Q, 2)] 
            Ā = qr_pullback_square_deep(Q̄, R̄, A, Q, R)
        else
            # partition A = [X | Y]
            # X = A[1:m, 1:m]
            Y = A[1:m, m + 1:end]
    
            # partition R = [U | V], and we don't need V
            U = R[1:m, 1:m]
            if R̄ isa ChainRules.AbstractZero
                V̄ = zeros(size(Y))
                Q̄_prime = zeros(size(Q))
                Ū = R̄ 
            else
                # partition R̄ = [Ū | V̄]
                Ū = R̄[1:m, 1:m]
                V̄ = R̄[1:m, m + 1:end]
                Q̄_prime = Y * V̄'
            end 

            Q̄_prime = Q̄ isa ChainRules.AbstractZero ? Q̄_prime : Q̄_prime + Q̄ 

            X̄ = qr_pullback_square_deep(Q̄_prime, Ū, A, Q, U)
            Ȳ = Q * V̄ 
            # partition Ā = [X̄ | Ȳ]
            Ā = [X̄ Ȳ]
        end 
        return (NoTangent(), Ā)
    end 
    return QR, qr_pullback
end

The part of the rrule that calculates the foward pass just calculates the QR decomposition and returns the result of that. The pullback consumes the incoming gradients, \(\bar{R}\) and \(\bar{Q}\), which are stored as defined in the pullback for getproperty. The actual calculation for the case where the dimensions of the matrix A are equal, m=n, and the case where A is tall and skinny, m>n are different. But they can re-use some code. That is what the local function qr_pullback_square_deep does. The rest is mostly matrix slicing and the occasional matrix product. The implementation is borrows strongly from pytorch’s autograd implementation.

Finally, I need to verify that my implemenation calculates correct results. For this, I define two small test functions which each work on exactly one output of the QR factorization, either the matrix Q or the matrix R. Then I generate some random input and compare the gradient calculated through reverse-mode AD in Zygote to the gradient calculated by FiniteDifferences. If they are approximately equal, I take the result to be correct.

V1 = rand(Float32, (4, 4));

function f1(V) where T
    Q, _ = qr(V)
    return sum(Q)
end

function f2(V) where T
    _, R = qr(V)
    return sum(R)
end


res1_V1_ad = Zygote.gradient(f1, V1)
res1_V1_fd = FiniteDifferences.grad(central_fdm(5,1), f1, V1)
@assert res1_V1_ad[1] ≈ res1_V1_fd[1]

[1] Schreiber et al. A Storage-Efficient WY Representation for Products of Householder Transformations

[2] HJ Liao, JG Liu et al. Differentiable Programming Tensor Networks

[3] M. Seeger, A. Hetzel et al. Auto-Differenting Linear Algebra

[4] Walter and Lehman Walter and Lehmann, 2018, Algorithmic Differentiation of Linear Algebra Functions with Application in Optimum Experimental Design

To see this, let’s go the other way around. Given our computational graph we start at the output \(yy\) and calculate derivatives with respect to intermediate values. For convenience we define the notation

that is, the symbol under the bar is the variable we wish to take the derivative with respect to. Now let’s dive right in and take derivatives from the example. We start out with

The first expression is often called the seed gradient and is trivial to evaluate. In order to evaluate \(\bar{v}_5\) we had to use the chain-rule.

Continuing, we now have to evaluate \(\bar{v}_4\). Let’s do it first using the chain rule:

where we have used the chain rule and that \(v_6 = v_6(v_5(v_4), v_4)\). Looking at the computational graph, we see that we had to split the product using a plus-sign precisely at a position where there are multiple paths from the origin \(y\) to the target node \(v_4\), one via \(v_6\) and one via \(v_5\). Now as trivial as this example is, real-world programs are often much more complicated. And to calculate the partial derivatives, the formulas can become ever more complex.

Now the bar notation is here to make our life easier. Essentially, it gives us an easy way to replace chain-rule evaluation with local values that are stored in all child-nodes of the target node \(v_i\) in \(\bar{v}_i\). Put in another way, they are used as cache-values when traverseing the graph from right-to-left, as we do in reverse-mode AD. This is illustrated in the sketch below:

We see that there are two gradients incoming to the \(v_4\) node: \(\bar{v}_5 \partial v_5 / \partial v_4\) and \(\bar{v}_6 \partial v_6 / \partial v_4\). Assuming that \(\bar{v}_6\) and \(\bar{v}_5\) are already evaluated, we just need the partial derivatives \(\partial v_5 / \partial v_4\) and \(\partial v_6 / \partial v_4\) to find \(\bar{v}_4\):

And indeed, we have recovered the same rule as when applying the chain rule. When we continue to calculate derivatives \(\bar{v}_i\), we see that this notation becomes very handy. For example, using the chain rule we have

where we again need to use the product rule. Using the bar-notation however, the derivative becomes trivial

given that both \(\bar{v}_4\) and \(v_1\) have been evaluated. But in typical reverse-mode use this is the case, as we have completed one forward pass through the graph and we have traversed the barred values from right to left.

For completeness sake, let’s calculate the Reverse-mode AD evaluation trace as we did for the forward mode.

As a first step, we are evaluating all the un-barred quantities \(v_i\) before we begin the backward pass, where we evaluate all the \(\bar{v}_i\). We also see that we need the values of the child nodes to calculate the barred values. For example, to calculate \(\bar{v}_3\) we need the value of \(v_1\), which is a child-node of \(v_4\), when evaluating the partial derivative \(\partial v_4 / \partial v_1\).

We also see that this procedure would allow us to calculate the derivatives \(\partial y / \partial v_i\) for an arbitrary number of \(v_i\)s in a single reverse pass. This is the exactly the behaviour that makes reverse-mode automatic differentiation so powerful in context of gradient-based optimization. In that situation we want to take the derivative of a scalar loss function with respect to a possible large amount of parameters.

Reverse-mode AD in higher dimensions: Jacobian-Vector products

We are often working in higher dimensional spaces, hidden layers in multilayer perceptrons for example can have hundreds or even thousands of features. It is therefore instructive to take a look how reverse-mode AD works in this setting. For this we are looking at a new example:

The computational graph for the program \(y = g(f(x))\) is just

Now the formulas for the higher-dimensional case arise naturally from the ones derived in the previous section. Since we are discussing reverse mode, we start at the right. The bar quantities for each element of the first hidden layer is given by:

Using vector notation, this equation can be written as

Now \(\bar{f}\) is the vector that we get as an output from that calculation and will be propagated left-wards. We obtain it by calculating a vector-jacobian product. In practice, one usually does not calculate the jacobian matrix, but it is more efficient to calculate the vjp directly.

Moving leftwards to the graph, we use the same procedure to calculate \(\bar{x}\):

Combining these results we find

This equation is evaluated left-to-right, so that it requires to calculate only vector-matrix products.

As a final note, these vector-jacobian products are often called pullbacks in the context of automatic differentiation software. A notation used for them is

Let’s unpack the expression \(\mathcal{B}^{u}_{f}(\bar{y})\). The upper index denotes simply that the pullback is a function of \(u\), since it needs the values from the forward-pass at the node to be evaluated. The lower index \(f\) denotes that it depends on the mapping \(f\), since this is the mapping from \(x\) to \(u: f(x) = u\). Finally, the argument \(\bar{y}\) denotes that the pullback needs the incoming gradient from the right.

In practice, automatic differentiation software uses the chain rule to split the computational graph of a program into finer units until it can identify a pullback for a segment of the computational graph. In some cases this may be the desired way to work and in some cases, a user-defined backpropagation rule may be desireable.

Summary

We have introduced reverse-mode automatic differentiation. By starting with a gradient at the end of the computational graph, this mode allows to quickly calculate the sensitivity of an output with respect to an arbitrary large amount of intermediate values. To aid the necessary computations for this, the bar-notation has been introduced. Finally, we motivated that reverse-mode AD only needs to calculate vector-jacobian products and introduced the notation of pullback as the primitive object of reverse-mode AD. And finally, please check out the references below which I used for this blog-post.

References

[1] C. Rackauckas 18.337J/6.338J: Parallel Computing and Scientific Machine Learning

[2] S. Radcliffe Autodiff in python

[3] R. Grosse Lecture notes CSC321

[4] A. Griewank, A. Walther - Evaluating Derivatives - SIAM(2008)

To see why this is extremely useful, consider a traditional data fitting problem. We have a model \(f_\theta: \mathbb{R}^m \rightarrow \mathbb{R}^n\), where \(\theta\) are the model parameters and observations \(y_i \in \mathbb{R}^n\) taken at \(x_i \in \mathbb{R}^{m}\). To fit the model on the data we can use a loss function, for example the mean-squared error

and tune the parameters \(\theta\) in order to minimize \(\mathcal{L}\). This process is also called learning or solving an inverse problem.

A common method to tune the parameters \(\theta\) is gradient descent. Given a learning rate \(\eta\) one iterates

until \(\mathcal{L}\) is approximately at a local minimum. Applying the chain rule, we can expand the derivative as

To perform the optization procedure we need to calculate the derivatives described by the last term in the equation above. This is where automatic differentiation comes into play.

A Toy Example

To better understand what automatic differentiation does and how it works let us consider the toy example

In the following we will investigate how to calculate the derivative of \(f\) with respect to its inputs \(x_1\) and \(x_2\). In the context of the previous section, we can think of \(x_1\) as a parameter \(\theta\) that we wish to tone. While the method we present here is as simple as possible - real-world examples of tunable models often have thousands, millions, or even billions of parameters - it serves well to introduce the fundamental way in which automatic differentiation works. For the work below we are following [1] closely.

This mathematical formula for \(f\) can be decomposed into a succession of elementary functions, like \(+\), \(\sin\), \(\exp\), \(*\). That is, we can define intermediate variables like \(v_1(a,b) = a \cdot b\), \(v_2(a) = \cos a\) etc, which are only defined once. Each of these variables is also initialized by applying a simple expression to previous variables. Applying this procedure to our toy example gives us the following sequence:

We now define the so-called Tangential Derivative \(\dot{v_i} = \partial v_i / \partial x_1\). For the tangential derivative, we are fixing the denominator to \(x_1\) so that it describes the sensitivity of the functions output with respect to the given input. Let us calculate this derivative for the v’s now:

Note that we have to use \(\dot{v}_{-1} = 1\) and \(\dot{v}_{0} = 0\) in order to calculate \(v\dot{v}_1\). That is, the seed determines which part of the sum is non-zero. From the expressions above we can alos see that we will need to evaluate the derivatives starting at \(v_1\). That is, we begin at the start of the function, calculate simple derivatives and push this calculation through. Therefore this is also called forward mode automatic differentiation.

In practice, forward mode AD is often implemented by operator overloading. An implementation would need seed an initial derivative \(x_i = 1\) and then calculate the program and the derivatives simultaneously. This approach works for our toy example as we see from the evaluation trace shown below:

It is illustrative to visualize how the values and derivatives are propagated in the computational graph.

Here the intermediate values are shown in yellow and the gradient values are shown in turquoise. Note that both, values and derivatives are propagated together from the left, the input, to the output on the right. The implication of this is that forward-mode autodiff can calculate the sensitivities of all outputs y with respect to one input x in a multidimensional setting.

Summary

Automatic differentiation is a set of tools that allows to calculate the derivatives of computer programs that return numerical values. Here we looked at the so-called forward mode of automatic differentiation. It works by de-composing a program into simple functions for which we know the derivative. We can initialize a program and set the input variable for which we want to know the output’s derivatives for. Then we can calculate the function value and the derivative of the output with respect to the chosen value in one sweep.

This works well for situations where we have a small number of inputs and a large number of outputs. But in data-fitting problems we often end up with the reverse situation: We have a large number of outputs and few, or maybe only a single output. In these situations one may want to use reverse-mode automatic differentiation.

References

[1] A. Griewank, A. Walther - Evaluating Derivatives - SIAM(2008)

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

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

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:

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.

A common task in scientific computing is to solve a system of linear equations

with a matrix \(A \in \mathcal{R}^{d \times d}\) that gives the coefficients of the system, a right-hand-side vector \(b \in \mathcal{R}^{d}\), and a solution vector \(x \in \mathcal{R}^{d}\).

Instead of solving the linear problem directly, one often solves the preconditioned problem. We write this problem by inserting a one in the linear system and re-parenthesing:

\(\left( A P^{-1} \right) \left( P x \right) = b\) In this equation we have changed the coefficient matrix from \(A\) to \(A P^{-1}\) and the solution vector is now \(P x\) instead of \(x\). To retrieve the original vector, simply calculate \(x = P^{-1}y\). The matrix \(P^{-1}\) is called the preconditioner.

The dependence of \(P^{-1}\) on \(x\) is by choice. There is no deeper reason for why the preconditioner should depend on the initial guess for the iteraion. We choose to include this dependency here to foreshadow later applications, where such a dependency may be useful. In practice, this choice here is rather limiting. Since the matrix \(A_0 P^{-1}\) needs to be positive definite, and by construction only \(\widetilde{b}\) is positive definite, the method as constructed here is only valid for \(x \approx 0\).

The Gauss-Seidel method is an iterative method to find the solution of a linear system. Given an initial guess \(x_0\), update the entries of this vector following an iterative scheme and after some, or many, iterations, \(x_0\) solves the equation \(A x_0 = b\). Wikipedia has a nice and instructive page on this scheme.

A good preconditioner gives you a converged solution after fewer iterations. In other words, with a good preconditioner you are closer to the true solution of the linear system after N iterations than you are with a bad preconditioner. Choosing a good preconditioner depends on the problem at hand and can become a dark art. Let’s make it even more dark and train a neural network to be a preconditioner.

Modeling preconditioner as a neural network

As a first step, let’s model \(P^{-1}\) as a multi-layer perceptron (technically a single-layer):

with a weight matrix \(W \in \mathcal{R}^{d^2 \times d}\), a bias \(\widetilde{b} \in \mathcal{R}^{d^2}\), and a ReLU \(\sigma\). With \(x \in \mathcal{R}^{d}\) the matrix dimensions are chosen such that \(P^{-1}\) is of dimension \(d^2\). Simply reshape it to \(d \times d\) to have it act like a matrix.

Here I’m discussing a simple test case and am working with the Gauss-Seidel iterative solver. For real problem one would probably use a different iterative algorithm, but it serves as a proof-of-concept. Anyway, for Gauss-Seidel to work, the system coefficient matrix \(AP^{-1}\) needs to be positive-definite. How do we do this for our neural network? A simple hack is to consider only vectors whose entries are close to zero.

That way we get away with requiring only \(\widetilde{b}\) to be positive-definite. We get the last property by letting \(\widetilde{b} = b_0 b_0^{T}\) and sampling the entries of \(b_0\) as \(b_{ij} \sim \mathcal{N}(0, 1)\). In a similar fashion, we sample the entries of the weight-matrix \(W\) as \(w_{i,j} \sim \mathcal{N}(0,1)\).

How can we train a preconditioner

To make this preconditioner useful, the weights \(W\) and bias term \(b\) need to be optimized such that the residual after \(N\) iterations is as small as possible. And this needs to be true for all vectors from a training set. Using automatic differentiation, we can train \(W\) and \(b\) using gradient descent like this:

• Pick an initial guess \(x_0\) with \(x_{0,i} \sim \mathcal{N}(0, 0.01)\)
• Build the preconditioner \(P^{-1} = \sigma(W, x_0, \widetilde{b})\)
• Calculate 5 Gauss-Seidel steps. Let’s call the solution here \(y_5\).
• Un-apply the preconditioner: \(x_5 = P^{-1} y_5\)
• Calculate the distance to the true solution \(\mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} \left( x_{\text{true}, i} - x_{5,i} \right)^2\)
• Calculate the gradients \(\nabla_{W} \mathcal{L}\) and \(\nabla_{B} \mathcal{L}\)
• Update the weight matrix and bias using gradient descent: \(W \leftarrow W - \alpha \nabla_{W} \mathcal{L}\) and \(b \leftarrow b - \alpha \nabla_{b} \mathcal{L}\). Here \(\alpha\) is the learning rate

The approach here is to directly back-propagate from the loss function \(\mathcal{L}\), through the numerical solver, to the parameters of the preconditioner \(W\) and \(b\). We are not working with an offline training and test-data set, but the data is taken directly from the numerical calculations. This way we directly capture the reaction of the numerical solver to updates proposed by gradient descent.

The derivatives \(\nabla_\theta \mathcal{L}\) can be calculate using automatic differentiation packges, such as Zygote.

Implementation

# Test differentiation through control flow

# Use a iterative conjugate solver with preconditioner

using Random
using Zygote
using LinearAlgebra
using Distributions
using NNlib

I copy-and-pasted the Gauss-Seidel code from wikipedia:

function gauss_seidel(A, b, x, niter)
    # This is from https://en.wikipedia.org/wiki/Gauss%E2%80%93Seidel_method
    x_int = Zygote.Buffer(x)
    x_int[:] = x[:]
    for n ∈ 1:niter
        for j ∈ 1:size(A, 1)
            x_int[j] = (b[j] - A[j, :]' * x_int[:] + A[j, j] * x_int[j]) / A[j, j]
        end
    end
    return copy(x_int)
end

The block below sets things up.

Random.seed!(1)
dim = 4

# Define a matrix
A0 = [10.0 -1.0 2.0 0.0; -1 11 -1 3; 2 -1 10 -1; 0 3 -1 8]
# Define the RHS
b0 = [6.0; 25.0; -11.0; 15.0]
# This is the true solution
x_true = [1.0, 2.0, -1.0, 1.0]

# Define the size of the training and test set
N_train = 100
N_test = 10

# Define an initial state. Draw this from a narrow distribution around zero
# We need to do this so that the preconditioner eigenvalues are positive
x0 = rand(Normal(0.0, 0.01), (N_train, dim))

# Define an MLP. This will later be our preconditioner
# The output should be a matrix and we work with 2dim as the size for the MLP
W = rand(dim*dim, dim)
# For Gauss-Seidel to work, the matrix A*P⁻¹ needs to be positive semi-definite.
bmat = rand(dim, dim)
# We know that any matrix A*A' is positive semi-definite
bvec = reshape(bmat * bmat', (dim * dim))
# Now Wx + bmat is positive semi-definite if x is very small
P(x, W, b) = NNlib.relu.(reshape(W*x .+ b, (dim, dim)))
# Positive-definite means positive Eigenvalues. We should check this.
@show eigvals(A0*P(x0[0,:], W, bvec))

This function will serve as our loss function. It evaluates the NN preconditioner and then runs some Gauss-Seidel iterations. Finally, the iterative solution approximation is transformed back by applying \(P^{-1}\).

function loss_fun(W, bmat, A0, b0, y0, niter=5)
    # W - Weight matrix for NN-preconditioner
    # bmat - Bias vector for NN preconditioner
    # A0: Linear system coefficient matrix
    # b0: RHS of linear system
    # y0 - initial guess for Linear system. Strictly, this is x0. But we call it the same
    # assuming that P⁻¹x0 = x0.
    # niter - Number of Gauss-Seidel iterations to perform
    
    loss = 0.0
    nsamples = size(y0)[1]
    
    for idx ∈ 1:nsamples:
        # Evaluate the preconditioner
        P⁻¹ = P(y0[idx, :], W, reshape(bmat * bmat', (dim * dim)))
        # Initial guess
        # Now we solve A(Px)⁻¹y = rhs for y with 3 Gauss-Seidel iterations
        y_sol = gauss_seidel(A0 * P⁻¹, b0, y0[idx, :], niter)
        # And reconstruct x
        x = P⁻¹ * y_sol
        loss += norm(x - x_true) / length(x)
    end

    return loss / nsamples
end

Now comes the fun part. Zygote calculates the partial derivatives of the loss function with respect to its input, in our case \(W\) and \(b\). Given the gradients, we can actually update \(W\) and \(b\).

# Number of epochs to train
num_epochs = 10
loss_arr = zeros(num_epochs)
# Store the Weight and bias matrix at each iteration
W_arr = zeros((size(W)..., num_epochs))
bmat_arr = zeros((size(bmat)..., num_epochs))
# Learning rate
α = 0.005

for epoch ∈ 1:num_epochs
    loss_arr[epoch], grad = Zygote.pullback(loss_fun, W, bmat, A0, b0, x0)
    res = grad(1.0)
    # Gradient descent
    global W -= α * res[1]
    global bmat -= α * res[2]
    # Store W and b
    W_arr[:, :, epoch] = W
    bmat_arr[:, :, epoch] = bmat
end

Finally, let’s evaluate the performance

# This functions returns a vector with the residual of the iterative
# solution to the true solution at each step
function eval_performance(W, bmat, A0, y0, b0, niter=20)
    # Instantiate the preconditioner with the initial guess
    P⁻¹ = P(y0, W, reshape(bmat * bmat', (dim * dim)))
    y_sol = copy(y0)
    loss_vec = zeros(niter)
    for n ∈ 1:niter      
        # Initial guess
        # Now we solve (Px)⁻¹y = rhs for y with 3 Gauss-Seidel iterations
        y_sol = gauss_seidel(A0 * P⁻¹, b0, y_sol, niter)
        # And reconstruct x
        x = P⁻¹ * y_sol
        loss_vec[n] = norm(x - x_true) / length(x)
    end

    return loss_vec
end

# Get the residual at each Gauss-Seidel iteration
sol_err = zeros(20, num_epochs)
for i ∈ 1:num_epochs
    # Calculate the loss at each iteration, averaged over the training data
    sol_avg = zeros(20)
    for idx ∈ 1:N_test
        sol_here[:] += eval_performance(W_arr[:, :, i], bmat_arr[:, :, i], A0, x0, b0)
    end
    sol_err[:, i] = sol_here[:] / N_test
end

Results

A plot says many words, so here we go

The plot shows the average residual to the true solution vector as a function of Gauss-Seidel iterations. Training for 1 epoch, the residual decreases as a power law for all 20 GS iterations. The longer we train the preconditioner, the faster the residual decreases. Remember that we trained only for 5 iterations. But in the plot we see that preconditioner GS scheme proceeds at an accelerated rate of convergence, even after the fifth iteration. So for this toy example, the NN preconditioner performs quiet well.

Finally, a word on what we learn. Since we only take small, non-zero vectors, we update mostly the bias term and not the weight matrix. We can verify this in the code:

julia> (W_arr[:, :, end] - W_arr[:, :, 1])
16×4 Array{Float64,2}:
  9.2298e-7    3.56583e-6    9.63813e-6   -3.62526e-6
  7.02612e-6   1.83826e-6    4.87793e-6    1.06628e-5
 -1.18359e-5  -5.52383e-6   -1.00905e-5   -2.49149e-5
  4.2548e-6   -4.93141e-7   -5.37068e-6    1.46914e-5
 -4.80919e-6  -1.97348e-5   -5.14656e-5    1.82806e-5
 -2.99721e-5  -1.85244e-5   -4.37366e-5   -3.33036e-5
  5.28254e-5   5.17799e-5    0.000115429   8.27112e-5
 -1.70783e-5  -1.11531e-5   -8.03195e-6   -5.73307e-5
  6.09361e-6   3.89891e-5    8.93109e-5   -3.54379e-5
  5.58404e-5   5.60942e-5    0.000112264   6.67546e-5
 -8.3024e-5   -0.000129324  -0.00020506   -0.000139751
  2.75378e-5   3.47865e-5    1.79561e-5    0.000101882
 -2.89687e-6  -2.14357e-5   -4.965e-5      2.24125e-5
 -3.26029e-5  -3.11671e-5   -5.59135e-5   -4.08413e-5
  5.10635e-5   7.48976e-5    0.000101988   9.20182e-5
 -1.64133e-5  -1.88765e-5   -3.41413e-6   -6.0346e-5

The average matrix element of W has changed only little during learning, on average by about 0.00001. Looking at how much the elements of b have changed during training

julia> (bmat_arr[:, :, end]*bmat_arr[:, :, 1]') - (bmat_arr[:, :, 1]*bmat_arr[:, :, 1]')
4×4 Array{Float64,2}:
 -0.000674214   0.00163171   0.0012703    0.000217951
  0.00631572    0.00148005  -0.0105803    0.00372173
 -0.050246     -0.0453774   -0.0181303   -0.0511274
  0.011839      0.0128648   -0.00398132   0.00899787

we find that they have changed more, on average by a factor of 100 more than entries of W. But as discussed earlier, including x0 in the preconditioner is a modeling choice which one does not have to make.

Conclusions

To summarize, we proposed to use a Neural Network to accelerate an iterative solver for linear systems by acting as a preconditioner matrix. We propose that the weights of the neural network can be optimized by automatic differentiation in reverse mode. By putting the solver in the loop, the training and inference steps couple to the simulation in a very simple way.

One drawback of the method as written here is that we are limiting ourselfes to initial guesses \(x \sim 0\). This is due to the requirement of the Gauss-Seidel scheme that the linear system is positive definite. In more practical settings this can be circumvented by either using different parameters to be passed to \(P\) than \(x\). Alternatively one can use an iterative solver that doesn’t pose such restrictions, such as Jacobian-Free Newton Krylov or Conjugate Gradient etc.