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
What crosses the wire in knowledge distillation besides the winning class? This experiment extracts a class-similarity kernel from teacher outputs and trains a student on pairwise relations alone—no labels, class names, or target probabilities. On Fashion-MNIST, the student recovers much of the label-trained geometry and approaches the spectrum of the transferred relation matrix.
-
Constructing the Fashion-MNIST Network, in JAX/Flax NNX
Train a small backbone and place its prototype head, then remove training entirely: implement fixed Sobel orientation channels, pool them into 343 named features, and classify with a constructed Yat head.
-
How Much of a Fashion-MNIST Network Can You Build by Hand?
Construct the prototype head on random and learned features, then replace the backbone with named edge and corner measurements. On Fashion-MNIST the zero-training pipeline reaches 83.3%, versus 85.7% for the matched trained model.
-
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
In these training runs, representation geometry moves through three recognizable regimes: random, organized into local clusters, and globally structured around separated class means. Interactive experiments test when loss plateaus coincide with those reorganizations—and when schedules change the order.
-
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
A label-similarity kernel can be turned into a target codebook by spectral embedding: retain its leading eigenmodes, scale by their square roots, and spend a finite dimension budget. Interactive experiments move that designed geometry from a simplex toward a taxonomy, then compare it with the class-mean and within-class spectra measured in trained representations.
-
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, and one geometric audit. Watch the losses organize the same 2D points while separating opposition, orthogonality, simplex packing, and statistical independence.