Model Editing

Jun 27, 2026

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 its depth.