OpenEBFM: Energy-based Foundation Models Are Scalable Architecture

Yuan Pu; Yudong Lu; Ting Wang; Zhicong Hong; Qinuo Li; Xuansheng Wu; Na Zou; Xia Hu

cs.CLOct 7, 2026

OpenEBFM: Energy-based Foundation Models Are Scalable Architecture

Authors: Yuan Pu, Yudong Lu, Ting Wang, Zhicong Hong, Qinuo Li, Xuansheng Wu, Na Zou, Xia Hu,

Organizations: Shanghai Artificial Intelligence Laboratory

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

Figures & tables

Panel of Figure 1 — Figure 1: Scaling behavior of EBFMs on the DCLM dataset.

Panel of Figure 2 — Figure 2: Generation process of EBFMs. (a) Energy steering: The energy head estimates the energy of the latent state and backpropagates the energy gradient with respect to $\boldsymbol{z}$ to obtain a low-energy latent state $\boldsymbol{z}^{*}$ for token generation. (b) Token generation: The generative head uses the optimized latent state $\boldsymbol{z}^{*}$ to generate the next token.

Figure 3: EBFMs parallelize energy steering for training tokens at all positions through re-designed attention masks and positional indexes. Left Panel: the re-designed attention mask where 0 indicates blocking the cross-token-attention. Right Panel: the forward path with re-designed attention mask and positional index.

Hyper-parameter	EBFM	Hyper-parameter	EBFM	Hyper-parameter	EBFM
Tokenizer	Nanochat	MCMC Steps	1	Weight Decay	0.2
Vocabulary Size	32,768	MCMC Step Size	300	Warmup Ratio	0.0057
Context Length	2,048	$\alpha$ Learning Rate	1.8	Warmdown Ratio	0.65
Global Batch Size	512	Muon Learning Rate	0.02	Muon Momentum	0.95
Optimizer	Muon + AdamW	Embedding Learning Rate	0.3	Muon $\beta_{2}$	0.95
Gradient Clipping	1	Final LR Fraction	0.05	AdamW $\beta_{1}$ , $\beta_{2}$	0.8, 0.95

Table 1: Hyper-parameter setting for EBFM pre-training.

Panel of Figure 4 — Figure 4: Validation error during pre-training on the (a) Climbmix, (b) DCLM, and (c) Fineweb datasets; (d) Scaling law of EBFMs versus different training budgets

Panel of Figure 5 — Figure 5: Scaling behaviors of EBFM with respect to model depth and training FLOPs.

Panel of Figure 6 — Figure 6: Validation error during supervised fine-tuning. (a) $d20$ models (b) $d26$ models

Pretrain data	Model	Depth	AGIEval	ARC		BIG-bench					BoolQ	CSQA	COPA	CoQA	HellaSwag		LAMBADA	OBQA	PIQA	Winograd		Avg.
Pretrain data	Model	Depth	LSAT-AR	C	E	CSAlg	Dyck	LangID	Ops	WikiQA					10-shot	0-shot				WSC	WinoG
	GPT	$d12$	26.96	23.20	42.80	40.80	11.20	25.80	11.90	28.20	55.60	21.80	56.00	15.00	31.00	30.20	31.20	26.40	57.00	56.41	50.20	33.77
		$d20$	23.91	26.60	47.60	41.00	13.20	26.20	15.71	34.20	57.60	28.00	62.00	19.80	34.00	34.00	35.80	25.40	59.60	59.34	46.00	36.31
		$d26$	26.52	27.80	46.60	37.40	14.20	26.00	17.14	37.00	50.40	25.00	59.00	20.80	33.40	32.60	35.00	27.20	60.00	57.51	51.60	36.06
		$d12$	21.74	23.29	44.36	42.12	13.70	25.63	18.57	30.75	56.24	37.84	58.00	16.47	30.47	30.75	28.18	27.20	58.16	54.58	49.33	35.13
		$d20$	24.35	26.02	50.80	41.59	17.50	24.83	21.90	42.55	56.61	25.72	61.00	26.92	36.70	36.60	38.93	27.00	61.21	64.47	50.12	38.67

Table 2: Zero-shot accuracy (%) on 19 question-answering datasets after pre-training.

Panel of Figure 7 — Figure 7: (a) Zero-shot accuracy versus model size (b) Zero-shot accuracy versus training tokens on DCLM dataset

Model	OpenWebText-val	PTB	Wikitext	LM1B	Lambada	AG News	Pubmed	Arxiv	Average
EBFM	0.93	1.46	1.10	1.24	1.53	1.34	1.34	1.44	1.35
GPT	0.94	1.51	1.11	1.26	1.48	1.35	1.39	1.51	1.37
MDLM	1.17	1.53	1.23	1.38	1.48	1.52	1.41	1.50	1.44

Table 3: Validation BPB of EBFMs, GPTs, and MDLM pre-trained on the OpenWebText datasests, where EBFMs and GPTs are trained for only ONE epoch with

9.06

B tokens, and MDLM is trained for TWO epoches with

18.12

B tokens. Average values are estimated for zero-shot settings.

Language Modeling

Language modeling develops systems that learn patterns in text and other sequences to understand, predict, and generate language.

See all 522 papers

Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models
Mar 12, 2026Samy Jelassi, Mujin Kwun, Rosie Zhao +5Fine-Tuning Language Models LLM Fine-Tuning
Harvard University, Kempner Institute · MBZUAI · Microsoft Research New England
OpenLanguageModel: Readable and Composable Small-Language-Model Pretraining for Education and Research
Jul 18, 2026Tavish Mankash, Vardhaman Kalloli, Keshava Prasad +1Language Model Pretraining Small Language Models
FAIRC · Plaksha University
Scaling Transformer Energy-based Models For High Quality Image Synthesis
Oct 22, 2026Jungsoo Lee, Alexi Gladstone, Yilun DuImage Generation Energy-Based Models
MIT · UIUC · Harvard University

Mar 12, 2026J/K move · Enter open · S save

Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models

Samy Jelassi, Mujin Kwun, Rosie Zhao, Yuanzhi Li+4

Harvard University, Kempner Institute · MBZUAI · Microsoft Research New England

Fine-Tuning Language Models LLM Fine-Tuning RL Fine-Tuning

Open paper AskPDF arXiv

Cross-entropy (CE) training provides dense and scalable supervision for language models, but it optimizes next-token prediction under teacher forcing rather than sequence-level behavior under model rollouts. We introduce a feature-matching objective for language-model fine-tuning that targets sequence-level statistics of the completion distribution, providing dense semantic feedback without requiring a task-specific verifier or preference model. To optimize this objective efficiently, we propose energy-based fine-tuning (EBFT), which uses strided block-parallel sampling to generate multiple rollouts from nested prefixes concurrently, batches feature extraction over these rollouts, and uses the resulting embeddings to perform an on-policy policy-gradient update. We present a theoretical perspective connecting EBFT to KL-regularized feature-matching and energy-based modeling. Empirically, across Q&A coding, unstructured coding, and translation, EBFT matches RLVR and outperforms SFT on downstream accuracy while achieving a lower validation cross-entropy than both methods.

Abstract

Language Modeling

Explore similar work

Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models

OpenLanguageModel: Readable and Composable Small-Language-Model Pretraining for Education and Research

Scaling Transformer Energy-based Models For High Quality Image Synthesis