Ten Years Later: Rebuilding My 2016 Neural Network in Nx and Axon

Ten years ago, I wrote one of the most popular posts on this blog, Machine Learning: A Simple Neural Network. It was a toy problem, a single neuron with three inputs, and I derived the gradient by hand.provenance

Much has changed since then. At the personal level my career has gone through ups and downs, and my stack of choice has now switched to Elixir for almost everything I do. The Elixir ecosystem that struggled with number crunching in 2016 has grown to include Nx and Axon, and matured enough for serious machine learning and numerical computing on the BEAM.

So I think it's time to revisit that old post, and see what has changed in the last ten years and what we can do now with these new libraries in our hands. I'll start with Nx and then switch to Axon, to see how much of the math and complexity these libraries take off our shoulders.

Why do we care?

In 2026, we have access to large language models like Claude and GPT, which can perform complex tasks without us needing to understand the underlying mechanics. Those frontier models are built on the same principles as the simple neural network we will be exploring.

The 2016 neuron, in one paragraph

The original neuron we built in 2016 was a toy, chosen because it was easy to run and easy to verify. Four training rows, three inputs each, one output; the output always equals the leftmost input, and the whole exercise is for the neuron to figure out that rule on its own. There's a fifth row, [1, 0, 0], held out; the answer should be 1.

Everything was seeded, so the run reproduces. Same three starting weights every time:

[[-0.16595599]
 [ 0.44064899]
 [-0.99977125]]

If we run the untrained neuron on the held-out row and it answers 0.45860596 — a shrug, closer to "I don't know" than to 1. We need to train our neuron, through the training loop, 10,000 times: predict, measure the error, adjust.

# 2016, NumPy — the update I derived by hand
output = sigmoid(dot(inputs, weights))
error = training_outputs - output
adjustment = dot(inputs.T, error * sigmoid_derivative(output))
weights += adjustment

Those four lines are the whole idea of training:

• run the neuron(sigmoid)
• measure how wrong it is(error)
• push each weight toward a smaller error (adjustment)

We do this a little at the time over thousands of passes. Think about it as a dumb-down version of gradient descent. For our example we used a big round number, 10,000 iterations; here's where the three weights ended up:

# after training
[[ 9.67299303]
 [-0.2078435 ]
 [-4.62963669]]

Those weights are the rule the neuron worked out on its own: a strong positive weight on the first input, negatives on the other two, which adds up to "the first column decides the output."

Rebuild #1: the neuron in raw Nx, minus the derivative

Let's start with Nx. If you have used NumPy or JAX, Nx won't be completely unfamiliar to you. Nx is Elixir's library for numerical computing, built around tensors, multi-dimensional arrays of numbers, and the operations that run over it.

Historically, Elixir and the BEAM have been slow at number crunching next to languages like Python or C. Nx brings array programming to the BEAM directly, along with the two neat optimizations, automatic differentiation, so you get gradients without deriving them, and pluggable backends that compile the same code to fast CPU or GPU kernels through XLA.

Now we can express our original neuron example in Elixir. The entire thing fits in one module, and the only change from 2016 is that we don't have to derive the derivative by hand. The grad function does that for us.

defmodule Neuron do
  import Nx.Defn

  # forward pass: 3 inputs -> 1 output through a sigmoid
  defn predict(weights, inputs) do
    Nx.sigmoid(Nx.dot(inputs, weights))
  end

  defn loss(weights, inputs, targets) do
    preds = predict(weights, inputs)
    Nx.mean(Nx.pow(targets - preds, 2))
  end

  defn update(weights, inputs, targets, lr) do
    gradient = grad(weights, fn w -> loss(w, inputs, targets) end)
    weights - lr * gradient
  end
end

Few things to call out here:

• defn marks a numerical function Nx can compile and differentiate, and inside it functions like grad are available.
• predict/2 is the same sigmoid-of-a-dot-product as 2016.
• loss/3 is the mean squared error, the average of the squared gaps between prediction and target.mean-squared-error
• update/4 is the training step, which computes the gradient of the loss with respect to the weights and takes a step in that direction scaled by the learning rate lr.

For our training loop, it is a plain reduce over the same 10,000 iterations:

