Machine Learning
Machine learning research notes: interpretability, kernels, contrastive learning, attention, and representation geometry, each with live interactive visualizations.
-
I Removed Every MLP from Gemma 4 12B
Deleting every feed-forward branch removes 8.5 billion parameters and makes local Gemma 4 inference roughly three times faster. It also turns a capable language model into a machine that emits the audio control token forever.
-
A Velocity Ledger for Transformers, in JAX/Flax NNX
A runnable companion: the pre-norm Transformer block as a forward-Euler step, then the residual-stream velocity ledger as one line of Flax NNX state (mu = 0 recovers plain), the ngpt-lite retraction variant, best-val early-stopped training, and the depth telemetry (path length and turning angle per sub-update). Four parameter-matched char-level GPTs that tie on quality and split on dynamics: the ledger's residual-stream path is a third as long and half as sharp.
-
Transformers With a Velocity Ledger
A pre-norm Transformer's residual stream is forward Euler: x += Attn(norm x); x += MLP(norm x). So the whole integrator dictionary transfers, and the same question follows: does a velocity ledger in the residual stream do for a Transformer what it did for a ResNet? The answer splits. On quality, four variants tie. On dynamics, the ledger changes everything: the residual-stream path through depth gets dramatically shorter and straighter, reaching the same answer by a calmer journey. Same destination, gentler road.
-
Calibrating a Bounded Net, in JAX/Flax NNX
A runnable companion: build the matched Yat and ReLU MLPs in Flax NNX with the same softmax head, then measure their honesty. The reliability diagram and ECE, temperature scaling fit on a held-out split, NLL and Brier, and the two out-of-distribution channels, kernel-field magnitude versus softmax confidence, all in JAX with every number from a real three-seed run.
-
Building the Second Layer by Hand, in JAX/Flax NNX
A runnable companion: build a whole second feature layer by hand in JAX, on top of the hand-built first. Named min-AND combinations of layer-1 edges (junctions, continuations, bends, stripes) feed the same constructed Yat head, no training anywhere. It reproduces the flat rung: 83.3% at layer 1, 82.9% with both, 78.8% from relations alone, and counts the combinatorial wall of 224 pairwise and 4,630 three-way types where construction stops.
-
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.
-
Editing a Deep Equilibrium Network, in JAX/Flax NNX
A runnable companion: build the weight-tied Yat equilibrium operator in Flax NNX, then teach a class by appending rows to the readout (F untouched, exact) or into the shared dynamics (one paste, present at every depth). Measure local Jacobian slopes, temper the edit gain, audit 520 old fixed points, test multiple starts, watch a finite-prefix edit evaporate, and forget by masking. Every number is from a real run.
-
Skip Connections With Inertia, in JAX/Flax NNX
A runnable companion: the residual block as a forward-Euler step, then the momentum residual network as a Flax NNX module with one extra state, a velocity the blocks write into. Train both on the rings task a first-order flow cannot separate exactly, watch the training crystallize, and run the trained network exactly backward until floating point, amplified by 1/mu per layer, steals the past.
-
When 80% Should Mean 80%
A network hands you a probability with every answer, and the number is the part you act on. So when this series' bounded, self-explaining kernel network says 80%, is that a measurement or a mood? Five posts of evidence say it should be the honest one. This post puts that reputation through a lie-detector test, reliability diagrams, expected calibration error and temperature scaling against a matched ReLU MLP on Fashion-MNIST, and what the test found is the post.
-
How Far Down Can You Build?
One hand-built feature layer matched a trained backbone at 83.3% on Fashion-MNIST, and real networks are deep. Conveniently, the recipe for a second layer has been on the shelf for half a century: vision science says edges assemble into junctions, continuations, bends and stripes. This post takes the recipe down and follows it, builds layer 2 entirely by hand with every dimension still nameable in one sentence, and measures exactly where construction stops, and why.
-
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.
-
Edit One Operator, Edit Every Depth
One post taught and forgot classes by editing rows of a Yat network, with proofs that nothing else moved. Another melted the stack of layers into a single operator iterated to a fixed point. This is the collision. Every one of those editing proofs rested on a pasted row entering the score once, as one term in one sum, and in an equilibrium network there is no once: whatever you paste is applied at every depth and fed back into its own input, and every fixed point is free to drift. So did melting the stack melt the editability? This post pastes, deletes, and measures: every guarantee that survives is either proved inside the recursion or measured against the real run, fixed point by fixed point.
-
Your Skip Connection Is Half of Newton
A residual block x + F(x) is one forward-Euler step: depth is time, the block is a velocity, position moves directly. That is half of Newtonian mechanics. A planet does not update position from force; force updates velocity, velocity updates position, and that split is why orbits are stable. So what does the missing half cost a deep network? We let the physics make three predictions about trained networks, then check all three live in the page. One of them comes back stranger than we wrote it.
-
A Network That Is a Fixed Point, in JAX/Flax NNX
A runnable companion: build the Yat deep-equilibrium network in JAX/Flax NNX. One shared operator F(z;x)=tanh(A·φ_W(z)+Ux+z0), solved by damped iteration and trained with implicit differentiation. Measure residual convergence, local Jacobian norms, and sensitivity to initialization instead of assuming a global contraction. Plus a weight-tied maze operator that reaches 99.5% on grids larger than training by iterating longer.
-
Your Network Is a Stack of Layers. It Could Be a Fixed Point.
A deep network makes you choose its depth before you have seen the problem, and gives every layer its own weights. Share one Yat-kernel operator across depth and the stack becomes a single equation: the answer is the fixed point reached by iteration. On the measured test trajectories, the solver converges from widely separated starts and the local Jacobian norm stays below one. The same twenty-four prototypes describe every step, reaching 98.2% on two moons from 1,700 shared parameters.
-
A White-Box Kernel FFN in JAX/Flax NNX
A runnable companion: build a transformer whose feed-forward block is a finite learned-center kernel expansion. Train it on tinyshakespeare, then read each memory slot, attribute outputs exactly, edit one slot, and test peak kernel response as an abstention score.
-
The MLP Block Can Be a Kernel Memory
Replace an MLP activation with a kernel and its feed-forward block becomes an explicit learned-center expansion. Its slots can be read, attributed, and edited, but this architectural parameterization is not the classical representer theorem.
-
What Can a Weight Be?
A kernel is a spectral price list: it decides which functions are affordable, and regularization sets the budget. Compare Sobolev, Gaussian, spherical, and finite examples, then connect their eigenvalues to kernel ridge shrinkage and effective dimension.
-
Where a Weight Lives, in JAX/Flax NNX
A runnable companion: build the representer-theorem weight in JAX. A positive-definite kernel, the Gram matrix, a single linear solve for the coefficients, and the weight comes out as a combination of the data, f = sum alpha_i k(x_i, .). A linear weight cannot separate nested rings; the placed kernel weight does, read purely through the kernel as a similarity-weighted vote of the data.
-
Where Does a Weight Live?
A standard neuron's weight and its input never actually meet: one is a point you can see, the other an arrow off in its own space, joined only by a shadow. This is what a reproducing kernel Hilbert space fixes: it gives input and weight one shared address, where the optimal weight is built from the data itself and sits right next to it. Four interactive panels.
-
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.
-
Editing a Network by Hand, in JAX/Flax NNX
A runnable companion: build the prototype Yat-MLP in Flax NNX, then add a class by concatenating a few prototype rows and forget a class by masking them out, with no gradient steps. Class-incremental learning that matches a from-scratch build, and exact machine unlearning, both as array edits you can read. Every number is from a real run on Fashion-MNIST.
-
A Kernel's Price List, in JAX
Construct valid periodic kernel spectra, audit truncated RKHS norms for convergence, then solve kernel ridge regression and read regularization as spectral shrinkage, effective dimension, and a measured generalization curve.
-
Your Network Is a List of Pictures. You Can Edit It.
If a neuron is a labelled picture, a classifier is a list of them, and a list is something you edit. Add a class to a trained-free Yat-kernel network by placing twenty pictures, and it recognizes that class at 95% with zero gradient steps. Delete a class by removing its pictures, and it is forgotten exactly, the other classes untouched. Class-incremental learning with no penalty and machine unlearning that is instant and exact, both falling out of the architecture rather than bolted on.
-
The Yat-Kernel MLP in JAX/Flax NNX
Build a finite bank of Yat-kernel sections in JAX/Flax NNX, verify the kernel, train it on two moons and Fashion-MNIST, inspect exact prototype contributions, and test the initialization control that separates visible centers from noisy ones.
-
What a Finite Kernel Buys an MLP
Replace the activation with a finite bank of learned kernel sections. The resulting MLP exposes prototypes, exact layer-local contributions, measurable geometry, and the conditions those claims require, then tests the construction on arithmetic and Fashion-MNIST.
-
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.
-
Q and K Projections in JAX/Flax NNX
A runnable companion to Why Attention Needs Q and K Projections: build scaled dot-product attention with separate query and key projections in Flax NNX, pull the bilinear form B = W_Q W_Kᵀ out of the module, split it into a symmetric metric and an antisymmetric directed part, wire a toy induction head, add RoPE, and measure the low-rank budget and the gauge freedom, all in plain JAX.
-
Why Attention Needs Q and K Projections
The dot product in attention is not enough by itself. Without learned query and key projections, attention can only compare tokens in the residual stream’s native geometry. With a shared projection it learns a symmetric metric. With separate Q and K projections, the score becomes a learned bilinear form x_iᵀW_QW_Kᵀx_j: directional, role-aware, low-rank, and different per head. That bilinearity is what lets attention ask one kind of question and let tokens advertise another kind of answer.
-
The Prototype Readout in JAX/Flax NNX
A runnable companion to The Readout is a Convex Combination of Prototypes: read the columns of W_out as output prototypes in Flax NNX, measure the convex/conic/affine/linear regimes numerically, then build a Nadaraya–Watson kernel readout that is convex by construction (nonnegative weights that sum to one, a point that never leaves the prototype hull), with the nonnegativity-vs-positive-definiteness distinction checked in code.
-
The Readout is a Convex Combination of Prototypes
The second linear map in a transformer MLP is a dictionary of output prototypes, one per hidden unit. If the hidden activations are nonnegative and normalized, W_out reads the active neurons as a convex combination of output prototypes. Two independent constraints, nonnegativity and summing to one, sort the readout into four regimes: convex, conic, affine, and linear. This reframes the MLP readout as the same object that makes attention legible (a weighted sum over named basis elements), connects it to feed-forward key-value memories and modern Hopfield retrieval, and shows when a kernel makes it convex by construction.
-
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.
-
Cheap Attention in JAX/Flax NNX
A runnable companion to Cheap Attention: implement positive-feature linear attention in JAX and Flax NNX, watch the all-pairs ledger turn into a shared feature state, and see where the N×N matrix disappears.
-
Cheap Attention: Linear-Time Kernel Approximation
A 128K-token context creates billions of pairwise questions per attention head. But the N×N matrix is not the essence of attention; it is the receipt for an infinite feature map we never wrote down. Approximate that feature map with random features, reassociate the sum, and softmax attention becomes linear-time kernel attention.
-
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.
-
Self-Attention as Kernel Regression in JAX/Flax NNX
A runnable companion to Attention is Explainable Because it is a Kernel: build scaled dot-product attention from scratch in Flax NNX, prove in code that it is exactly a Nadaraya–Watson kernel smoother, watch the separate q/k projections break positive-definiteness numerically, swap the exp-dot-product kernel for Gaussian, Yat, and linear kernels to see which keep the weights a convex partition of unity, read the temperature as a kernel bandwidth, and train a single head end-to-end to route to a marked token.
-
What Attention Weights Can Explain
Self-attention has the normalized weighted-average form of a compatibility smoother. That exposes exact routing arithmetic, but it does not make attention weights causal explanations or guarantee a Mercer kernel on tokens.
-
What Activations Do to Geometry
ReLU, GELU, and their relatives enter a layer's Jacobian as an input-dependent row scaling. Here is when that scaling erases directions, when it merely distorts them, and what the usual repairs actually guarantee.