Representation Learning
Representation learning geometry: neural collapse, the simplex and Welch bound, contrastive objectives, and what a good latent space actually looks like.
-
Distillation as Kernel Transfer, in JAX/Flax NNX
A runnable companion: the five-run distillation experiment in JAX/Flax NNX. Train a teacher CNN, extract its class-similarity kernel S = E[softmax(z/T) softmax(z/T)ᵀ], train a student on nothing but pairwise relations (no labels, no soft targets), and measure it against the label ceiling and the random floor with a linear and a nearest-centroid probe. Every number is from a real run, with six GIFs that animate the kernel assembling, the temperature dial, the handoff, the spectrum inheritance, the probe race, and the inherited mistakes.
-
Distillation Is a Geometry, Not an Answer Key
Knowledge distillation has a standing puzzle: Hinton's student recognized 98.6% of the digit 3s in the test set after training on a transfer set with every 3 deleted. An answer key cannot do that, so what actually crosses the wire? This post gives dark knowledge a data type, a class-similarity kernel, and runs the experiment that isolates it: a student trained on nothing but pairwise relations, no labels, no soft targets, no class names, measured against the label-trained ceiling and the random floor. With live experiments: watch the kernel accumulate from single outputs, turn the temperature knob on how much geometry leaks, train a relational student in the page, and watch whose spectrum the student grows into.
-
Constructing the Head on Learned Features, in JAX/Flax NNX
A runnable companion: train a small conv backbone in Flax NNX, then on its frozen features build a constructed Yat head with no gradient steps and compare. The constructed head lands within a couple of points of the trained one, and even a random backbone's features sort at 73% while its trained head is at chance.
-
You Only Have to Train the Features
Leave a convolutional network's weights at their random starting values and build a Yat head on its features by hand: the trained head on that random backbone sorts at chance while the constructed one reaches 74%. On a properly trained backbone the constructed head reaches 83.2% against 85.7% for the trained one. The accuracy lives in the representation; the classifier, and its edits, are furniture you place. This maps the boundary between what you must optimize and what you can construct.
-
The Three States of Information, in JAX
A runnable companion to The Three States of Information: train tiny models in JAX and measure the three states directly: the feature-covariance spectrum collapsing from high-rank (random) to a C−1-mode frame (structured), the distributional simplicity bias that fits low-order structure first (organized), the neural-collapse simplex where class-mean cosines lock onto −1/(C−1), and the alignment/uniformity split of contrastive learning running on two separate clocks. Four live JAX visualizations, every number an eigenvalue or a loss.
-
The Three States of Information
Representations learned by a network pass through three states, like matter: random (high-entropy, no structure), organized (clusters, local order), and structured (a maximally-separated simplex, global order). The transitions between them are exactly the loss plateaus you see when training: the flat stretch is where the representation reorganizes before that reorganization shows up in the loss. Built from live in-browser training runs.
-
Latent on the Spectrum, in JAX
A runnable companion to Latent on the Spectrum: build a codebook as the spectral embedding of a label kernel in JAX (classical MDS with square-root eigenvalue scaling), watch a flat spectrum become the simplex and a graded one become the horseshoe, measure kernel-target alignment, split a representation into its between-class prototype frame and within-class information spectrum, and watch neural collapse grind the information to zero.
-
Latent on the Spectrum: Why Cats Sit Closer to Dogs Than to Cars
The regular simplex is the perfect codebook only when classes are strangers, and real labels are not strangers. A latent space is a lossy, finite-dimensional encoding of a label-similarity kernel: the codebook is the top eigenmodes of that kernel, the information rides in the modes below them, and the Welch bound sets the geometry of that channel. A follow-up to the Welch-bound post with live in-browser experiments: steer a codebook from simplex to taxonomy, spend a dimension budget, watch neural collapse grind the information spectrum to zero, read dark knowledge off a wandering feature, and see a structured codebook make better mistakes.
-
Auditing Latent Space Geometry in JAX
A runnable companion to the Welch-bound latent-space post: generate GIFs and implement the JAX metrics that tell you whether embeddings are collapsing, wasting rank, forming a simplex, or pressing against the Welch floor.
-
What Makes a Good Latent Space? The Welch Bound and the Simplex
The hidden codebook inside representation learning: why collapse happens, why opposition is a trap, why class means form a simplex, and why the Welch bound sets the best geometry when too many concepts share too few dimensions.
-
Organizing Randomness: Contrastive Learning in JAX
A block-by-block JAX + Optax implementation of six contrastive losses, each watched as a real animated GIF turning random 2D points into organized embeddings. The runnable companion to "Untangling the Moons."
-
Untangling the Moons: A Visual History of Contrastive Learning
Eight contrastive losses, twenty years of history, one interactive playground. Watch pair, triplet, InfoNCE, CLIP, SupCon, SigLIP, alignment+uniformity, and cosine→0 organize 2D points, and see which ones know when to stop.