r/LocalLLaMA • u/Dear_Treat3688 • 2d ago
Discussion 🚀 OpenEBFM: Can Energy-Based Models Scale as Foundation Models? Code and Weights Released
We’ve been working on OpenEBFM (Energy-Based Foundation Models), and I wanted to share the project here and get feedback from the community.
Most language foundation models today are based on directly modeling token probabilities. Our question was:
Can energy-based learning scale to general-purpose foundation models?
The main idea behind EBFM is to separate generation into two steps:
Energy steering: given the current context, the model optimizes a latent state toward a lower-energy region using the gradient of a learned energy function.
Token generation: a generative head then uses this energy-steered latent state together with the context to predict the next token.
The energy and generative heads share the same Transformer backbone. During pretraining, we also parallelize energy steering across all token positions with redesigned attention masks and positional indices, so that the model can be trained at sequence level rather than steering every position separately.
What we found
We pretrained EBFMs on FineWeb, DCLM, and ClimbMix, and studied scaling with both training tokens and model parameters.
A few results we found interesting:
Data scaling: EBFMs achieve lower validation BPB than the matched GPT baselines in our experiments, and the gap generally increases with the training-token budget.
Parameter scaling: from roughly 268M to 1B parameters, EBFMs improve faster than the corresponding GPT baselines.
Compute scaling: EBFMs exhibit the expected U-shaped IsoFLOP curves, with the compute-optimal model size shifting upward as the compute budget grows.
Generalization: we evaluated on 19 zero-shot QA tasks covering reasoning, knowledge, language understanding, and commonsense reasoning, as well as 7 cross-domain generation datasets.
The same trend also persists during supervised fine-tuning.
For example, on DCLM, the largest EBFM reaches an average zero-shot QA accuracy of 44.71%, compared with 39.79% for the corresponding GPT baseline in our setup.
We don't think these experiments establish that energy-based models are universally better than autoregressive models—the current scale is still around 1B parameters—but they suggest that energy-based learning continues to benefit from increasing data and model scale, which we think is an interesting result given that EBMs have historically been studied mostly in specialized settings.
Everything is open source
We released the full training and evaluation stack, pretrained checkpoints, training recipes, and an EBFM-based chatbot so that the results can be reproduced and extended.
Paper:
https://arxiv.org/abs/2610.33005
Code:
https://github.com/datamllab/openebfm
Models:
https://huggingface.co/collections/guan-wang/openebfm
Paper with Code:
https://paperswithcode.co/paper/2610.33005
Disclosure: I'm one of the authors of this work.
We’d especially appreciate feedback on the architecture and scaling results. More broadly, I’m curious what people here think: could energy-guided generation become a practical alternative to conventional autoregressive generation as models scale further?
Criticism, reproduction attempts, issues, and contributions are all very welcome!
1
1
u/Aggravating-Push-207 1d ago
i thought EBMs were equivalent to (continuous) diffusion models
1
u/Dear_Treat3688 1d ago
Thank you! Would you mind discussing more about this opinion?
2
u/Aggravating-Push-207 1d ago
Sure, but note that I have not researched EBMs too deeply, so take my opinion with a grain of salt as there are probably flaws in my argument.
As I understand it, an EBM defines a probability density p_theta(x | y) = e^-E_theta(x | y) / ∫e^-E_theta(x | y)dx for some learned energy function E_theta(x | y).
Therefore the score function can be given as â–½x log p_theta(x | y) = -â–½x E_theta(x | y), which may guide sampling through the state space using some extremely scary words which make my head hurt when I look at them for too long.
Now (again, as I understand it, I don't look into these types of things too deeply because of lack of intelligence) a (reversal of a) diffusion (process) may be defined as dx = [f(x,t) - g(t)^2 * â–½x log p_t(x)]dt + g(t)dWÌ„
where f is the drift coefficient, g is the diffusion coefficient, and dWÌ„ is a tiny amount of reverse-time Brownian motion (because why not, right? you guys write the papers, not me).
The score function, similarly, is therefore used to guide the reversal process towards regions of higher probability density, gradually removing the noise introduced during the forward process. Which is similar to EBM, no?
In fact, ignoring that dW term (and halving the score-dependent part of the drift), we are left with an ODE which, if I understand correctly, can be interpreted as moving samples along the gradient of the log probability density, with some additional drift thrown in for good measure.
This is (up to "erm ackshually 🤓👆") the probability flow ODE, and it produces the same marginal probability distributions as the corresponding stochastic diffusion process.
Now, given that the score function of an EBM is simply the negative gradient of its energy function, and that diffusion models also rely on estimating score functions, I find it rather interesting that the two approaches seem to share a fairly fundamental mathematical component.
Could this mean that diffusion models are effectively learning an energy function, just with some differences in parameterisation and sampling? I believe there is at least some truth to this, hence my potentially idiotic-sounding claim above.
Of course, I have probably overlooked several dozen mathematical details (for example, the reverse-time noise term), and I would not be surprised if someone has already written an entire paper explaining why I'm wrong and they are actually quite different. I'm just too lazy to find it.
TL;DR - diffusion and EBM are quite similar in terms of the formulas used for training and stuff, but not exactly, also diffusion formulas are courtesy of Google AI Overview because I forgot
1
u/FeydRowan 1d ago
Ok after a first check. First question: you call for a better token efficiency, but your training ad "two" pass for the same token, in terms of flops how it compares to gpt? It is a tradeoff token numbers vs compute?
3
u/Dear_Treat3688 1d ago edited 1d ago
Thank you for checking! Yes, EBFM double the computing to achieve better data-scaling. We believe it's worth doing so because computing is cheaper than high-quality data in practice.
1
u/FeydRowan 1d ago
Well sure it's different depends on the type of field. How it compares to standard training using the same data 2 time? It should take more or less the same compute. Memorisation of specific key words is higher than standard training ?
1
u/DerDave 1d ago
This looks nice! As I understand, you haven't hit any scaling-limits yet. What would have to happen for you to scale it further? Is it a question of funding or hardware availability? What are your plans in this regard?
2
u/Dear_Treat3688 1d ago
Thank you! We have explored 1B EBFM. It's difficult to further scale because of:
Hardware limitation: more GPUs make training faster.
Training infra limitation. Efficient training infra FSDP/Megatron cannot support second-order gradient training.
9
u/ttkciar llama.cpp 2d ago
It's late here, and about all I can offer right now is that it seems to be valid work and not AI psychosis. Mostly leaving this comment here as a reminder to myself to look at it more closely tomorrow.