Random Features
Random features: approximating an infinite-dimensional kernel feature map with a finite sketch, the lever behind linear attention.
-
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 banks, and eight random prototypes drifting through training without ever becoming pictures.
-
How Many Random Neurons Buy a Trained One?
Freeze a bank of randomly initialized Yat units and train only the linear readout. The induced random-feature kernel is a Monte Carlo average; under finite variance its estimation error has the familiar square-root scaling. This post measures that exponent, counts how many frozen units buy each rung of an accuracy ladder, and then unfreezes the centers to measure what feature learning adds.
-
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.