XOR Neural Network in Flux.jl

One of my favorite introductory neural network examples is XOR. It illustrates how linear neurons can learn nonlinear functions when combined with an activation function, which helped me intuitively understand part of why neural networks have gone on to become so successful.

The idea is to use a minimal neural network to predict the output of the exclusive OR logic operator (XOR), which outputs true only if exactly one of its inputs is true. XOR has a nonlinear decision boundary, so it can’t be reproduced exactly by any combination of purely linear functions.

The book Deep Learning by Ian Goodfellow, Yoshua Bengio, and Aaron Courville has a very nice worked example of this. Most code translations I’ve seen of it set up an uninitialized network and run a training procedure to show that the network can learn to predict XOR outputs with high accuracy. I wanted to replicate the example from the book directly, where the authors specified weights, biases, and an activation function, and then illustrated the linear algebra.

Not having tried Flux.jl in a few years and having spent some time in PyTorch land, I also wanted to reacquaint myself with Flux.jl’s updated API.

Linear Algebra Illustration

XOR can be defined in matrix form as follows:

X = Float32[
    0 0;
    1 0; 
    0 1;
    1 1
];

y = Float32[
    0; 
    1; 
    1; 
    0
];

As described above, this means that y is true if exactly one of the two inputs is true.

To show how this can be represented with two hidden neurons, an activation, and an output neuron, let’s work through the linear algebra from the book step by step, using the same solution the authors specified.

The hidden layer neurons are defined by weights W and bias c, and the output layer weights are defined by w. The output layer therefore implicitly has zero bias.

To start with, we multiply the input X by the hidden layer weights W:

W = [1 1; 1 1];
c = [0 -1];

w = [1; -2];

X*W

4×2 Matrix{Float32}:
 0.0  0.0
 1.0  1.0
 1.0  1.0
 2.0  2.0

Neural network biases are applied element-wise, so we broadcast the hidden layer bias c across the result of X*W:

X*W .+ c

4×2 Matrix{Float32}:
 0.0  -1.0
 1.0   0.0
 1.0   0.0
 2.0   1.0

Next, we apply the relu activation to the result. relu, or rectified linear unit, is a common neural network activation function (a transformation applied to the output of a neuron). If the neuron’s output is positive, relu outputs it directly. Otherwise, it outputs zero:

using Flux

relu(X*W .+ c)

4×2 Matrix{Float32}:
 0.0  0.0
 1.0  0.0
 1.0  0.0
 2.0  1.0

Finally, we multiply the above by the output layer weight vector:

relu(X*W .+ c) * w

4-element Vector{Float32}:
 0.0
 1.0
 1.0
 0.0

The combination of these operations recreates y exactly:

all(relu(X*W .+ c) * w .== y)

true

Okay, we’ve replicated the book example. Now how do we recreate this as a neural network, rather than a matrix representation of one?

Specifying Weights in Flux.jl

Let’s now define this as a neural network using Flux.jl. We won’t learn the weights and biases through training as with most other examples; instead we’ll specify the exact same ones used above, to illustrate how to map them appropriately.

In Flux.jl, as with many other neural network APIs, networks are declared as sequences of layers. Flux.jl calls these sequences Chains. For the XOR example, we’ll use linear Dense layers and the relu activation function to recreate the book example.

nn = Chain(
    Dense([1 1; 1 1], [0; -1], relu),
    Dense([1 -2], false)
);

Flux.params(nn)

Params([[1 1; 1 1], [0, -1], [1 -2]])

In most cases, we’d define our Dense layers by mapping an input dimension to an output dimension, optionally providing an initialization function. But here we’ve used an alternative interface to input our desired weights directly. Both options are documented, although it did take me a little while to get to the correct syntax shown above.

In the first layer, [1 1; 1 1] corresponds to W, [0; 1] to c, and we specify the relu activation function.

In the output layer, the weights [1 -2] are w, and we pass false to indicate that we don’t have a bias in this layer.

The biases are transposed relative to the original problem definition. We also need to transpose the input and output for the neural network to match the results of the linear algebra, because Flux.jl (and most neural network libraries) define their dense layer something like this: y = σ.(W * x .+ bias). That is, the first operation will be W*X instead of X*W.

nn(X')'

4×1 adjoint(::Matrix{Int64}) with eltype Int64:
 0
 1
 1
 0

With appropriate transpositions, we get the correct output from this network:

all(nn(X')' .== y)

true

And that’s it! We’ve confirmed that W, c, and w combine with the relu activation to produce an exact solution to the XOR problem, both in matrix form and when implemented as a simple neural network.

Other References

Post image courtesy of the Flux.jl project.