By the end of this lesson you will have a real, runnable, full training loop in MLX — define a small two-layer MLP, generate a synthetic regression dataset (line + noise), train the model with SGD, watch the loss go down, run inference. About 30 lines of code, all of it production-shaped, none of it toy.
The point isn't to learn an interesting model — it's to feel every joint of the MLX training loop click into place. Every later track will assume you've felt this once.
The pieces
nn.Module — base class for any model. You subclass it, define __init__ with submodules and __call__ with the forward pass. There is no separate forward method — Python's __call__ is enough.
nn.Linear, nn.relu — the layer types and activation we need.
mlx.optimizers.SGD — the simplest optimizer. The .update(model, grads) method applies the gradients in place.
nn.value_and_grad(model, loss_fn) — the canonical pattern for getting both loss and gradient with respect to the model's parameters in one call.
mx.eval(model.parameters(), optimizer.state) — at the end of each step, materialize the updated weights and the optimizer state. This is the lazy-eval boundary that lets MLX fuse the step.
The whole loop
The code block below is the entire training pipeline. Read it once for shape — note how short it is — then run it. Watch the loss numbers drop. Note that the final inference call returns a value within shouting distance of the true target (2.0 × 1.0 + 0.5 = 2.5).
What to take from this
This is not a toy because the shape is real. A 70B model fine-tune in Track 4 will use exactly this skeleton — model subclass, nn.value_and_grad, optimizer.update, mx.eval at the boundary. The loop's complexity scales with what you train; its shape doesn't.
Code
End-to-end training loop — 2-layer MLP on synthetic regression·python
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
class MLP(nn.Module):
def __init__(self, in_dim=1, hidden=16, out_dim=1):
super().__init__()
self.l1 = nn.Linear(in_dim, hidden)
self.l2 = nn.Linear(hidden, out_dim)
def __call__(self, x):
return self.l2(nn.relu(self.l1(x)))
# Synthetic data: y = 2x + 0.5 + small noise
mx.random.seed(42)
N = 256
x = mx.random.uniform(-3, 3, (N, 1))
true_w, true_b = 2.0, 0.5
y = true_w * x + true_b + mx.random.normal((N, 1)) * 0.1
model = MLP()
mx.eval(model.parameters()) # materialize once before the loop
optimizer = optim.SGD(learning_rate=0.05)
def loss_fn(model, x, y):
return ((model(x) - y) ** 2).mean()
loss_and_grad = nn.value_and_grad(model, loss_fn)
for step in range(200):
loss_v, grads = loss_and_grad(model, x, y)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state)
if step % 50 == 0:
print(f'step {step:3d} loss {float(loss_v):.5f}')
# Verified output (2026-05-03):
# step 0 loss 9.29228
# step 50 loss 0.03690
# step 100 loss 0.02353
# step 150 loss 0.01696
# final loss: 0.013964570127427578
# Quick inference check
xt = mx.array([[1.0]])
print('model(1.0) ≈', float(model(xt).item()), ' (true 2*1+0.5 = 2.5)')
# Verified: model(1.0) ≈ 2.544 — within noise of true target.
Run the training loop. Confirm the loss drops and the final inference is close to 2.5. Then change two things, one at a time: (a) increase hidden to 64 and re-run, (b) drop the mx.eval(model.parameters(), optimizer.state) line and re-run. For (a), does the loss drop faster? For (b), does training still work, and how does memory behavior look on Activity Monitor? Two sentences total. The (b) experiment is the lesson 4 lazy-eval discussion made tactile.
Progress
Progress is local-only — sign in to sync across devices.