inputs  = Nx.tensor([[0, 0, 1], [1, 1, 1], [1, 0, 1], [0, 1, 1]], type: :f32)
targets = Nx.tensor([[0], [1], [1], [0]], type: :f32)
weights = Nx.tensor([[-0.16595599], [0.44064899], [-0.99977125]])
unknown = Nx.tensor([[1.0, 0.0, 0.0]])

trained =
  Enum.reduce(1..10_000, weights, fn _, w ->
    Neuron.update(w, inputs, targets, 1.0)
  end)

Neuron.predict(trained, unknown)

We can seed it with the identical 2016 weights and the untrained neuron should give the same value of 0.45860596:

iex> Neuron.predict(weights, unknown)
#Nx.Tensor<
  f32[1][1]
  [
    [0.4586059]
  ]
>

Now, if we run our training loop, the weights will be updated 10,000 times. The final prediction on the held-out input [1, 0, 0] should be very close to 1:

iex> trained = Enum.reduce(1..10_000, weights, fn _, w -> Neuron.update(w, inputs, targets, 1.0) end)
iex> Neuron.predict(trained, unknown)
#Nx.Tensor<
  f32[1][1]
  [
    [0.99999976]
  ]
>

TryPress Step once. Then raise the learning rate: it first overshoots but still settles; eventually each overshoot grows.

Rebuild #2: Axon, the neuron as a layer

Now if where doing neural networks for production applications, there is an even better option than NX alone: Axon. Axon is a library built on top of Nx that provides a high-level interface for defining and training neural networks. It abstracts away the low-level details of tensor operations and gradient calculations, allowing you to focus on the architecture of your model.

Everything we did in the Nx version can be expressed in Axon with just a few lines of code. The forward pass, loss calculation, and training loop are all handled by Axon, making it much easier to build and train neural networks:axon-dense

model =
  Axon.input("inputs", shape: {nil, 3})
  |> Axon.dense(1, activation: :sigmoid)

Axon.input/2 takes the name and shape of our in puts; the nil is the batch dimension, left open so the same model takes one row or ten thousand. Axon.dense(1, activation: :sigmoid) is the equivalent of our hand-crafted neuron: one output unit, a weight per input, a sigmoid on top.

To train the model, we can use Axon.Loop, which provides a convenient way to define the training loop. We specify the loss function, optimizer, and metrics, and then run the training loop with our data:

training_data = [{inputs, targets}]

trained_state =
  model
  |> Axon.Loop.trainer(:mean_squared_error, :sgd)
  |> Axon.Loop.metric(:mean_absolute_error)
  |> Axon.Loop.run(training_data, %{}, epochs: 10_000, compiler: EXLA)

Axon.predict(model, trained_state, unknown)

• Axon.Loop.trainer/3 takes the loss and the optimizer as atoms. :mean_squared_error is the same loss I wrote out by hand in the Nx version, :sgd is stochastic gradient descent.
• From those it builds the entire gradient loop for you: forward pass, grad, weight update, repeat. Axon.Loop.metric/2 bolts on mean absolute error, so you get a number to watch that isn't the loss itself.
• Axon.Loop.run/4 runs it and hands back an %Axon.ModelState{}, the trained parameters.axon-training-state

That's it. The Axon version is shorter, cleaner, and easier to read than the Nx version, and it still runs on the same underlying numerical kernels through XLA.

A word on performance

If you are planning to do any serious numerical computing with Nx or Axon, you should be aware of the performance implications of the backend you choose. By default, Nx operations run on the Nx.BinaryBackend, which is a pure-Elixir implementation and can be quite slow for large computations.

To take advantage of XLA's optimized kernels, you need to set the default backend to EXLA.Backend in your configuration

# config/config.exs
config :nx, default_backend: EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA

EXLA.Backend routes tensor operations through XLA, the same compiler PyTorch and JAX lean on; :default_defn_options does the same for compiled defn functions. exla-backend

Without XLA, the same code runs on a pure-Elixir interpreter, which is much slower, by several orders of magnitude.

10 years later

Elixir now has a full-featured numerical computing stack, and the same toy neuron I built in 2016 can be expressed in a few lines of code; it what I would argue is one of the most readable and maintainable modern languages.

Expressing the same ideas in a few lines of code shows how far the Elixir ecosystem has come in the last ten years. The libraries have matured, and the abstractions they provide allow us to focus on the problem at hand rather than the underlying mechanics, or reimplementing complex math.