RSSAmplifier

Blog

Records of the !mmortal Data Scientist

Machine learning research notes by Taha Bouhsine — neural network interpretability, kernel methods, contrastive learning, attention mechanisms, RKHS, and transformer architectures. Long-form pieces with interactive visualisations.

tahabouhsine.comRSS feed ↗69 posts

Latest posts

Lazy Training a Yat Network in JAX/Flax NNX

A runnable companion to the lazy-training post: the Yat layer with its frozen softplus scalars, the one-line NNX filter that trains a readout while the bank stays frozen, the per-arm learning-rate bracketing, and the movement telemetry that caught the anti-lazy power law. With the run's own prototype trajectories: a Gram matrix crystallizing as neurons accumulate, two readouts racing on frozen…

How Many Random Neurons Buy a Trained One?

Freeze a bank of Yat units at random initialization, noise prototypes that will never move, and train nothing but the linear readout on top. This should not work, and the reason it does is one of the oldest facts in the kernel literature wearing a new activation: a frozen random layer is a Monte Carlo estimate of a kernel, one neuron per sample, converging at the square-root rate. This post…

Measuring Attention's Geometry in JAX/Flax NNX

A runnable companion to the geometry-of-attention post: the telemetry that exports a trained head's actual query and key vectors, the offline replay that recomputes every row's winner under rescaled queries, the linear program that asks whether any key sits inside the others' convex hull, and the census of who owns whom through training. Plus the toy laws animated from their real formulas:…

The Geometry of Attention Is a Choice of Kernel

An attention head is only ever evaluated at the tokens of the sequence, but its formula accepts any query vector at all. So between the tokens there is a whole continuous space, and the head silently carves it into territories: for every possible query, some key takes the top weight. This post draws that map for the two attention laws this series has trained head to head. The dot-product law…

Softmax-Free Attention in JAX/Flax NNX

A runnable companion to the compatibility-kernel post: the attention module where one branch computes softmax and the other computes the Yat kernel with no exponential anywhere, the parameter-matched training harness, the telemetry that measured our bounded-scores belief dead, and the checkpointed attention maps. Every number is from the real Kaggle runs.

The Kernel Between the Roles

The QK post in this series ended on a construction it refused to build: keep the query and key roles, but replace the bilinear-then-exponential score with a genuine Mercer kernel between them. This post builds it, with the kernel's full form: a per-head learned bias inside the square, the term the universality theorem requires, and a per-head learned softening. Because the kernel is nonnegative by…

An Error Controller for a Trained Net, in JAX

A runnable companion to the depth-on-demand post: the leapfrog classifier trained with lax.scan at fixed depth, the step-doubling controller that re-renders it to tolerance at inference, the honest work accounting (probes included), and the measured tol^(-1/3) power law. Every figure is rendered from the real Kaggle run.

Depth on Demand

The last post made depth a resolution: layers are time steps of a learned flow, and running more of them just renders the same trajectory finer. But every camera knows not to spend equal film on empty sky. This post gives a trained network the integrator's next tool, an error controller that chooses its own step size per input, with no retraining: the same weights, rendered to tolerance. The…

Reversible Backprop as a custom_vjp in JAX

A runnable companion to the memory post: the momentum block and its exact inverse, the custom_vjp whose backward pass reconstructs the trajectory instead of storing it, XLA's memory_analysis as the measuring instrument, and the (1/mu)^L float cliff reproduced in numpy float32. Every number is from the real Kaggle run.

Backprop Without the Memory

Training memory is a tax nobody chose: backprop must hold every activation of the forward pass hostage until the backward pass consumes it, so depth costs memory even when it costs little compute. This post spends the invertibility this series proved two posts ago: a momentum residual block can be run backward, so the backward pass can recompute the past instead of storing it. Measured on the same…

Building the Energy-Conserving Net in JAX/Flax NNX

A runnable companion to the energy-conservation post: the HNN pendulum field as the symplectic gradient of one learned scalar, the plain field model it beats on drift, and the leapfrog classifier whose residual block is a kick-drift-kick step of a learned potential, all as Flax NNX modules with lax.scan doing depth. Every figure is rendered from the real Kaggle run.

