Constructing the Head on Learned Features, in JAX/Flax NNX

The explainer split a network into a backbone that builds features and a head that classifies, trained only the backbone, and built the head by hand. This is the implementation in JAX. The backbone is a small Flax NNX convolutional network with a normal training loop; the head is a constructed Yat vote placed on the frozen features, with no gradient steps. The whole point is to run both and see how little the head needs. Every number is from a real run; the script is scripts/jax_train_features.py.

Prefer a notebook? Run the whole thing on Kaggle: the same code, block by block, with the math written out and every figure reproduced from a real run.

The two stages, as two modules#

Where does the explainer’s split land in code? In two modules. The backbone is the part we train, two conv-and-pool blocks down to a 64-dimensional feature; the head is a plain linear layer for the trained baseline, and later a placed Yat vote for the constructed version.

import jax, jax.numpy as jnp, optax
from flax import nnx

class Backbone(nnx.Module):
    """Two conv-pool blocks to a 64-d feature. This is the only part that gets trained."""
    def __init__(self, rngs):
        self.c1 = nnx.Conv(1, 16, (3, 3), padding='SAME', rngs=rngs)
        self.c2 = nnx.Conv(16, 32, (3, 3), padding='SAME', rngs=rngs)
        self.lin = nnx.Linear(32 * 7 * 7, 64, rngs=rngs)

    def __call__(self, x):                               # x: [N,28,28,1]
        x = nnx.max_pool(nnx.relu(self.c1(x)), (2, 2), (2, 2))
        x = nnx.max_pool(nnx.relu(self.c2(x)), (2, 2), (2, 2))
        return nnx.relu(self.lin(x.reshape(x.shape[0], -1)))   # [N,64]

class Net(nnx.Module):
    def __init__(self, rngs):
        self.feat = Backbone(rngs); self.head = nnx.Linear(64, 10, rngs=rngs)
    def __call__(self, x): return self.head(self.feat(x))

The head you build instead of train#

The constructed head needs no module of its own; it is a function of the frozen features. Z-score them, place twenty k-means prototypes per class, and classify a test point by the nearest prototype under the Yat kernel. Not one gradient step is involved.

from sklearn.cluster import KMeans

def build_head(Ftr, ytr, Fte, per=20):
    mu, sd = Ftr.mean(0), Ftr.std(0) + 1e-6
    Ztr, Zte = (Ftr - mu) / sd, (Fte - mu) / sd
    W, vote = [], []
    for c in range(10):
        W += list(KMeans(per, n_init=2, random_state=0).fit(Ztr[ytr == c][:2500]).cluster_centers_)
        vote += [c] * per
    W = np.array(W, 'float32'); vote = np.array(vote)
    dot = Zte @ W.T
    d2 = (Zte ** 2).sum(1, keepdims=True) + (W ** 2).sum(1) - 2 * dot
    ker = (dot + 0.5) ** 2 / (d2 + np.median(d2) * 0.1)
    score = np.full((len(Zte), 10), -1e9)
    for c in range(10): score[:, c] = ker[:, vote == c].max(1)
    return score.argmax(1)

The surprise, before any training#

What happens if the backbone never trains at all? That is the cleanest version of the claim, so run it first. Take a backbone with random weights, never trained, and build the head on the garbage features that come out. You would expect nothing, and the trained linear head on those features confirms it: it sits at chance. But the constructed head sorts them anyway.

rnd = Net(nnx.Rngs(0))
con_rand = accuracy(build_head(feats(rnd, Xtr), ytr, feats(rnd, Xte)), yte)
print(con_rand)   # 72.7%

72.7% from a network that never saw a gradient. Random convolutional features are a known free lunch: even an untuned stack of filters keeps enough class structure for a head that only has to read it.

Train the backbone, compare the heads#

Now the ordinary training loop, but on the backbone only. Flax NNX makes the step a few lines: a loss, its gradient, an optimizer update.

net = Net(nnx.Rngs(0)); opt = nnx.Optimizer(net, optax.adam(1e-3), wrt=nnx.Param)

@nnx.jit
def step(net, opt, xb, yb):
    def loss(net):
        return optax.softmax_cross_entropy_with_integer_labels(net(xb), yb).mean()
    l, g = nnx.value_and_grad(loss)(net); opt.update(net, g); return l

for ep in range(6):
    perm = np.random.RandomState(ep).permutation(len(Xtr))
    for i in range(0, len(Xtr) - 256, 256):
        j = perm[i:i + 256]; step(net, opt, Xtr[j], jnp.asarray(ytr[j]))

After six epochs the trained head reaches 87.7%. Freeze the same backbone, throw the trained head away, and build a fresh Yat head on its features: 85.9%, with no head training at all. Training the classifier bought under two points; the other seventy-five came from the backbone.

trained = accuracy(net, Xte, yte)                       # 87.7%
con = accuracy(build_head(feats(net, Xtr), ytr, feats(net, Xte)), yte)   # 85.9%

So the curve runs from 72.7% on random features to 85.9% on trained ones, all read by a head that was never trained, always a point or two under the trained head and never more. (The explainer’s 74 / 83.2 / 85.7 are from its reference run, the PyTorch pipeline in scripts/construct_vs_optimize.py; the numbers on this page are this post’s own JAX run. The story is identical, the last digits move with the framework and the seed.) The accuracy lives in the representation. The classifier, and its edits, are furniture you place.

The strength of random convolutional features is from Saxe et al. (2011); the prototype head from Chen et al. (2019); Fashion-MNIST from Xiao et al. (2017); the Yat kernel from Bouhsine (2026). The conceptual companion is You Only Have to Train the Features.

@misc{bouhsine2026trainthefeaturesjaxflaxnnx,
  author       = {Bouhsine, Taha},
  title        = {Constructing the Head on Learned Features, in JAX/Flax NNX},
  year         = {2026},
  month        = {jun},
  howpublished = {\url{https://tahabouhsine.com/blog/train-the-features-jax-flax-nnx/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}