OpenEBFM: Energy-based Foundation Models Are Scalable Architecture

AI summary

Overview

• Research area: Natural language processing — specifically, a new foundation-model architecture based on energy-based models (EBMs), evaluated for scaling behavior and general-purpose task performance.
• Technical level: Advanced. The paper assumes familiarity with energy-based modeling, MCMC-style latent optimization, Transformer scaling laws, IsoFLOP analysis, and bits-per-byte (BPB) evaluation.
• Scope: The paper introduces Energy-Based Foundation Models (EBFMs), which combine an energy head and a generative head over a shared Transformer backbone, pre-trains them on three large corpora at depths of 12, 20, and 26 layers (up to 1.03B parameters), and compares them against GPTs and diffusion language models on scaling curves and 26 downstream tasks.

What This Paper Is About

Energy-based models assign a low energy score to good solutions and a high score to bad ones, but prior EBMs have been built and tested for one narrow task or domain at a time, and they produce scores rather than text. The authors ask whether EBMs can be turned into general-purpose foundation models that actually generate language while still using energy evaluation to guide generation. Their answer is EBFM: a single architecture in which an energy head first refines a latent state toward low energy, and a generative head then uses that refined state to predict the next token.

Key Contributions

1. An energy-based foundation architecture. EBFM integrates energy steering and token generation inside one Transformer, with an energy head (a linear layer) and a generative head (an embedding layer with softmax) sharing the same backbone, so energy evaluation directly refines generation rather than replacing it.
2. A parallel training scheme for energy steering. The authors redesign the attention mask and positional indices so that latent states at all sequence positions are steered simultaneously in one forward pass, making large-scale pre-training feasible instead of requiring a separate optimization per position.
3. Scaling-law evidence against GPTs and MDLMs. EBFMs are pre-trained on Climbmix, DCLM, and Fineweb across depths 12/20/26 and from 268M to 1.03B parameters, with data-scaling, depth-scaling, and IsoFLOP analyses, plus supervised fine-tuning curves.
4. Broad downstream evaluation and full open release. Zero-shot results are reported on 19 question-answering datasets and 7 generation datasets, and the model weights (OpenEBFM), training recipe, infrastructure, and logs are released for reproduction.

Main Findings

• Faster data scaling than GPTs. On Climbmix, EBFMs with 12, 20, and 26 layers match GPTs trained on 7B tokens after pre-training on 4.5B, 3.5B, and 5.1B tokens respectively; the paper reports this as requiring 0.67B, 1.49B, and 2.24B fewer training tokens. On DCLM and Fineweb, the reported savings are 1.04B and 1.61B fewer tokens.
• The advantage grows with budget. On DCLM, as training tokens increase from 65M to 10B, EBFM validation error falls faster than GPT's, so the gap widens with more training data.
• Better parameter scaling. Measured on DCLM with 5B and 10B training tokens, EBFM validation BPB drops faster than GPT's as model size grows from 268M to 1.03B parameters.
• GPT-like IsoFLOP behavior. IsoFLOP curves on DCLM are U-shaped, and the compute-optimal point shifts down and to the right as the budget increases, meaning larger EBFMs reach lower validation error.
• Stronger zero-shot QA average. On the 19 QA tasks, EBFM beats GPT at the 20- and 26-layer depths on every corpus; for example, at depth 26 the averages are 44.71 vs 39.79 on DCLM, 43.93 vs 39.47 on FineWeb, and 40.32 vs 36.06 on ClimbMix. One exception is the 12-layer FineWeb setting, where GPT's average (36.29) edges out EBFM's (35.86).
• Best zero-shot generation BPB. Pre-trained on OpenWebText (EBFM and GPT for one epoch on 9.06B tokens; MDLM for two epochs on 18.12B tokens), EBFM achieves an average BPB of 1.35 versus 1.37 for GPT and 1.44 for MDLM — reported as 0.02 and 0.09 lower BPB, respectively. Note that GPT is better on Lambada (1.48 vs 1.53).
• More efficient fine-tuning. Starting from each architecture's optimal pre-training checkpoint on the Nanochat-SFT mixture (SmolTalk, MMLU, GSM8K), EBFMs for depths 20 and 26 reach lower validation error while consuming fewer SFT tokens.

Methodology in Plain English

The model is a Transformer with two outputs. One is an energy head that reads a context plus a latent vector and returns a single number: how compatible they are. The other is a generative head that reads the same context plus the latent vector and predicts the next word.

At inference, the model does two things in sequence. First, energy steering: it starts the latent vector from random Gaussian noise and takes a fixed number of gradient-descent steps on the energy score, adjusting only the latent vector while the model's weights stay frozen. The authors deliberately drop the random noise usually used in MCMC sampling because they found deterministic updates more stable in large-scale pre-training. Second, token generation: the generative head uses the refined, low-energy latent vector to produce the next token. The intuition is that the latent vector is a controllable summary of the context, and pushing it toward low energy yields a representation better aligned with what should come next.