A Network That Conserves Energy

A pendulum never forgets its energy, and a trained network has no such number to forget. This post builds a residual network whose hidden state carries a conservation law by construction: the block is a symplectic step of a learned energy, so the quantity is held by the architecture, not encouraged by a loss. The learned pendulum keeps its energy to 0.6% where a plain field model leaks 36%, the…

Running the Survival Trial, in JAX/Flax NNX

A runnable companion to the survival-model trial: the Yat DeepSurv trunk in Flax NNX, the Cox partial-likelihood loss, an LR-fair per-model training loop with best-epoch selection, the concordance / integrated-Brier / time-dependent-AUC evaluation, and the classical baselines wired through sksurv and lifelines. Every figure is a real number from a real run across five datasets, with the prototypes…

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…

Solving It and Descending It, in JAX/Flax NNX

A runnable companion to the solve-vs-descend post: the Yat kernel and its Gram matrix, the exact kernel ridge solve via Cholesky, the same kernel as a Flax NNX module trained by AdamW with LR sweeps and best-epoch selection, the measured timing wall, minibatching through 511k rows, and the conv trunk the solve can never train. Every number is from the real Kaggle runs.

The White-Box Survival Model on Trial

A classical kernel machine is beautiful and needs a solve that will not minibatch or compose. A deep net trains on anything and its risk score is a fog. What if one thing had the training of a net and the theory of a kernel? We put a deep Yat-kernel survival model on trial across five real datasets against Cox, penalized Cox, and Random Survival Forest: it trains with plain gradient descent, lands…

Transformers With a Velocity Ledger

A pre-norm Transformer's residual stream is forward Euler: x += Attn(norm x); x += MLP(norm x). So D1's whole dictionary transfers, and the same question follows: does a velocity ledger in the residual stream buy in a Transformer what it bought in a ResNet? The answer splits. On quality, four variants tie. On dynamics, the ledger changes everything: the residual-stream path through depth gets…

One Kernel, Fitted Twice

Kernel methods gave us the theory everyone still wants back, and the field abandoned them over one procedure: the O(n³) solve over an n by n Gram matrix, which cannot minibatch, cannot scale, and cannot sit under other layers. So we took one Mercer kernel and fitted it twice: once by the classical exact solve, once by plain gradient descent on a bank of prototypes. The two machines agree, to a…

A White-Box DeepSurv, in JAX/Flax NNX

A runnable companion: the Cox partial-likelihood loss, a standard DeepSurv, and a Yat-kernel DeepSurv whose log-risk decomposes into prototype patients, all in Flax NNX. Concordance evaluation, exact convex attribution, cohort deletion, and OOD abstention as short array operations. Every number is from a real run on METABRIC breast-cancer survival.

A Risk Model That Names Its Reasons

A survival model tells an oncologist a patient is high-risk, and she has to act on the number without being able to ask why. What would it take for the risk score to name its reasons? We build a Yat-kernel DeepSurv on breast-cancer survival, match a standard DeepSurv on concordance, and get a risk score that decomposes exactly into the prototype patients this one resembles, a model you can read,…

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…

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…

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 the contraction certificate with power iteration and bisect one gain to restore it, audit the drift of 520 old fixed points, watch a layer-only edit evaporate, and forget…

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…

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…

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…

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…

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…

The Price List, in JAX/Flax NNX

A runnable companion to the price-list post: kernel ridge in JAX, the representer solve (K + lambda I) alpha = y, the RKHS-norm bill alpha^T K alpha, the effective dimension d_eff = sum lambda_k/(lambda_k + lambda) from the Gram spectrum, and a generalization sweep that draws the U-curve. Every number and every figure is from one analytic solve, no gradient descent.

Why Regularization Is a Price List

