The Three States of Information, in JAX

A state is only a state if some number can tell you which one you are in. The explainer named three that a representation passes through while training (random, organized, structured) and argued the loss plateaus are the transitions between them; what follows is four tiny JAX runs, one instrument each. The effective rank of the feature covariance falling from high rank toward C−1. The simplicity bias fitting a near-linear boundary before it wraps the moons. The class-mean cosines locking onto the Welch value −1/(C−1). And the contrastive loss splitting into two forces that turn out to resolve on two different clocks, in the wrong order.

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.

Everything below is plain JAX on tiny synthetic data (a few hundred points, a two-layer MLP), so the training loop and the linear algebra both fit on screen. A shared step function does the work:

import jax, jax.numpy as jnp, optax

def make_step(loss_fn, opt):
    @jax.jit
    def step(p, st):
        g = jax.grad(loss_fn)(p)
        up, st = opt.update(g, st, p)        # adamw needs params for decoupled decay
        return optax.apply_updates(p, up), st
    return step

State as spectrum: effective rank falls toward C−1#

Can a single number say which state a representation is in? The best candidate is the effective rank of the feature covariance, the participation ratio $(\sum_i \lambda_i)^2 / \sum_i \lambda_i^2$ of its eigenvalues. At initialization no direction is preferred yet, the spectrum is spread over many modes, and the effective rank sits high, far above C−1 (the random state; a finite ReLU network is not perfectly isotropic, so it starts high rather than at the full d). As training squeezes out within-class variation, a few modes dominate and the rest are ground down; the spectrum collapses onto the C−1 directions that separate the classes (the structured state). Organized is the in-between, where a handful of modes have emerged but the tail has not yet vanished.

def feature_spectrum(Z):
    """Z: (N, d) features -> sorted eigenvalues of the feature covariance."""
    Zc = Z - Z.mean(0)
    cov = Zc.T @ Zc / Z.shape[0]
    ev = jnp.clip(jnp.linalg.eigvalsh(cov), 0.0, None)
    return jnp.sort(ev)[::-1]

def effective_rank(ev):
    return (ev.sum() ** 2) / (jnp.sum(ev ** 2) + 1e-12)   # participation ratio

Train a small classifier on C=5 Gaussian blobs with a d=12 feature layer and watch the effective rank fall from 7.6 toward C−1 = 4:

ev0 = feature_spectrum(feats(params, X))
print("init effective rank:", float(effective_rank(ev0)))     # 7.56 (well above C-1)
# ... train ...
evT = feature_spectrum(feats(params, X))
print("final effective rank:", float(effective_rank(evT)))    # 3.86, close to C-1 = 4

Organized first: the simplicity bias#

Why does organized come before structured? Because networks fit low-order structure first: the mean and covariance of the data before its higher-order correlations, the distributional simplicity bias documented from small CNNs to LLMs (Refinetti et al., 2023; Belrose et al., 2024). On two moons it is visible to the eye: the decision boundary starts nearly linear (the coarse, low-order split) and only gradually grows the curvature needed to wrap the moons.

from sklearn.datasets import make_moons

Xn, yn = make_moons(320, noise=0.18, random_state=0)
Xn = (Xn - Xn.mean(0)) / Xn.std(0)

def net(p, x):
    h = jax.nn.relu(x @ p["W1"] + p["b1"])
    h = jax.nn.relu(h @ p["W2"] + p["b2"])
    return h @ p["W3"] + p["b3"]

loss_fn = lambda p: optax.softmax_cross_entropy_with_integer_labels(
    net(p, X), y).mean()

A cheap proxy for boundary complexity is its length on a grid: count how often the argmax prediction flips between neighbouring cells. Accuracy saturates fast; that length keeps climbing well after, the signature of higher-order structure arriving late:

def boundary_complexity(p, grid, shape):
    pred = net(p, grid).argmax(-1).reshape(shape)
    flips = jnp.abs(jnp.diff(pred, axis=0)).sum() + jnp.abs(jnp.diff(pred, axis=1)).sum()
    return flips / shape[0]

Structured: the simplex and the Welch value#

Where does the collapse land? Not at an arbitrary arrangement: in the balanced classification limit the class means converge to a simplex equiangular tight frame, every pair of class-mean directions at the same angle, and that angle is the Welch bound value $-1/(C-1)$, the most-separated arrangement possible for C points on a sphere (Papyan, Han & Donoho, 2020; Welch, 1974). The diagnostic is the distribution of off-diagonal class-mean cosines: it should concentrate onto that single value.