Training the naive way would mean re-running this optimization separately for every token position, which is far too slow. The authors' key engineering trick is to concatenate the input tokens with a full block of latent vectors and design a custom attention mask plus positional indices (1, 2, 3, …, S) so that each position's latent state attends only to its own prefix and itself. This lets all positions be steered in one forward pass, with the loss being standard next-token negative log-likelihood computed from the steered latents. The released configuration uses a context length of 2,048, a vocabulary of 32,768, a global batch size of 512, and 1 MCMC step, with Muon + AdamW as the optimizer.

Evaluation proceeds in four layers: BPB on validation sets of the three pre-training corpora, IsoFLOP and depth-scaling curves, supervised fine-tuning curves on the Nanochat-SFT mixture, and zero-shot accuracy on 19 QA datasets plus zero-shot BPB on 7 generation datasets.

Why This Matters

Impact on research. The paper reframes energy-based models from task-specific "critics" into a scalable generative backbone, and it reports scaling curves that are competitive with GPT-style Transformers while using fewer training tokens. If the token-efficiency results hold at larger scales, it suggests that pre-training efficiency is partly an architectural choice rather than only a matter of data volume. The full release of weights, recipe, and infrastructure makes the claim checkable rather than aspirational.

Potential real-world applications (the paper reports no deployed systems; these follow from the described properties):

• Compute- and data-constrained pre-training, where a given quality target must be reached with a smaller token budget.
• Controllable or constraint-aware text generation, since the energy score is an explicit, optimizable signal over candidate latent states.
• Domain adaptation via supervised fine-tuning, which the paper demonstrates is more data-efficient from an EBFM initialization.
• Research infrastructure for energy-based methods, since the released training code and logs let others build domain-specific EBMs rather than single-task ones.

Industry relevance. The headline claim — matching GPT quality with fewer tokens on a fixed corpus — maps directly onto training-cost economics. The paper is equally explicit about the counterweight: EBFM training currently depends on distributed data parallelism because mainstream model-parallel stacks (FSDP, Megatron-LM, DeepSpeed ZeRO) do not readily support the second-order differentiation EBM training requires. That means model size is capped by single-GPU memory (the authors cite a single H200), which is a practical adoption barrier any industrial-scale user would hit immediately.

Future Directions

1. Training infrastructure for second-order gradients. Extending FSDP, Megatron-LM, or DeepSpeed ZeRO to support model sharding with efficient higher-order gradients is the stated prerequisite for scaling beyond the 1.03B-parameter, depth-26 configuration, which is currently bounded by one H200's memory.
2. Block-wise generation. The paper generates one token at a time; moving to block-wise energy steering would let the model score whole spans and enforce semantic coherence and safety constraints over longer ranges.
3. Efficient inference and training tailored to EBFM. Because the architecture differs from GPTs, existing acceleration tools (GPTQ, Flash Attention, Page Attention) do not transfer cleanly, leaving room for architecture-specific kernels and quantization.
4. Instruction and domain fine-tuning. The authors call for fine-tuning EBFMs on domain-specific instructions and data in areas such as scientific reasoning, code generation, and finance, and for testing whether the pre-training advantages carry into those specialized settings.

Target Audience

Researchers and engineers working on foundation-model architectures, pre-training efficiency, and scaling laws will get the most from this paper, particularly those interested in alternatives to the standard autoregressive next-token paradigm. It is also relevant to practitioners who care about token-efficient pre-training and to anyone already working with energy-based or score-based methods who wants to see them applied at language-model scale. Readers without background in EBMs, MCMC-style latent optimization, or IsoFLOP analysis will need to consult the cited background material first.

Authors’ abstract

Energy-based models have shown promising properties for domain-specific tasks but remain challenging in modeling general-purpose scenarios. To bridge this gap, we design energy-steered models (ESMs) for foundational language modeling. ESMs integrate energy steering and token generation within a single architecture to leverage the energy-guided search to refine token generation. Specifically, ESMs first optimize a latent state over the energy landscape, then leverages the refined latent state for token generation. By performing such energy-guided pre-training on large-scale corpora, ESMs learn generally effective energy landscape where favorable tokens correspond to lower energy values. To explore the scaling behavior of ESMs, we pre-train ESMs on three large-scale corpora, then evaluate across 26 downstream tasks spanning a broad range of domains. Our results demonstrate that ESMs have promising scaling behaviors with increasing model size and training data, and outperform GPTs on most downstream tasks. These findings highlight ESMs as a promising architecture for large-scale pre-training and general purpose problem solving. We fully release the model weights and training infrastructure to enable the community to completely reproduce and extend our results. GitHub: https://github.com/datamllab/openebfm HuggingFace: https://huggingface.co/collections/guan-wang/openebfm