The representer theorem says the optimal weight is a sum over prototypes, but it does not explain why that sum generalizes. The answer is the RKHS norm: a price list that charges each prototype by its eigenvalue, and regularization is just tightening the budget. Four panels show the knob turning.

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 to its fixed point by damped iteration, trained not by backprop-through-iterations but by implicit differentiation, the adjoint (I−Jᵀ)u=∂L/∂z* run by the same contraction. Plus the weight-tied maze operator that extrapolates from 11×11 to 27×27 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 all of them instead, and the stack collapses into a single equation: the answer is the fixed point the state settles into. Training makes that operator a contraction, so the settling point is unique and reached from anywhere, the network decides…

A White-Box FFN: the Representer Theorem in JAX/Flax NNX

A runnable companion: build a transformer whose feed-forward block is a Yat kernel, so the FFN is exactly a representer sum over learned key-value memory slots. Train it on tinyshakespeare, then do four things you cannot do to an opaque ReLU FFN: read each memory slot, attribute an output to the slots that wrote it, edit one slot and watch generation change, and read off when the memory is out of…

The MLP Block Is a Representer Theorem

After the 3Blue1Brown attention video you can read half a transformer: you can see which token attends to which. The other half, the MLP block, stays a black box. But attention is legible because it is a kernel, a vote by similarity, and if you make the MLP a kernel too, its output becomes the same thing: a representer-theorem vote over learned prototypes. Then the whole transformer explains…

What a Weight Can Be, in JAX/Flax NNX

A runnable companion: compute the price list of a kernel in JAX. The eigenvalues are the kernel's spectral density, found with an FFT; the RKHS norm of a weight is a sum over them. A corner is affordable only under a Sobolev kernel, and the same numbers place the Yat kernel: universal and smooth, roomier than a Gaussian but not a Sobolev space.

What Can a Weight Be?

Once a kernel gives a weight a home, a second question follows: what is the weight allowed to be? Not all reproducing kernel Hilbert spaces are the same. A Sobolev space lets the weight have a sharp corner; a Gaussian's space forbids it; on normalized data the home is a sphere graded by spherical harmonics. A kernel is secretly a price list for roughness, and that list decides everything. Four…

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.

The Hand-Built Network, in JAX/Flax NNX

A runnable companion: build the training-free image classifier from the post in JAX. The feature extractor is pure JAX (Sobel gradients, orientation binning, patch pooling); the classifier is a Flax NNX module holding k-means prototypes that votes with the Yat kernel. Nothing is trained, and it reproduces the 83.3% on Fashion-MNIST.

You Don't Even Have to Train the Features

The last post trained a backbone and built the classifier by hand. This one builds the features by hand too: oriented-edge and corner detectors pooled over a grid of patches, the way computer vision worked for decades. Feed those to the same constructed Yat head and, with nothing trained anywhere, it matches the trained backbone on Fashion-MNIST point for point, within a couple of points of a…

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…

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.

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…

Your Neuron Is a Picture, in JAX/Flax NNX

A runnable companion: build the prototype MLP from the post in Flax NNX, train it on Fashion-MNIST, and watch the neurons. Pull the prototypes out as images, read a prediction as a vote over pictures, see the model abstain on out-of-distribution digits, check that random-init prototypes classify but stay noise, and track the prototypes migrating through a UMAP fit on the dataset as they train.

Your Neuron Is a Direction. It Should Be a Picture.

Why should a neuron store a direction when it could store a thing? A direction is not a referent you can point at, which is why MLPs are opaque. Put the Yat kernel where the activation was, train on Fashion-MNIST, and every neuron becomes a prototype that lives in pixel space, literally a picture, so the network reads its own predictions: this looks like that, no saliency method required.

The Yat-Kernel MLP in JAX/Flax NNX

A runnable companion to What a Finite Kernel Buys an MLP: build a layer whose unit is the Yat kernel instead of a linear map plus an activation, assert it is positive definite and nonnegative, write down its exact finite feature map, train it end-to-end on two moons with no activation function, and measure the lazy-loading sparsity, the bounded off-distribution response, the RKHS capacity, and the…

What a Finite Kernel Buys an MLP

Replace the activation function with a finite, explicit, positive-definite kernel, the Yat kernel, and an MLP stops being a stack of linear maps glued by a nonlinearity. It becomes a kernel machine, with locality, attribution, geometry, capacity control, and a feature map you can write down.