A Perceptron from Scratch in Mojo

This notebook starts with hello world, then builds the smallest useful learning system we can inspect completely. We will train one perceptron to reproduce the logical AND gate without NumPy or a machine-learning framework.

# The notebook uses a Python kernel. This import registers the %%mojo cell magic.
import mojo.notebook

1. Check the toolchain

Every Mojo notebook cell contains a complete program with a main() entry point.

%%mojo
def main():
    print("Hello from Mojo. Let us teach one neuron.")

Hello from Mojo. Let us teach one neuron.

2. Make the training data explicit

The AND gate returns 1 only when both inputs are 1. Our four examples are the entire problem, not a sample of it.

x1	x2	target
0	0	0
0	1	0
1	0	0
1	1	1

%%mojo
def main():
    var x1: List[Float64] = [0.0, 0.0, 1.0, 1.0]
    var x2: List[Float64] = [0.0, 1.0, 0.0, 1.0]
    var targets: List[Int] = [0, 0, 0, 1]

    print("x1 x2 | target")
    for i in range(len(targets)):
        print(x1[i], x2[i], "|", targets[i])

x1 x2 | target
0.0 0.0 | 0
0.0 1.0 | 0
1.0 0.0 | 0
1.0 1.0 | 1

3. Define prediction

A perceptron first calculates activation = x1*w1 + x2*w2 + bias. Its step function returns 1 when the activation is zero or positive and 0 otherwise. With every parameter initialised to zero, every row initially receives the same prediction.

%%mojo
def predict(x1: Float64, x2: Float64, w1: Float64, w2: Float64, bias: Float64) -> Int:
    var activation = x1 * w1 + x2 * w2 + bias
    return 1 if activation >= 0.0 else 0

def main():
    var w1: Float64 = 0.0
    var w2: Float64 = 0.0
    var bias: Float64 = 0.0

    print("Initial prediction for (0, 0):", predict(0.0, 0.0, w1, w2, bias))
    print("Initial prediction for (1, 1):", predict(1.0, 1.0, w1, w2, bias))

Initial prediction for (0, 0): 1
Initial prediction for (1, 1): 1

4. Correct one mistake

The first row should be zero, but the untrained model predicts one. Let error = target - prediction. We update each weight with learning_rate * error * input and update the bias with learning_rate * error. Because both inputs are zero, this example changes only the bias.

%%mojo
def predict(x1: Float64, x2: Float64, w1: Float64, w2: Float64, bias: Float64) -> Int:
    var activation = x1 * w1 + x2 * w2 + bias
    return 1 if activation >= 0.0 else 0

def main():
    var x1: Float64 = 0.0
    var x2: Float64 = 0.0
    var target: Int = 0
    var w1: Float64 = 0.0
    var w2: Float64 = 0.0
    var bias: Float64 = 0.0
    var learning_rate: Float64 = 0.2

    var guess = predict(x1, x2, w1, w2, bias)
    var error = target - guess
    print("prediction:", guess, "error:", error)
    print("before -> w1:", w1, "w2:", w2, "bias:", bias)

    w1 += learning_rate * Float64(error) * x1
    w2 += learning_rate * Float64(error) * x2
    bias += learning_rate * Float64(error)
    print("after  -> w1:", w1, "w2:", w2, "bias:", bias)

prediction: 1 error: -1
before -> w1: 0.0 w2: 0.0 bias: 0.0
after  -> w1: 0.0 w2: 0.0 bias: -0.2

5. Train on all four rows

One pass through the dataset is an epoch. The program below repeats the update rule for four epochs, counts mistakes, and then prints the learned truth table. No model state is hidden: w1, w2, and bias are the complete model.

%%mojo
def predict(x1: Float64, x2: Float64, w1: Float64, w2: Float64, bias: Float64) -> Int:
    var activation = x1 * w1 + x2 * w2 + bias
    return 1 if activation >= 0.0 else 0

def main():
    var x1: List[Float64] = [0.0, 0.0, 1.0, 1.0]
    var x2: List[Float64] = [0.0, 1.0, 0.0, 1.0]
    var targets: List[Int] = [0, 0, 0, 1]
    var learning_rate: Float64 = 0.2
    var w1: Float64 = 0.0
    var w2: Float64 = 0.0
    var bias: Float64 = 0.0

    for epoch in range(4):
        var mistakes: Int = 0
        for i in range(len(targets)):
            var guess = predict(x1[i], x2[i], w1, w2, bias)
            var error = targets[i] - guess
            if error != 0:
                w1 += learning_rate * Float64(error) * x1[i]
                w2 += learning_rate * Float64(error) * x2[i]
                bias += learning_rate * Float64(error)
                mistakes += 1

        print("epoch", epoch + 1, "mistakes:", mistakes, "w1:", w1, "w2:", w2, "bias:", bias)

    print("\nLearned AND gate")
    print("x1 x2 | target prediction")
    for i in range(len(targets)):
        print(x1[i], x2[i], "|", targets[i], predict(x1[i], x2[i], w1, w2, bias))

epoch 1 mistakes: 2 w1: 0.2 w2: 0.2 bias: 0.0
epoch 2 mistakes: 3 w1: 0.4 w2: 0.2 bias: -0.2
epoch 3 mistakes: 3 w1: 0.4 w2: 0.2 bias: -0.4000000000000001
epoch 4 mistakes: 0 w1: 0.4 w2: 0.2 bias: -0.4000000000000001

Learned AND gate
x1 x2 | target prediction
0.0 0.0 | 0 0
0.0 1.0 | 0 0
1.0 0.0 | 0 0
1.0 1.0 | 1 1

6. Read the result

By the fourth epoch the model makes no mistakes. The learned boundary is 0.4*x1 + 0.2*x2 - 0.4 >= 0. Only (1, 1) crosses it.

This is genuinely learning, but it is still limited. One straight decision boundary cannot represent XOR. A later notebook can solve that by composing several perceptrons into a tiny multilayer network.