Solving XOR: Why Networks Need a Hidden Layer

# Chapter 8: XOR, the problem one neuron can't solve, and the hidden layer that fixes it.
# Uses the Tensor class from Chapter 6, saved as tensor.py (see Chapter 7).
import numpy as np
from tensor import Tensor

# The staircase light: a switch at the bottom and one at the top. The light is on
# when exactly one of them is up. That's XOR.
X = np.array([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])   # (bottom, top), 1 = up
Y = np.array([[0.0], [1.0], [1.0], [0.0]])                         # 1 = light on


def cross_entropy(P):
    return -(Y * P.log() + (1 - Y) * (1 - P).log()).mean()


def fit(params, forward, lr=0.5, steps=2000):
    """Chapter 7's loop, full batch: forward, zero the grads, backward, step."""
    for _ in range(steps):
        loss = cross_entropy(forward())
        for p in params:
            p.grad = np.zeros_like(p.data)
        loss.backward()
        for p in params:
            p.data = p.data - lr * p.grad
    P = forward()
    return float(cross_entropy(P).data), P.data.ravel()


# === 1. One neuron does its best, and its best is a coin flip =====================
rng = np.random.default_rng(0)
w, b = Tensor(rng.normal(size=(1, 2))), Tensor(np.zeros(1))
loss, P = fit([w, b], lambda: (Tensor(X) @ w.T + b).sigmoid(), steps=5000)
print(f"one neuron:  loss {loss:.4f} (ln 2 = {np.log(2):.4f}), predictions {P.round(3)}")

# === 2. A two-layer network, built by hand ========================================
W1 = np.array([[1.0, 1.0],      # h1 counts the switches that are up
               [1.0, 1.0]])     # h2 ... minus 1: it only fires when both are up
b1 = np.array([0.0, -1.0])
w2 = np.array([1.0, -2.0])      # light = h1 − 2·h2
H = np.maximum(0, X @ W1.T + b1)
print("\nhand-built:  input   hidden (h1, h2)   h1 − 2·h2")
for x, h in zip(X, H):
    print(f"            {x.astype(int)}     {h}           {h @ w2:+.0f}")
p = 1 / (1 + np.exp(-(10 * (H @ w2) - 5)))       # scaled into probabilities
print("as probabilities:", p.round(4), " loss", round(float(-np.mean(Y.ravel() * np.log(p) + (1 - Y.ravel()) * np.log(1 - p))), 4))

# === 3. Let gradient descent find its own hidden layer ==============================
def mlp(hidden, seed, init="random"):
    rng = np.random.default_rng(seed)
    if init == "random":
        W1 = Tensor(rng.normal(size=(hidden, 2)))
        W2 = Tensor(rng.normal(size=(1, hidden)))
    else:                                           # every weight the same: a symmetric start
        W1 = Tensor(np.full((hidden, 2), 0.5))
        W2 = Tensor(np.full((1, hidden), 0.5))
    b1, b2 = Tensor(np.zeros(hidden)), Tensor(np.zeros(1))
    params = [W1, b1, W2, b2]

    def forward():
        H = (Tensor(X) @ W1.T + b1).relu()
        return (H @ W2.T + b2).sigmoid()
    loss, P = fit(params, forward)
    active = (np.maximum(0, X @ W1.data.T + b1.data) > 0).sum(axis=0)   # inputs each hidden unit fires on
    return loss, P, active, W1.data

loss, P, active, _ = mlp(4, seed=0)
print(f"\ntrained 2-4-1 net (seed 0): loss {loss:.4f}, predictions {P.round(3)}")

# === 4. It doesn't always work: 20 random starts per width ==========================
print("\nhidden units   solved   final losses of the first 10 starts")
for hidden in [2, 3, 4, 8]:
    runs = [mlp(hidden, seed) for seed in range(20)]
    solved = sum(np.all((P > 0.5) == (Y.ravel() == 1)) for _, P, _, _ in runs)
    print(f"{hidden:12}    {solved:3}/20   {[round(l, 3) for l, _, _, _ in runs[:10]]}")

# What the two kinds of failure look like, for 2 hidden units
for seed in [0, 2]:
    loss, P, active, _ = mlp(2, seed)
    print(f"seed {seed}: loss {loss:.4f}, predictions {P.round(3)}, inputs each unit fires on: {active}")

# === 5. A symmetric start never breaks the tie ====================================
loss, P, _, W1 = mlp(4, seed=0, init="same")
print(f"\nall weights 0.5: loss {loss:.4f}, predictions {P.round(3)}")
print("first-layer weights after training (every row identical):\n", W1.round(4))