What if next-token prediction wasn't a single forward pass, but a tiny optimization problem?
Introducing: nanoEBM a tiny transformer that learns to think harder by doing gradient descent on its own predictions.
You can start training on your Mac now - it comes < 400 lines
A few days ago I wondered if LMs could think harder without CoT prompting but by doing gradient descent on some objective function - this led me to energy-based models[1] where you learn some smooth and differential function to find a good y for any x. [1] https://atcold.github.io/NYU-DLSP20/en/week07/07-1/.
The core idea is to learn a E(x,y) where "low energy" is text we want and vice versa. During inference we get a prediction from the LM, compute its energy, take a gradient step to find a better prediction until we find a better answer without additional forward passes. In theory this could enable implicit CoT as every next-token prediction helps us descend down a learned E for a global sequence. For this implementation I don't learn globally but token(in this case char) by token.
Ok, but aren’t we just learning our dataset? E can be anything so we need use MC-sampled negatives (contrastive divergence) to approximate the gradient; pushing energy down on data, up on our samples. Because the model’s unnormalized(important later), computing E is exponentially hard; the curse of dimensionality.
nanoEBM is a tiny, from-scratch, ~10M-param char Transformer (67 vocab, 6 layers, 384 dim, 6 heads) with a linear energy head. Its based off @karpathy's minGPT with an additional 400 lines for our EBM. Next-token prediction is a tiny optimization where we compute energies per token once, unroll K small gradient steps on the logits to lower expected energy, then sample. Training defaults to cross entropy on these refined logits plus contrastive divergence from earlier-- to push down energy on real tokens, push up energy of logit‑space negatives (via Langevin against the batch’s fixed energy), with gradients flowing end-to-end through the unrolled thinking steps. https://huggingface.co/sdan/nanoebm-think-11m
On simple tasks like char-level Shakespeare I found it did assign low energy for real data and create a smooth landscape to minimize energy. However the data was too simple, was trained for too short, and our model size was too tiny to find marginal improvements in perplexity with increasing sizes in tokens trained. If you asked someone to think harder about Shakespeare I'm sure there's not much harder someone can think.
Hopefully I can run nanoEBM on harder more verifiable tasks in math, coding, where reasoning/thinking helps. The original energy-based transformer paper this is inspired by https://alexiglad.github.io/blog/2025/ebt/ trains on 7b model and found its ~35% more data efficient(which was my original inspiration here).
So why isnt EBMs more popular? My understanding is that its extremely expensive at test-time and too unstable at train time. Napkin math tells us for just a 4 "thinking steps" per token we would increase FLOPs on test-time by 1.4x compared to AR decoding which is insane at scale.
As AR LLMs need to answer one token at a time (forcefully?) reasoning, there could be another avenue to implicitly let models think with EBMs; since the energy/scoring function loosely lets us learn a verifier that determines whats good and how to find it.
nanoEBM runs on your mac/GPU easily, implemented minimal/extensibly in < 400 lines. A nice property of EBMs is you can combine multiple energy functions (product of experts) and normalize to get a usable distribution. Curious what people will train and merge!
Also big appreciation for high insanely quick @modal notebooks were shutdown and stood-up @ekzhang1 . Made training much easier.