def equiangularity(Z, y, C):
    """Off-diagonal cosines between (centered) class means; target is -1/(C-1)."""
    mu = jnp.stack([Z[y == c].mean(0) for c in range(C)])
    mu = mu - mu.mean(0)
    mun = mu / (jnp.linalg.norm(mu, axis=1, keepdims=True) + 1e-9)
    G = mun @ mun.T
    iu = jnp.triu_indices(C, 1)
    cosines = G[iu]
    return cosines, -1.0 / (C - 1)            # (off-diagonals, Welch target)

cos, target = equiangularity(feats(params, X), y, C)
print("mean cosine:", float(cos.mean()), " target -1/(C-1):", target)
print("equiangularity std:", float(cos.std()))   # -> 0 as the frame tightens

Two clocks: alignment and uniformity#

Does the road through the states always run in the same direction? The supervised runs above went local-order-first: clusters, then the global frame. Contrastive learning is where that order can flip. Wang & Isola (2020) decompose the contrastive loss into alignment (positive pairs pulled together, the force behind local order) and uniformity (features spread evenly over the sphere, the force behind global order), and nothing says the two resolve at the same speed. In this run the augmentations are strong, and the race runs backwards: uniformity resolves fast, and while it spreads it flings positive pairs apart, so alignment climbs an early hump before a second, slower phase pulls positives back together. Global order arrives before local order, the mirror image of the explainer’s gas-liquid-crystal route, and both routes end at the same structured code.

def align_loss(za, zb):
    return jnp.sum((za - zb) ** 2, axis=1).mean()          # ‖f(x) − f(x⁺)‖²

def uniform_loss(z, t=2.0):
    d2 = jnp.sum((z[:, None, :] - z[None, :, :]) ** 2, -1)  # pairwise sq dist
    n = z.shape[0]
    mask = 1.0 - jnp.eye(n)
    mean_off = (jnp.exp(-t * d2) * mask).sum() / (n * (n - 1))
    return jnp.log(mean_off)                                # Wang–Isola uniformity

def contrastive_loss(p):
    za, zb = enc(p, A), enc(p, B)                           # two augmented views
    return align_loss(za, zb) + uniform_loss(jnp.concatenate([za, zb], 0))

The encoder ends z / ‖z‖, projecting onto the unit circle (a 2-D sphere) so the geometry is drawable:

def enc(p, x):
    h = jax.nn.relu(x @ p["W1"] + p["b1"])
    h = jax.nn.relu(h @ p["W2"] + p["b2"])
    z = h @ p["W3"] + p["b3"]
    return z / (jnp.linalg.norm(z, axis=1, keepdims=True) + 1e-8)

How the animations are computed#

All four are generated with Python, JAX, and matplotlib, every metric (eigenspectrum, effective rank, equiangular cosines, alignment/uniformity losses) recomputed live each frame from the current parameters: the spectrum walking from random to structured as effective rank falls from $d$ to $C-1$; the two-moons boundary picking up low-order structure first; the class-mean cosines descending to $-1/(C-1)$, the simplex; and the contrastive run resolving uniformity and then alignment on two clocks.

None is a benchmark. The objects are deliberately tiny so the state stays legible, but the four metrics (participation ratio, equiangularity, alignment, uniformity) are the same probes the explainer recommends logging during a plateau on a real run, and they cost a few lines each to add.

The point#

The explainer made the case in pictures and prose; this companion turns each state into a number you can watch fall or concentrate. Random is a spread, high-rank spectrum. Organized is local, low-order structure fit first: dominant modes emerging, a near-linear boundary. Structured is the C−1-mode frame, the equiangular simplex at −1/(C−1), the sphere uniform and aligned at once, however the two forces split the work of getting there. Three states, four instruments, and a plateau you can now name while it is happening.

References: three-states framing and the contrastive split from Wang & Isola (2020); neural collapse from Papyan, Han & Donoho (2020); the simplicity bias from Refinetti et al. (2023) and Belrose et al. (2024); saddle-to-saddle dynamics from Saxe, McClelland & Ganguli (2013); the Welch bound from Welch (1974).

@misc{bouhsine2026threestatesofinformationjax,
  author       = {Bouhsine, Taha},
  title        = {The Three States of Information, in JAX},
  year         = {2026},
  month        = {jun},
  howpublished = {\url{https://tahabouhsine.com/blog/three-states-of-information-jax/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}