PYTHON · NUMPY

Train for ternary: QAT matches FP32

Train a 2-H-2 MLP on XOR-quadrants with quantization-aware training and a straight-through estimator so that, evaluated with weights restricted to {-1,0,+1}, it reaches >=90% test accuracy — matching its FP32 twin.

01Challenge

Everyone "knows" you quantize a network AFTER training and pay for it in accuracy. That framing is the whole reason ternary looks like a compromise. It isn't. If the forward pass is ternary WHILE you train, the network learns weights that were always meant to be three-valued — and it lands as a peer of full precision, not a degraded copy. In this lesson you make a net whose every weight is -1, 0, or +1 match its FP32 twin on a non-linear task. That is the crux of the frontier claim: ternary is not an approximation you tolerate; trained right, it is FP32's equal, and it gets there having thrown away the multiplier array.

Predict firstCan a network whose weights are only {-1, 0, +1} ever match its full-precision twin's accuracy on a task a linear model can't solve?

02Model

Here is the sleight of hand that makes multiply-free networks real. A ternary weight matrix cannot hold a gradient — round() has zero slope almost everywhere, so ordinary backprop dies on contact. Quantization-Aware Training keeps two copies of every weight: a full-precision SHADOW that accumulates tiny gradient nudges, and its ternarized image {-1,0,+1} that the forward pass actually multiplies by. The Straight-Through Estimator is the trick that couples them: on the way forward you ternarize; on the way back you pretend the ternarize was the identity function and let the gradient flow straight into the shadow. The shadow drifts continuously; its ternary image snaps to three levels; and because you train THROUGH that snap, the network discovers a configuration of -1s, 0s, and +1s that solves the task outright. The one place people quietly cheat is the scale. If you let the ternarize threshold be a LEARNABLE parameter, the optimizer finds it can lower its loss by zeroing more and more weights — the zero-ratio runs away toward ~90% and the model collapses; this is exactly the failure BitNet-family QAT reports. A NON-LEARNABLE absmean scale, s = mean(|W|) recomputed every forward, removes that degree of freedom: the threshold simply tracks the weight distribution and self-regulates, so the network keeps a healthy mix of all three states. Do that, and the payoff is the frontier beat of this whole course — a net carrying nothing but {-1,0,+1} matches full precision on a task a linear model can't touch, having replaced every multiply with an add, a subtract, or a skip.

03Guided practice

Worked: The forward pass must consume ternary weights, so forward calls ternarize(W1) and ternarize(W2) and multiplies with those — never the raw shadow weights. ternarize(W) computes a scale s, divides, rounds, and clips to {-1,0,+1}. Watch what the scale controls: np.round(W/s) sends any weight with |W| < s/2 to 0 and the rest to ±1, so s is literally the rounding threshold measured in weight units.

Your turn: The starter ships s = 1.0. Our shadow weights are clipped to [-0.8, 0.8], so with s = 1.0 every |W| < 0.5 rounds to 0 — the matrix goes (nearly) all-zero, the net emits a constant, and accuracy sits at chance. Replace it with the non-learnable absmean, s = np.mean(np.abs(W)), recomputed fresh on every forward call so the threshold tracks the current weight scale. Do NOT hoist s out of the function or make it a trained parameter.

Independent: Run it. The STE is already wired (the backward pass sends da1 = dz2 @ W2t.T straight to the shadow weights, treating ternarize as identity), and the shadow weights are clipped each step for stability. When s adapts, the ternary net should print a test accuracy at or above 0.90 with a zero-ratio that settles into a healthy mix rather than collapsing to ~1.0.

04Feedback

PASS WHEN Prints PASS with ternary test acc >= 0.90 (about 0.985 at SEED=3) and a zero-ratio around 0.47 rather than ~1.0.

If your run fails, check:

FAIL with acc near 0.50 and zero-ratio near 1.00: you left s = 1.0. Shadow weights are clipped to [-0.8, 0.8], so with a threshold of 1.0 every |weight| < 0.5 rounds to 0 — ternarize returns an (almost) all-zero matrix, the net outputs a constant, and you get chance accuracy. Fix: s = np.mean(np.abs(W)) so the threshold tracks the weight scale.
Accuracy stuck around 0.75 with zero-ratio drifting toward ~0.9: the scale has effectively become learnable or stale — e.g. you computed s once outside forward() and reused it, so it no longer tracks the shrinking/growing weights. That is the zero-ratio collapse BitNet-family QAT warns about. Fix: recompute s = np.mean(np.abs(W)) fresh inside every forward call.
NaNs or accuracy frozen at chance from step one: you likely killed the gradient through ternarize (multiplying da1 by a rounding mask, or backpropagating through np.round). The Straight-Through Estimator treats ternarize as the identity on the backward pass — send dz2 @ W2t.T straight to the shadow weights, unmasked.

05Retrieve & space

Why must the absmean scale be recomputed inside every forward pass instead of being stored as a trained (learnable) parameter? Name the specific failure a learnable scale triggers.
round() has derivative zero almost everywhere, so gradients cannot pass through ternarize by ordinary backprop. Explain how the Straight-Through Estimator still lets the shadow weights learn, and which copy of the weights the gradient is applied to.
Transfer: you ship this trained net to an edge device that stores only the ternary weights, not the FP32 shadow. Which numbers does the device keep, which does it discard, and what arithmetic operation has it eliminated from the hot loop?

06Mastery & project

Passing proves you can train a network whose weights are only {-1,0,+1} to match a full-precision net on a non-linear task — that you can wire the STE (ternary forward, identity backward into an FP32 shadow) and use a non-learnable absmean scale to dodge zero-ratio collapse. You have shown, with your own computed number (~0.985 at SEED=3), that ternary is a peer of FP32, not a lossy afterthought.

This is module 2 of PAI-240: the AI crux. Module 1 established WHY the multiply is the enemy and what three states buy you; here you proved a ternary net can be TRAINED to full-precision parity. That parity is the license for everything downstream — the later modules cash it in, mapping these add/subtract/skip weights onto in-memory-compute hardware and the Institute's joule-ledger research program, whose scaling laws are working TARGETS, not settled results. The goal there is to match FP32 accuracy at a fraction of the energy that Horowitz's op-energy numbers imply a multiplier array must burn (FP32 multiply ~3.7 pJ vs add ~0.9 pJ @45nm) — and this lesson is the accuracy half of that bet, proven.

GO FURTHER · THE LIVE PLAYGROUND

Drive the live ternary policy ↗ — FP32 vs ternary vs log, trained live
Read the whitepaper ↗ — Multiply-Free Neural Computation (TR-2026-02)

← Absmean: rounding a net to three states BitVLA: the trit ships in production →

PYTHON · NUMPY

Train for ternary: QAT matches FP32

A 2->16->2 MLP on XOR-quadrants (label 1 when x*y<0 — the two off-diagonal quadrants, a problem no linear model solves). You train with QAT+STE: the forward pass ternarizes the weights, the backward pass updates FP32 shadow weights as if ternarize were the identity, and shadow weights are clipped each step. Everything is numpy and deterministic (SEED=3); a one-line np.seterr keeps numpy's spurious matmul FPE flags out of your output. Fill the one gap — the ternarize scale — and evaluate with the ternary weights. PASS when ternary test accuracy >= 0.90.

You read

the arrays and values already in scope

You change

the code you write in each cell

Fixed

the dataset and the checks that grade you