Energy-Based Transformers are Scalable Learners and Thinkers
Authors: Alexi Gladstone
Ganesh Nanduru
Md Mofijul Islam
Peixuan Han
Hyeonjeong Ha
Aman Chadha
Yilun Du
Heng Ji
Jundong Li
Tariq Iqbal
Venue: ICLR 2026 Oral (2026)
Keywords: Energy-Based Models
System 2 Thinking
Reasoning
Verification
Scaling
Transformers
Generative Modeling
Primary area: generative models
TL;DR: We introduce Energy-Based Transformers, a scalable new approach for learning how to think from unsupervised learning, generalizing current System 2 Thinking/reasoning approaches.

Abstract: Inference-time computation, analogous to human System 2 Thinking, has recently become popular for improving model performance. However, most existing approaches suffer from several limitations: they are modality-specific (e.g., working only in text), problem-specific (e.g., verifiable domains like math and coding), or require additional supervision/training on top of unsupervised pretraining (e.g., verifiers or verifiable rewards). In this paper, we ask the question “Is it possible to generalize these System 2 Thinking approaches, and develop models that learn to think solely from unsupervised learning?” We find the answer is yes, by learning to explicitly verify the compatibility between inputs and candidate-predictions, and then re-framing prediction problems as optimization with respect to this verifier. Specifically, we train Energy-Based Transformers (EBTs)---a new class of Energy-Based Models (EBMs)---to assign an energy value to every input and candidate-prediction, enabling predictions through energy minimization until convergence. To support this approach, we introduce several key techniques for stable and parallelizable training, which enable the emergence of strong System 2 Thinking capabilities and scalable EBMs. Across discrete and continuous modalities, we find EBTs outperform the Transformer++ approach, scaling up to 35% faster during pretraining, and improving inference-time performance by up to 29%. EBTs also surpass Diffusion Transformers on image denoising while requiring 99% fewer forward passes. Moreover, System 2 Thinking with EBTs yields larger performance gains on data that is farther out-of-distribution, and EBTs achieve better results than existing models on most downstream tasks despite achieving the same or worse pretraining performance, enabling EBTs to generalize better than existing approaches. Consequently, EBTs are a flexible and exciting new approach for scaling both the learning and thinking capabilities of models.

Published as a conference paper at ICLR 2026
ENERGY-BASED TRANSFORMERS ARE SCALABLE LEARNERS AND THINKERS
Alexi Gladstone Ganesh Nanduru', Md Mofijul Islam',, 13, Peixuan Han', Hyeonjeong Ha², Aman Chadha²4, Yilun Du', Heng J Jundong Li', Tariq Iqbal
UVA
UIUC Amazon GenAl Sttnnfrrd University 'Harvard University
gittub.com/alexxggd//EB1
ABSTRACT
Inference-time computation, analogous to human System 2 Thinking, has recently become popular for improving model performance. However, most existing ap- proaches suffer from several limitations: they are modality-specific (e.g., working only in text). wroblem-speciffe (c.g., verifiable domains like math and coding), or require additional supervision/rraining
top of unsupervised pretraining (e.g., verifiers or verifiable rewards). In this paper, we ask the question "Is it possi- ble to generalize these System Thinking approaches, and develop models that learn to think solely from unsupervised learning?" We find the answer is yes, by learning to explicitly verify the compatibility between inputs and candidate- predictions, and then re-framing prediction problems as optimization with respect this verifier. Specifically, we train Energy-Bass Based Transformers (EBTs) new class of Energy-Based Models (EBMs) to assign an energy value to every input and candidate- prediccioon,. enabling predictions through energy minimization until convergence. To support this approach, we introduce several key techniques for stable and parallelizable training, which enable the emergence of strong System 2 Thinking capabilities and scalable EBMs. Across discrete and continuous modal- ities, we find EBTs outperform the Transformer++ approach, scaling up to 35% faster during pretraining, and improving inference-time performance by up to 29%. EBTS also surpass Diffusion Transformers on image denoising while requiring 99% fewer forward passes. Moreover, System Thinking with EBTs yields larger performance gains on data that is farther out-of out-of-distribution, and EBTs achieve better results than existing models on most downstream tasks despite achieving the same or worse pretraining performance, enabling EBTs to generalize better than existing approaches. Consequently, EBTs are a flexible and promising new approach for scaling both the learning and thinking capabilities of models.
INTRODUCTION
In psychology, human thinking is often classified into two different types: System (thinking fast) and System 2 (thinking slow) Evans (2011); Frankish (2010); Kahneman (201 Kahneman et al. (2002). System thinking is characterized by quick, intuitive and automatic responses, relying on previous experience to solve simple or familiar problems. Alternatively, System 2 Thinking is slow and deliberate, requiring effort to solve complex problems that go beyond automatic pattern recognition, such in mathematics Current models perform well on tasks suitable for System thinking Li et al. (2025c), but continue to struggle with tasks that demand System 2 capabilities Mirzadeh al. (2024); Yan et al. (2025).
-uittf-distiibutionn situations Goel et al. (2000); Neys (2006).
As a result, System Thinking has become a growing research focus, driving the development of foundation models such 01 Jacchh et al. (2024), RI Guo al. (2025), and Claude Anthropic (2025). These "reasoning models" excel on math and coding benchmarks by increasing the time models
Correspondence Alexi Gladstone:
alexigladstone @ gmail.com. Work does not relate to position
Amazon.

---

Published as a conference paper at ICLR 2026
1
X
RNN
X1
X1
+ RNN
X2
X
Xt
L
X
R1+1
Feed
RNN
Diffusion Transformer
Energy
Forward
Based
Transformer
Xt
RNN
Transformer
8c+1
21+1
Noisc(2c+1)
Energy(&e+)
(a) AR Transformer
(b) RNN
(c) DīT
(d) EBT
Existing Autoregressive Approaches
Figure 1: Autoregressive Architecture Comparison. (a) Autoregressive (AR) Transformer is the most common,
with (b) RNNs becoming more popular recently Gu & Dao (2023); Peng al. (2023). Diffusion Transformers (DiTs) Li et al. (2025h); Peebles Xie (2023) are similar to EBT, being able dynamically allocate computation during inference. However, diffusion models are not trained as explicit verifiers, unlike EBTs.
spend thinking. However, publicly available information from the open-source R1 model Guo et al. (2025), suggests that the Reinforcement Learning (RL) approach for training these models only works domains where rule-based rewards can easily verify answers, such as math and coding. This limits applicability, and often harms performance tasks such as writing OpenAl (2024); Su et al. (2025). Moreover, recent evidence indicates this approach may not foster new reasoning patterns, restricting performance tasks requiring exploration Yue et al. (2025). For discussion on Related Works, see Section E.
As one of the primary goals of AI is to figure out how can create systems that learn to think on their own on any problem type. these approaches ultimately bring about the following core research question: "Can rely entirely unsupervised learning to develop System Thinking?" Such capability could enable generalization of current System Thinking approaches to many problems, many modalities, and avoid the reliance human, reward, or model supervision.
We argue and our empirical results suggest that the answer to this question is yes, but that there are several limitations in existing models that prevent this general Thinking from emerging. Particularly, when comparing the qualities of human System 2 Thinking with current modeling approaches (Figure Table 1), we observe several key differences, outlined below as two key Facets of System Thinking:
Facet 1: Dynamic Allocation of Com- putation. Humans naturally allocate varying amounts of effort to different tasks depending difficulty, which is widely supported by psychology and neu- roscience Ditterich (2006); Kahneman (2011); Rougier et al. (2005).2 For ex- ample, a decision regarding whether to change careers generally takes more time than deciding what to eat.
Table 1: Architectures and Cognitive Facets. For each pre- diction,
(FF) Transformers and RNNs gener- ally' have finite amount of computation. DiTs (Diffusion Transformers) can increase inference computation by denois- ing longer, but lack explicit prediction verification. In contrast, EBMs support dynamic computation through flexible iteration, and give energy scalar for prediction verification.
Dynamic Compute Allocation (Facet 1)
Prediction Verification (Facet 2)
Arch.
Facet 2: Verification of Predictions. In
FF Trans.
RNNs
addition to allocating computation, hu- man thinking also benefits from the abil- ity to verify predictions Alkouri (2016);
DiTs
EBTs
Loesche al. (2018). which guide decisions about when to stop thinking to select the most accurate predictions. This also supports more dynamic inference time behavior, such as early stopping when a prediction is known to be correct, or allocating more compute when problem is difficult. For more information on additional Facets, please refer to Section F.
Recent works attempt enable dynamic computation per prediction Geiping al. (2025); Hao al. (2024), these approaches are generally not modality agnostic and have not heen widely adopted.
We refer to this dynamic computation at the granularity of each prediction, meaning current L.MMs built with AR transformer/RRNNG cannot dynamically allocate compute per token. See Section G for more info.
2

---

Published as a conference paper at ICLR 2026
Text Pretraining Objective: Predict Next Token
Video Pretraining Objective: Predict Next Frame
Context
Content
The caught the
Energy Based
Transformer
Energy
Energy Bosed Transformne
Energy
Linear Projector
Predictions Across Stsps
Prodictions Actoss Stops Step
Step 0
Step
Step 2
Step
Step 4
Step
Step 2
Step
Step
Thinking Process:
Minimize Energy Predicted
Thinking Process: Minimize Energy Predicted Frume using Gradient Descent Iteratively
Distrihution using Gradient Descent Iteratively
Figure 2: EBT for Autoregressive Modeling. Each blue box corresponds to different prediction at each step of the thinking process, where the initial prediction starts as random. At each step, new prediction is fed into the model, which gives energy scalar for the prediction's 's current compatibility (unnormalized likelihood)) with the context (Facet 2). Then, the gradient of this energy with respect to the prediction is calculated and used to update prediction. This gradient descent update is done iteratively to refine the prediction until convergence the predicted energy, which allows for dynamic use of computation (Facet 1).
To achieve the facets described, we propose viewing thinking as optimization with respect to a learned verifier, which evaluates the compatibility (unnormalized probability) between an input and candidate prediction (Figure 2). Specifically, we train Energy-Based -ased Models (EBMs) to learn an energy (unnormalized probability) landscape over all possible energy indicates higher compatibility (Facet 2). Thinking then corresponds to starting from an initial random prediction and refining it through optimization along the energy landscape until convergence (visualized Figure 3). This naturally enables dynamic compute allocation (Facet I) in the form of more challenging problems utilizing additional optimization steps.
pairs, where lower
While this thinking perspective promising, Energy- Models (EBMs) have struggled with scalability Du & Mordatch (2019), with no known foundation EBMs. This stems from issues with training instability and long training times Arbel et al. (2020); Du & Mordatch (2019). To address these challenges,
introduce Bnergy-Based Transformers (EBTs), Transformers specifically for EBMs. We further propose practical training improvements, theoretical insights into EBM training scalability, and novel energy landscape regularization techniques that improve System 2 Thinking.
To assess learning and thinking scalability, we compare EBTs to the Transformer++ (autoregressive) and DiT (bidirectional) across discrete and continuous modalities. EBTs show up to 35% higher scaling rates than the Transformer++ across data, batch size, parameters, FLOPs, and depth. At inference, EBTs outperform existing models on System 2 Thinking for example, by improving language model performance by 29% more than the Transformer++, and by outperforming DiTs in image denoising with 99% fewer forward passes. We observe two key effects: (1) EBTs often outperform baselines at inference even with worse pretraining performance, demonstrating the importance of System 2 Thinkingg and (2) System Thinking yields greater gains on more out-of- distribution data, paralleling human thinking. We believe the EBT implementations, along with novel techniques for EBMs to maximize the learning and thinking scalability, will advance the EBM approach by addressing key challenges in stable, parallelizable, and efficient training.
2
ENERGY-BASED TRANSFORMERS (EBT) INTUITION
2.1 LEARNING TO VERIFY
Verifying solutions often substantially more tractable than generating them, distinction well- known in complexity theory Cook (2023); Goldwasser et al. (2019); Gödel (1956). For example, in solving a maze, verifying the correctness of given path is significantly easier than discovering such path. This asymmetry has been recognized and utilized for several decades, notably in the field of cryptography Goldwasser et al. (2019); Lavin et al. (2024); Rivest et al. (1978). EBMs are built on this principle that verification is easier than generation: rather than learning to generate directly, as in most existing approaches, EBMs learn to generate by optimizing predictions with respect to the learned energy function (shown in Figure 3).
3

---

Published as a conference paper at ICLR 2026
Context dog caught the
High Energy (Low Compatibllity)
Predictions
Siep
Step
Energy Landscape
Step
Energy (High Compatibilit)
Converged
Tall
Continue Thinking Process until Energy Convergence
Figure 3: Thinking Process Visualization. An example energy landscape and its optimization through gradient descent, interpreted as thinking. Here, the model predicts a distribution over text tokens, progressively shifting from an initial random distribution to the target distribution. At cach step, the EBM assigns an energy scalar indicating how compatible the current prediction is with the context, visualized the landscape's height (Facet 2). This scalar's convergence allows the model determine whether the prediction is adequate if further thinking is necessary (Facet 1). We include more detailed toy example Section C.3. Adapted from Li et al. (2018).
Recent works have attempted leverage verifiers Ma et al. (2025); Team (2023); Yao al. (2023), but these approaches decouple the verifier and generator, resulting in adversarial dynamics Ma et al. (2025) and challenges in scalability Yao al. (2023). For example, researchers combining tree search and LLMs required thousands (2023). In contrast, EBMs combine the verifier and generator into a single model, where the generator is defined implicitly by the gradient of the verifier Du & Mordatch (2019). We show that this coupling addresses scalability and adversarial issues (Figures 7b and B.4a).
millions of samples to achieve optimal performance Team
An additional advantage of verifiers is generalization. Because verification is usually casier than generation Swamy et al. (2025), prediction verification on Out- Of-Distribution (OOD) data is often easier than explicit prediction generation for OOD data Du al. (2022). This characteristic often results in better generalization of verifiers than explicit generators Du et al. (2022). This may explain why EBMs often generalize better than existing models Du et (2022; 2024), which we further support in our experiments (Figure 7a and Table 4).
2.2 LEARNING TO UNDERSTAND
This verifier-centric perspective also relates to a deeper limitation generative models, referred to as "The Generative AI Paradox West al. (2023)." Although current generative models achieve strong generative capabilities, they frequently lack basic discrimination skills, such as the ability to assess the plausibility coherence of their predictions Stojnić et al. (2023); West al. (2023), impeding their ability to engage in reasoning, planning, and decision Docision-mmaing Kambhampati et al. (2024); Yan et al. (2025). In contrast, EBMs offer potential solution to this challenge: EBMs generate by learning verifier (which is similar to a discriminator), they develop strong discrimination skills Wang al. (2023). Experimental results further support this observation (Table 4).
ENERGY-BASED TRANSFORMERS (EBT) APPROACH
3.1 ENERGY-BASED MODELS (EBM) BACKGROUND
Energy- Based Models (EBMs) assign scalar energy value to each configuration of input variables. enabling them to model the compatibility and interactions between variables, such as between context and cendidate-prediction. For probabilistic EBMs, this defines a probability distribution
using Boltzmann distribution (= = -Pfrr) where Z(0) f - (x) is the intractable
2(()
partition function involving integral over all possible values of z. To avoid the intractability of the partition function, it is common to work with unnormalized EBMs, which dispense of the partition function in favor of representing relative unnormalized probabilities. This formulation shifts the
4

---

Published as a conference paper at ICLR 2026
Algorithm 1: Training
Algorithm 2: Inference with Verification
Inputs: Context 2. Target y. EBM Es (s 8)
Inputs: Context æ, EBM (x. 5))
Hparams: Steps M. Step Size a, Loss J(.)
Hparams: Steps M, Step Size a, Samples N for j I...., NN do
Sample to ~ (O, I);
for i 0, ................................ M
do
Sample jo.5 (0, I);
Di+1 - U1 - V Eo(a, (s): f- JÚM. y):
for i 0, 0...., M do
0i+1, f- Di.5 a vn. (a, return 0" = =rgmmn, Eo(2,, (od.));
return L. update Es:
focus from addressing the partition function, to simply assigning low energy to the true data manifold and high energy elsewhere Dawid & LeCun (2024); Du & Mordatch (2019), offering benefits such as scalability to spaces where the true data manifold is thin and therefore a probabilistic EBMs would have an infinite score Dawid & LeCun (2024). In supervised or predictive self-supervised learning (e.g., classification, autoregressive modeling, masked modeling), unnormalized EBMs can be formulated as: p0( 9) xx e Fo(n.))), where the goal of the EBM is to learn to predict ý given x.""
3.2 SCALABLE EBM LEARNING
While EBMs offer a flexible modeling framework, training them scalably remains an open research problem. Two primary training approaches exist-contrastive and regularized methods LeCun (2022). Contrastive methods increase the energy of negative samples while decreasing the energy of positive samples. Due to the curse of dimensionality Dawid & LeCun (2024), where the volume of spaces grows exponentially with their dimension, contrastive methods struggle to scale because they must increase the energy of an exponentially higher number of negative samples.
An alternative is to frame EBM learning as an optimization problem Du al. (2022); Wang et al. (2023), which avoids the curse of dimensionality by implicitly regularizing the energy landscape, enabling scalable learning. In this approach, EBMs are trained to optimize an initial prediction to the ground truth solution through gradient descent, as shown in Figure 3. This pushes the energy landscape to have a local minima surrounding the ground truth solution, thereby regularizing the energy landscape to only have low energy the true data manifold. Intuitively, this optimization- based training approach is similar to GANS Goodfellow et al. (2014). During the forward pass, FB GAN discriminator by giving an energy "verification"; on the backward pass they can be seen a GAN generator by optimizing predictions through energy minimization to try and fool the discriminator.
be
Training EBMs to perform optimization can be formalized as follows. We begin with an EBM Eg. an initial prediction do. an input (context) for the model x, and seek to predict y. We aim to find the minimum energy (most compatible) given an x, which we search for using gradient descent:
ii+1 = D. aV (z,
(1)
where a is the step size (formalized in Algorithm 1). Then, the loss can be computed using any standard objective function. Importantly, this loss is backpropagated through the entire optimiza- tion process, requiring second-order derivatives (i.e., gradients of gradients). These are computed efficiently via Hessian-vest vector products, which scale linearly with model size Dagréou et al. (2024), similar to standard gradient descent in feed-forward models. More details and pseudocode can be found in Section I.2.
3.3 SCALABLE EBM THINKING
While this training approach is scalable, achieving smooth energy landscapes with single local minimum remains challenging on real-world problems. Because y is high-dimensional, the energy landscape spans high-dimensional space, and must remain well-shaped throughout. To address this, we found three key energy landscape regularization techniques to be helpful for learning smoother, more wehl-behaved energy landscapes, enabling strong thinking capabilities emerge during training.
First, we found a replay buffer helps simulate longer optimization trajectories, enabling energy landscapes to be well defined near their minimum. Second, a variant of Langevin Dynamics Du
In self-supervised learning, x some unmasked portion of the original and ý is the masked portion.
5

---

Published as a conference paper at ICLR 2026
EBT Scales 35.98% Foster Than Trrnsformer-
EBT Scales 28.46% Faster Than Transfommrrs
EBT Scales 5.29% Paster Than Transfommet
Transforter+
Terafornee
Validation Remploslt Gogg
Trenficmn-
Papiedry
Vilkacoo Paptesyy
EBT
&
50
20
Namber Toleme
Tranfienn Biudo allg)
(a) Scaling for data.
(b) Scaling for batch size.
(c) Scaling for depth.
Figure 4: Language Learning Scalability-Daals Batch Size, and Depth. A comparison between the scaling of the Transformer++ recipe Touvron et al. (2023) and EBTs across data, batch size, and depth during pretraining On all axes, EBTs out-scale the Transformer++ recipe significantly, indicating improved data efficiency. The improved depth scaling offers promise for reasoning, where depth is crucial Ye et (2024). These results suggest that EBTs offer promise at large data scale.
Mordatch (2019) was found be helpful for encouraging exploration of the energy landscape:
Di+1 ÿ a VV 6. Flg (1. ús) + *%, Y ~N(,
where a is the magnitude of the noise 7). Without this noise term, exploration is often limited to paths leading directly to the energy minimum, leaving other regions poorly defined. Third, varying the paths taken towards predicting solutions, by randomizing the gradient descent step size and number of optimization steps, significantly improved generalization. Together, these techniques improved the System Thinking capabilities of models, as confirmed by ablation experiments in Table 2.
With these techniques established, we explored two main Thinking approaches. First, corresponding to dynamic computation allocation (Facet 1), we conduct experiments that involve changing the number of steps taken for optimization of a single prediction. Second, corresponding to the ability to verify predictions (Facet 2), we generate N predictions from an EBM and choose the minimum energy prediction (BoN or Self-Verification, which is formalized in Algorithm 2). This is conceptually similar to Best of (BoN) sampling using language models Stiennon et al. (2020). However, EBMs generalize this approach to both discrete and continuous modalities and perform it on every single prediction, not just to entire sequences. We demonstrate performance improvements gained from both of these techniques in several Thinking experiments (Figures 7, 6, and B.6), which further confirm the importance of the described cognitive facets. This thinking process is formalized in Algorithm and we more formally define and justify usage of the term thinking in Section C.I.
3.4 ENERGY-BASED TRANSFORMERS (EBTs) ARCHITECTURE
Transformers excel across domains due to their parallelizability, stability, and scalability Borsos et al. (2023): Oquab et al. (2023); Radford et al. (2019); Vaswani et al. (2017). In contrast, Energy-Baser Models (EBMs) struggle with these aspects Du & Mordatch (2019); Du et al. (2020); Li et al. (2023), making Transformers natural fit for scaling EBMs. Consequently, we introduce Energy-Based Transformers (EBTs), Transformer implementations designed for EBMs. We developed two variants: a GPT-style Radford et al. (2018) causal decoder-only EBT, for autoregressive modeling, and a bidirectional EBT with full sequence attention Devlin et al. (2019); more details are in Section C.6. While the bidirectional EBT implementation is straightforward, the autoregressive EBT requires care to prevent information leakage (see Section C.6).
EXPERIMENTATION AND RESULTS
We experiment with EBTs across both Autoregressive (AR) Radford al. (2019) as well as bidirec- tional models Devlin et al. (2019) in discrete and continuous spaces.* 4 In discrete spaces, we focus on the language modeling objective. continuous spaces, we focus on vision tasks of next frame prediction (Section I) and image denoising. All models are pretrained from scratch, EBT's archi- tecture is incompatible with existing foundation models, and therefore cannot be fine-tuned. We focus
'Here, autoregressive and bidirectional refer to the procedure for generation. It's worth noting that autore- gressive models compatible with bidirectional attention, in Li et al. (2025b).

---

Published as a conference paper at ICLR 2026
Table 2: System Thinking Ablations. All energy landscape regularization techniques described in Section 3.3 and their impact on System 2 Thinking performance, measured by percent perplexity improvement. Thinking Longer denotes more optimization steps and Self- Verification denotes optimizing many predictions and choosing the best. Removing regularization, such as Langevin Dynamies, results in less energy landscape exploration, which improves single path performance (thinking longer) at the expense of self-verificatin performance.
Model
Thinking Longer f
Thinking Longer and Self-Verification 1
No Random Step Size
-1.47
0.19
No Random Num. Steps
0,00
9.65
No Langevin Dynamics
17.2
17.0
No Replay Buffer
14.8
17.8
Full System Configuration
7.19
18.7
FRT Scales s, 8.97% Fadter The Trandfrment
FRT Seales 8.97% Faster Than Transfomer
EBT Scales 0.02% Faster Than Transformert
Tranforrares
Thoodhmant
Valiiatio Peppesit Coop scak))
Tronetiome+4
EST
Waldamu sae (in solee
EBT
4.0
Nes Enbhkking Panzntas (K wole
Traning FLOP deg wald
Imbeddinng Dimondte (lng scale)
(a) Scaling for number of Parameters. (b) Scaling for number of FLOPs. (c) Scaling for the embed, dimension.
Figure 5: Language Learning Sealabilit-PPrameeeees FLOPs, and Width. Pretraining scaling comparisons between the Transformer++ recipe Touvron et al. (2023) and EBTs across model size (parameters), compute (FLOPs), and width (emhedding dimension). EBTs have an 8.97% higher scaling rate than the Transformer++ in FL.OP and parameter scaling and b), suggesting that EBTs offer promise as pretraining approach.
on two primary types of results. First, examine learning scalability, investigating how quickly models can fit the pretraining data, which is standard in pretraining Gu & Dao (2023); Hoffmann et al. (2022); Kaplan et al. (2020); Touvron et al. (2023). Second, we study thinking scalability, or how model performance changes as we scale the System 2 Thinking of models (Definition C. 1), measured with the Number of Function Evaluations (NFEs) Chen et al. (2018); Ma et al. (2025) (forward passes).
4.1 AUTOREGRESSIVE LANGUAGE MODELING EXPERIMENTS
In this section, we detail and discuss the results for all NLP experiments using Autoregressive (AR) Language Models trained to predict the next discrete token in a text sequence Radford et al. (2019). All language models are pretrained on the RedPajama text corpus Computer (2023); Weber et al.
(2024) 100B sample from HuggingFace using the GPT-NeoX tokenizer Black et al. (2022b) (as in Gu & Dao (2023).. Following existing pretraining work, we compare AR EBT with the standard Transformer++ recipe Gu & Dao (2023); Sun et al. (2024); Touvron et al. (2023).
For downstream evaluation, we used four key datasets in addition to the pretraining dataset, spanning reasoning, question answering, and syntax understanding. Ordered roughly by increasing perplexity difficulty, these include GSM8K Cobbe et al. (2021). 1), SQuAD Rajpurkar et al. (2016), BigBench Elementary Math QA Srivastava et al. (2022). and BigBench Dyck Languages Srivastava al. (2022). We focus on reasoning benchmarks due to their close alignment with System Thinking.
We conduct scaling experiments for six different axes- includinn data, batch size, depth, parameters, FLOPs, 3 and embedding dimension. The results for the data, batch size, and depth scaling are shown in Figure 4; and the results for parameters, FLOPs, and embedding dimension visualized in Figure 5. Across all axes, EBTs consistently have a higher scaling rate than the Transformer++ recipe, suggesting that EBTs offer promise at large scale
'The FLOPP calculation is nuanced and depends on specific hyperparameters, please refer to Section D.5. 'We cannot run larger-scale experiments due to limited computational resources, so larger-scale performance remains speculative.
7

---

Published as a conference paper at ICLR 2026
Table 3: Language Model Task Generalization Comparison. We conduct experiments aimed demonstrating the generalization EBTs. Despite having slightly higher pretraining perplexity, EBTs often achieve lower perplexity on downstream tasks than the Transformer++ indicating better generalization. All models are trained with the same amount of data and parameters, hut because EBTs at the current scale are less FLOP efficient (sec Figure 5b), they used more FL.PPs for this experiment. BB stands for BigBench.
Model
Pretrain
GSM8K +
SQUAD
BB Math QA &
BB Dyck $
Transformer++
31.36
49.6
52.3
79.8
131.5
EBT
33.43
43.3
53.1
72.6
125.3
Building on the learning results, we investigate EBTs for thinking at inference time. We found that the thinking capabilities of EBT emerge with sufficiently large data scale, and therefore, due to limited resources, we focus on conducting thinking experiments with smaller models trained on substantial amounts of data. In Table 2 conduct ablation studies to confirm the benefits of our energy landscape regularization techniques for System 2 Thinking on from the BigBench Dyck Languages benchmark Srivastava et al. (2022). We find that using all techniques yields the best System 2 Thinking performance when combining extended thinking and self-verifie reff-verificatin.. Additionally, the results show that randomizing the step size is crifical-removing it nearly eliminates thinking gains. In contrast, disabling Langevin Dynamics degrades combined performance but improves results without verification, offering a performance-comppe
Data
tradeoff.
Having established the importance of these landscape regularization techniques, in Figure 7, we ana- lyze the scalability of thinking with EBTs, where the results yield two main insights. First, as shown in Figure 7a, EBT's are able to improve performance by as much as 29% by increasing the amount of forward passes (thinking time), whereas the Transformer++ cannot improve performance. This aligns with our claims that because traditional feed-forward Transformers cannot dynamically allocate additional computation for each prediction being made, they are unable to improve performance for each token by thinking for longer.
Second, as demonstrated in Figure 7b, the thinking capabilities of EBTs scale, showing that EBTs are trained for longer, their ability to achieve improve- ments from verification improves, increasing up 12% - 14% from 4% - 8%. This suggests that EBTs trained at the same scale as modern foundation mod- els, such as the 15T tokens Llama3 Grattafiori et al. (2024) was trained on (+= 1000x the current scale), could have more substantial self-verification results.
Thinking Helps More on OOD Data
1PP imporegen By Tainkin
EBT With Max Thinking
20
18
16
14-
12
As System 2 Thinking humans is associated with generalization to novel scenarios, we conduct ex- periments directly aimed at measuring the effects of System 2 Thinking ure 6, we visualize the performance of EBTs on the datasets described, which have varying levels of Out- of of-Distribution (OOD) shift (measured as the ratio of downstream task perplexity to pretraining perplexity). We observe a strong linear trend: as the data becomes more OOD, thinking leads greater performance improvements. Therefore, these findings suggest that the benefits of EBTs' thinking are not uniform across all data, but scale positively with the magnitude of distributional shifts, highlighting thinking as critical mechanism for robust generalization. These findings align with observations in psychology. where humans rely on deliberate System 2 Thinking to tackle challenging OOD tasks.
1.0
1.5
2.0
2.5
3.0
3.5
4.0
OOD Magnitude Shift
Figure 6: OOD Thinking Performance. As data becomes
generalization. In Fig-
OOD, thinking with EBTs leads greater performance how thinking is critical for generalization to OOD data. Performance is measured on 5 datasets vary- ing in Out-of- Diitriibibbttii (OOD) magnitude shift, which measured the ratio of downstream per- plexity to pretraining perplexity. Max Thinking combines thinking longer and self-verification.
Next, we investigate the relation between OOD generalization and pretraining performance. To investigate this, we compare models with identical training setups with respect to data and parameters,
"Because we pretrained language models from scratch, and are unable to train models the size of modern foundation models, find models did not benefit from inference time techniques such Chain-of-Thon Thought.
8

---

Published as a conference paper at ICLR 2026
Performance Drop Out of Distribution Data
PPL Improvement Verification
Serf-Veeiiication Capabilities Scale During Training
Poplexity Increase (Lower la Better
45.0
Transformert=
14
BoN-5 Samples
42.5
EBT No Thinking
12
EBT Thinking Longer
EBT Seff- Verification
10
37.5
35.0
32.5
2.3
1.5
30
25
30
Number of Forward Passes
Tokens Trained On (B)
(a) OOD Thinking Performance Comparison.
(b) Verification Capabilities Scale Increases.
Figure EBT Thinking Analysis. (a) Mean performance degradation the standard Transformer++ recipe Tou- vron et al. (2023) and the Energy-Based Transformer (EBT) on four Dut-of-Distribution (OOD) datasets. The Transformer++, not being explicitly designed for inference-time computation, is unable to reduce perplexity per-token level. Alternatielly, EBTs can improve performance with more forward passes over single token/sample (Longer Thought) as well as generating many samples and choosing the minimum energy one (Self- Verificaation.. (b) The performance of EBTs with and without self-verification; as scale increases, the bene- fits from self-verification increase. These results suggest EBTs generalize OOD better than the Transformer++ because of their System capabilities, and that the thinking capabilities of EBTs improve during training.
Table 4: Image Denoising and Classification Comparison. For image denoising, EBTs significantly outperform DiTs Peebles & Xie (2023) in Peak Signal to Noise Ratio (PSNR), well as MSE, on both in-distribution and Dii-t---Dittibttiinn (OOD) data, while using 99% less forward passes. On image classification, EBTs also perform better than DiTs, yielding around 10x higher accuracy, suggesting that EBTs learn better image representations and therefore understand images better than DiTs.
In Distribution Noise
OOD Noise
Image Classification
Model
PSNR f
MSF Pixel
PSNR 1
MSE Pixel +
Top Acc.
Top Ace.
DiT
26.58
142.98 122.55
19.56 23.29
718.7 305.2
0.31% 5.32 %
1.36% 13.2 %
EBT
27.25
where EBTs have slightly worse pretraining perplexity than Transformer++ models. As shown in Table 3, despite achieving a higher pretraining perplexity, EBTs achieve lower (better) perplexity on most downstream tasks, suggesting stronger generalization, particularly to OOD data. Together, with the better learning scalability results, and knowing that improved pretraining performance usually leads to improved downstream task performance Chen et al. (2024); Isik et al. (2024), these results suggest that EBTs offer promise at scale during both pretraining and inference.
4.2 BIDIRECTIONAL IMAGE EXPERIMENTS
In addition investigating autoregressive EBTs, we explore the performance of EBTs trained bidirectionally. Following Chen et al. (2020); Du et al. (2022), models are trained to denoise images with fixed noise level. At inference, we test both the training noise level and higher OOD level. The results are in Table 4, where we observe that EBTs perform better than DiTs at both in and out of distribution image denoising across various metries. Following Chen et al. (2018), we plot the performance based on the number of forward passes (NFEs) in Figure B.6. These results demonstrate that EBTs perform better than DiT's while using 99% less denoising steps. Lastly, qualitative results for denoised ouitt--diitribution images for EBT compared to the DiT bascline are shown in Figure 8, demonstrating the improved visual quality of denoised images from EBTs.
In effort to understand whether the representations learned from denoising captured useful visual features, we perform a linear probe evaluation on ImageNet- 1k Russakovsky et al. (2015) of the models learned from denoising, following common practice in visual representation learning Oquab et al. (2023). The results are shown in Table 4, where the accuracy of EBTs around 10x higher than that of DiTs, demonstrating that EBTs learn better image representations than DiTs.

---

Published as a conference paper at ICLR 2026
sleps
sleps
GT Image
steps
steps
GT Image
Energy Based
Transformer
Diffuusio
Trrmerormer
100 steps
200 steps
300 sleps
GT Image
100 steps
steps
300 stops
GT Image
Forward Passes
Figure 8: Qualitative OOD Image Denoising. EBTs achieve better denoising quality during inference while using step for every 100 denoising steps of a DiT, or 99% less steps. The overall image quality of EBT denoised images less blurry than images denoised by DiT.
4.3 AATA-CONSTRAINED EXPERIMENTS
To determine how EBTs perform in data-constrained settings, we conduct experiments aimed at measuring generalization with small training set that is easy to overfit. We follow existing work benchmarking models on Sudoku Du et al. (2024), where the dataset from SAT-Net Wang al. 2) for the training set and the dataset from RRN Palm et al. (2018) as the test set. Models with discrete numbers in the range of 1 9 filled in, predict the rest of the Sudoku board that satisfies the necessary constraints. We report the accuracy on the test-set in Table 5, where we observe EBTs score the highest at 29.7%, RNNs score the second highest at 17.7%, and geed-Forward Transformers score the worst at 0.03%. More details on these experiments are discussed in Section D.
Table 5: Sudoku OOD Test Set Perfor- mance. We compare Feed-Forwwe Forward (FF.) Transformer, an RNN based on re- cent work Jolicocur-Maui Martineau (2025), and EBT on data-constrained al- gorithmic reasoning using Sudoku dataset Palm et al. (2018).
trained to, given initial Sudoku board
Architecture
Accuracy
Transformer
0.03%
RNN EBTs
17.7%
29.7 %
For more experiments on video, uncertainty estimation, and other topics, please refer to Section
LIMITATIONS AND CONCLUSION
Limitations. Despite demonstrating strong preliminary results, EBTs have several limitations. First, because EBTs generate predictions through an optimization process, they introduce additional hyperparameters. Second, while EBTs scale well up 800M parameters, larger models were unexplored due to resource constraints. Third, for distributions that are highly multimodal, such as images, EBTs with the current formulation struggle to capture the many modes, hence why we often combine EBT's with autoregression. Finally, current EBTss lag behind a large margin in -f..-efficieecy, posing high barrier to short-term adoption. Future researchers interested in leveraging EBTs will need weigh the tradeoff between using more computation and improved generalization and reasoning.
Transformers by
Conclusion. We introduced Energy- -ased Transformers (EBTs), a new approach that frames System 2 Thinking as an optimization procedure with respect to learned verifier (an Energy- Based Model), enabling System 2 Thinking to emerge across many problems and modalities from unsupervised learning. Across discrete and continuous modalities, scale at faster rate than the Transformer++ during pretraining across all measured axes, including data, batch size, depth, parameters, FLOPs, and width- with up 35% higher scaling rate. This suggests that EBTs offer promise at larger scale, even without System Thinking. With System Thinking, EBTs improve even further-iiccrasing performance by up 29% on text tasks, which we observe increases with data that is more Out-of-Distributtio (OOD). Comparisons to DiT's on image denoising also reveal significantly better thinking scalability: EBTs match or exceed DiT's performance with only 1% of the forward passes, EBTs also learn substantially better representations, achieving approximately 10x higher accuracy than DiTs. Ultimately, the improved scaling of EBTs during both training and inference positions them as a promising new approach.
results demonstrate that EBTs
10

---

Published as a conference paper at ICLR 2026
6 ACKNOWLEDGEMENTS
We extend special thanks to Jeonghwan Kim and Cheng Qian for their helpful discussions. This work based upon work supported by the U.S. National Science Foundation Graduate Research Fellowship Program under Grant No. DGE 21-46756, U.S. DARPA ECOLE Program No. #HR001 12390060, DARPA ITM Program No. FA8650-23-C-7316 NSF Molecule Maker Lab Institute, an Al Institute for Molecular Discovery, Synthesis Strategy. and Manufacturing funded by the U.S. National Science Foundation under Awards No. 2019897 and 2505932, the AI Research Institutes program by National Science Foundation and the Institute of Education Sciences, U.S. Department of Education through Award No. 2229873 AI Institute for Transforming Education for Children with Speech and Language Processing Challenges, and NSF NAIRR award. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the U.S. Government or the National Science Foundation. The U.S. Government authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein. This research used the Delta and DeltaAl advanced computing and data resources, which are supported by the National Science Foundation (award OAC 2320345 and award OAC 2005572) and the State of Illinois. Delta and DeltaAl are joint efforts of the University of Illinois Kirbana-Chapper and its National Center for Supercomputing Applications.
Reproducibility Statement We include the source code for full reproducibility of our results, which contains a detailed README on how to run and execute the code. We also include comprehensive experimental details for results in all Figures and Tables Sections 4 and D. This includes all necessary hyperparameters, seeding, and model configurations for full reproducibility.
REFERENCES
URL httltttss//eeeerreeee httpt:////seaarc..ggggg////o//// 21
2ttttappprrtt-22-chhllllll challi
Chirag Agarwal, Sree Harsha Tanneru, and Himabindu Lakkaraju. Faithfulness vs. plausibility: On the (un)
reliability of explanations from large language models. arXiv preprint arXi::240..04614, 2024. 34
Zaid Alkouri. Using contents and containers to investigate problem solving strategics among toddlers. 2016.
Anthropic. Claude 3.7 sonnet and claude code, 2025. URL. https: : //tw://wwwww/// http:://wwwwtt.oo.cccooooooo wwt.ttooo..ccooeeeoo..oo claude-3-1 3-nnoonnt.. Accessed: 2025-02-21. 1, 34
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. arXiv preprint
arXiv:2003.00033, 2020. 4.40
Anton Bakhtin, Yuntian Deng, Sam Gross, Myle Oit, Marc' Aurelio Ranzato, and Arthur Szlam. Residual energy-based models for text. Journal of Machine Learning Research, 22(40):1-41, 2021. 19,36
Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. The reversal curse: Llms trained on" is b" fail to learn" is a". arXiv preprint arXiv:2309. 2023. 19
Michael Betancourt. A conceptual introduction to hamiltonian monte carlo. arXiv preprint arXiv:1701.0244
2017. 20
Sumanta Bhattacharyya, Amirmohammad Rooshenas, Subhajit Naskar, Simeng Sun, Mohit lyyer, and Andrew McCallum. Energy-hased reranking: Improving neural machine translation using energy-based models. arXiv preprint arXiv:2009. 13267, 2020. 19,
Christopher M Bishop. Mixture density networks. 1994, 37
Sid Black, Stella Biderman, Fric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Cpt-neox-20b: An open-source autoregressive language model, 2022a. 31. 32
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et al. Gpt-neox-200: An open-source autoregressive language model. arXiv preprint arXiv:2204.067 06745, 2022b.
11

---

Published as a conference paper at ICLR 2026
Zalán Borsos, Raphaël Marinier, Damien Vincent, Eugene Kharitonov, Olivier Pietquin, Matt Sharifi, Dominik Roblck, Olivier Teboul, David Grangier, Marco Tagliasacchi, and Neil Zeghidour. Audiolm: language modeling approach to audio generation, 2023. 6
Adam Casson. Transformer flops.
2023.
URL httpp://aaaaccssoonccccoo/
trans former-flops 31,
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018. 7,9
Xinshi Chen, Hanjun Dai, Yu Li, Xin Gao, and Le Song. Learning to stop while learning to predict. In International conference machine learning, pp. 1520-15:0 1530- PMLR, 2020. 9
Yangyi Chen, Binxuan Huang, Yifan Gao, Zhengyang Wang. Jingfeng Yang, and Heng Scaling laws for predicting downstream performance llms. arXiv preprint arXiv:2410.08527, 2024. 9, 29
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, pp. 0278364924127:666 2023. 19
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jemry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 2021.
Together Computer. Redpajama: open dataset for training large language models, 2023. URL https: //github. com wtw.mm///tteeeeeeeooooooeeo P jmaa Data.
Stephen A Cook. The complexity of theorem-proving procedures. In Logic, automata, and computational complexity: The works of Stephen A. Cook, pp. 143-152. 2023.
Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How compute hessian-vector prod-
ucts? In ICLR Blogposts 2024, 2024. URL http:://////ictloooooppoooeeeee withhghhii///000//gggheeeeaa-a/ http:://iuuuuuoooo......iuoooood bench hvp/. https://iclt-blopooste
gighhu..b.... 1c/202//bllg/ 5, 27, 37,31,33
Anna Dawid and Yann LeCun. Introduction to latent variable energy-based models: path toward autonomous machine intelligence. Journal of Statistical Mechanics: Theory and Experiment, (024(10):10401, 2024. 5, 37, 38
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv: 180700//07 2018. 34
Chaorui Deng, Deyao Zhu, Kunchang Li, Shi Guang, and Haoqi Fan. Causal diffusion transformers for generative modeling. arXiv preprint arXiv:2412.1209 2024a.
Haoge Deng. Ting Pan, Haiwen Diao, Zhengxiong Luo, Yufeng Cui, Huchuan Lu, Shiguang Shan, Yonggang Qi. and Xinlong Wang. Autoregressive video generation without vector quantization. arXiv preprint arXiv 14169, 2024b. 21
Jacob Devlin, Ming- Wei Chang, Kenton Lec, and Kristina Toutanova. Bert: Pre-training of deep bidircctional transformers for language understanding, 2019,
Jochen Ditterich. Evidence for time-variant decision making. European Journal of Neuroscience, 24(12): 3628-3664 2006.
Rodney Douglas and Kevan AC Martin, Recurrent neuronal circuits in neocortex. Current biology. 17(13): R500,,500, 2007. 20
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019. 3, 4, 5, 6, 20, 40
Yilun Du, Shuang Li, Joshua Tenenbaum, Igor Mordatch. Improved contrastive divergence training of energy based models. arXiv preprint arXi::2012200336 2020. 6, 40
Yilun Du, Shuang Li. Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference Machine Learning, pp. 5570-5582, PMLR, 2022. 4, 9, 31, 32, 35
12

---

Published as a conference paper at ICLR 2026
Yilun Du, Conor Durkan, Robin Strudel, Joshua Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl- Dickstein, Arnadd Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-basee bagy- diffusion models and meme. In International conference on machine learning, pp. 8489 8510. PMLR, 2023. 34, 35, 37
Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion, arXiv preprint arXiv:2406. / / 2024. 4, 10, 25, 35, 36
Jonathan BT Evans. Dual-process theories reasoning: Contemporary issues and developmental applications. Developmental review, (12-)))86-102, 2011.
William A Falcon. Pytorch lightning. GitHub, 3, 2019, 28
Keith Frankish. Dual-process and dual-system thcories of reasoning. Philosophy Compass, 5110:9114-926 2010.
Angela D Friederici and Jürgen Weissenborn. Mapping sentence form onto meaning: The syntax semantic interface, Brain research, 144:500-58, 2007. 37
Jonas Geiping, Sean Mcl.cish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-lime compute with latent reasoning: recurrent depth approach. arXiv preprint urXiv:2502.051771 2025. 2, 23, 34, 39
Xavier Glorot and Yoshua Bengio. Understanding the difficulty training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256.
JMLR Workshop and Conference Proceedings. 2010. 41
Vinod Goel, Christian Buchel, Chris Frith, and Raymond J Dolan. Dissociation of mechanisms underlying
syllogistic reasoning. Neuroimage, 12(5):504-1 501444444, 2000. 1
Shafi Goldwasser, Silvio Micali, and Chales Rackoff. The knowledge complexity of interactive proof-systems. In Providing sound foundations for cryptography: On the work of shufi goldwasser and silvio micali, pp.
203-225.220 2019.
lan J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde- Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.5 5
Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski. Joanna Materzynska, Susanne Westphal, Heuna Kim. Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag. et al. The" something something" video database for learning and evaluating visual common sense. Proceedings of the IEEE international conference on computer vision, pp. 5842-5850, 2017. 21
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788. 2024. 8, 23, 24
Albert Gu and Tri Dao. Mamba: Linear-tims sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0m 200075222 2023. 2. 7, 20, 29, 30, 34
Yuchao Gu, Weijia Mao, and Mike Zheng Shou. Long-context autoregressive video modeling with next-frame prediction. arXiv preprint arXiv:2503. 22.225503,.,,,,,. 2025. 21
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma. Peiyi Wang, Xiao Bi, et al. Deepseek-r1; Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint 2xxv:25011.2944 2025. 1, 34
Kurt Gödel. Letter to john von neumann, 1956. URL tttpss:t////ooomooo cornell www.://moooo..oee...eeeeeeeeeeee api ctce//rr/b////// / teatttss//eeeeeeeeee.ttatt 4 6ae - 28d---57--------- 203-57dd--3e--bbltte bb6chl abate//ootee... Accessed: 2025-04-28. 3
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models reason in continuous latent space. arXiv preprint arXiv:2412.06769, 2024. 2, 34
Alvin Heng, Harold Soh, et al. Out-of-distribution detection with single unconditional diffusion model. Advances Neural Information Processing Systems, 37:43952-43974 2024. 36
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural
information processing systems, 33:6840-6851, 2020. 34
13

---

Published as a conference paper at ICLR 2026
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022. 7, 29
Benjamin Hoover, Yuchen Liang. Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horng Chau, Mohammed Zaki, and Dmitry Krotov. Energy transformer. Advances in Neural Information Processing Systems, 36, 2024. 36
John Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):555-2558, 1982. 39
Shengding Hu, Xin Liu, Xu Han, Xinrong Zhang, Chaoqun He, Weilin Zhao, Yankai Lin, Ning Ding, Zebin Ou, Guoyang Zeng, et al. Predicting emergent abilities with infinite resolution evaluation. arXiv preprint arXiv:2310.03262, 2023. 29
Kaiyi Huang. Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-omppeench: A comprehensive benchmark open- world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723-78747, 2023. 37
Tobias Höppe, Arash Mchrjou, Stefan Bauer, Didrik Niclsen, and Andrea Dittadi. Diffusion models for vidco prediction and infilling, 2022, 34
Itay Inbar and Luke Sernau. Time matters: Scaling laws for any budget. arXiv preprint arXiv:2406. 18922 2024. 30
Berivan Isik, Natalia Ponomareva, Hussein Hazimeh, Dimitris Paparas, Sergei Vassilvitskii, and Sanmi Koyejo. Scaling laws for downstream task performance of large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. 9
Md Mofijul Islam, Alexi Gladstone, Riashat Islam, and Tariq lqbal. Eqa-mx: Embodied question answering using multimodal expression. The Twelfth International Conference on Learning Representations, 2023.
21
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Cameey, al. Openai ol system card. arXiv preprint arXiv:2412.16720, 2024. 1, 24, 34
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991, 2022. 19
Alexia Jolieoeu-Marttnnau Less is more: Recursive reasoning with tiny networks. arXiv preprint
WwXiv:2510.0487/ 2025. 10,
Daniel Kahneman. Thinking, fast and slow. macmillan, 2011. 2, 25
Daniel Kahneman, Shane Frederick, et al. Representativeness revisited: Attribute substitution in intuitive judgment. Heuristics and biases: The psychology of intuitive judgment, 49(49-81):74, 2002.
Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Mudit Verma, Kaya Stechly, Siddhant Bhambri, Lucas Paul Saldyt, and Anil B Murthy. Position: Llms can't plan, but can help planning in llm-modulo frameworks. In Forty-first International Conference on Machine Learning, 2024. 4
Jared Kaplan, Sam McCandlish, Tom Hemighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001. 08361, 2020. 7, 29
Diederik P Kingma, Max Welling, et al. Auto-encoding variational bayes, 2013. 33,
Seijin Kobayashi, Simon Schug, Yassir Akram, Florian Redhardt, Johannes von Oswald, Razvan Pascanu, Guillaume Lajoie, and João Sacramento. When can transformers compositionally generalize in-context? arXiv preprint arXiv:2407.12273 2024. 37
Siddique Latif, Aun Zaidi, Heriberto Cuayahuitl, Fahad Shamshad, Moazzam Shoukat, and Junaid Qadir. Transformers in speech processing: A survey, arXiv preprint arXiv:2303.1160) 2023. 33
Ryan Lavin, Xuekai Liu, Hardhik Mohanty, Logan Nomma, Giovanni Zaarour, and Bhaskar Krishnamachari. survey on the applications of dero-knowledge proofs. arXiv preprint arXiv:2408.00 00043, 2024.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-00-27. Open Review, 62, 2022. 5, 26
14

---

Published as a conference paper at ICLR 2026
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, al. A tutorial energy-based learning. Predicting structured data, 1(0), 2006. 35
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape neural nets. Advances in neural information processing systems, 31, 2018. 4
Margaret Li, Sncha Kudugunta, and Luke Zettlemoyer. (mis) fitting: A survey of scaling laws. arXiv preprint 2tt.:2552..89999 2025a. 29
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. Advances Neural Information Processing Systems, 37:56424-56445, 2025b. 2,
Zengyi Li, Yubei Chen, and Friedrich T Sommer. Learning energy-based models niihh-dimensiona spaces with multiscale denoising-scon matching. Entropy, 25(10): 2023. 6, 40
Zhong-Zhi Zhi Li, Duzhen Zhang, Ming-Liang Zhang, Jiaxin Zhang, Zengyan Liu, Yuxuan Yao, Haotian Xu, Junhao Zheng, Pei-Jic Wang. Xiuyi Chen, et al. From system 1 to system 2: A survey of reasoning large language models. arXiv preprint arXiv:2502.177 12200...74199 2025c.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023. 34
Tianhe Lin, Jian Xie, Siyu Yuan, and Deqing Yang. Implicit reasoning in transformers is reasoning through shortcuts. arXiv preprint ar 2503::2503.77044 2025. 34
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrrence Zitnick. Microsoft coco: Common objects in context. In Computer vision-ECCV 2014: 13th European conference, zurich, Switzerland, September 6-12. 2014, proceedings, part v 13, pp. 740-755. Springer. 2014. 30
Fangfu Liu, Hanyang Wang, Yimo Cai, Kaiyan Zhang, Xiaohang Zhan, and Yueqi Duan. Video-tl: Test-time scaling for video generation. arXiv preprint 2rXiv:2503. 18942, 2025. 34
Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. In European Conference on Computer Vision, pp. 423-439. Springer, 2022. 35
Frank Loesche, Jeremy Goslin, and Guido Bugmann. Paving the way to cureka introducing "dira" as an experimental paradigm to observe the process of creative problem solving. Frontiers in Psychology, 9:1773, 2018. 2
Nanye Ma. Shangyuan Tong, Haolin Jia. Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, al. Inference-time scaling for diffusion models beyond scaling denoising steps. arXiv preprint Wtw.:2551.097322 2025. 4, 7, 24, 34, 35, 36
Rohin Manvi, Anikait Singh, and Stefano Ermon. Adaptive inference-time time compute: Lims can predict if they can do better, even did-generatiio. arXiv preprint arXiv:2410.02725, 2024. 24
Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar. Gsm-symbolic: Understanding the limitations of mathematical reasoning in large language models. arXiv preprint arXiv:2410.05229, 2024.
Erie Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan. Do deep generative models know what they don't know? arXiv preprint arXiv: 180000911, 2018. 37
Wim De Neys. Dual processing in reasoning: Two systems but one reasoner. Psychological science, 17(5): 428 433, 2006.
OpenAL Learning to reason with 1ls,, 2024. bearning-to-reason witbon-wwtthhliee Accessed: 2025-02-21. 2
URL hctps://openai.ee index/
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo. Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaacldin El-Nouby. Mahmoud Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Yao Huang. Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve. Hu Xu, Hervé Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2023. 6, 9, 33
15

---

Published as a conference paper at ICLR 2026
Rasmus Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. Advances in neural information processing systems, 31, 2018. 10
Thomas Pamr Giovanni Pezzulo. and Karl J Friston. Active inference: the free energy principle in mind, brain, and behavior. MIT Press, 2022. 25
Adam Paszke, Sam Gross, Francisco Massa, Adam Lever, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et Pytorch: imperative style, high high-performance deep learning library, Advances in neural information processing systems, 32, 2019. 26
Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation. 6(1):147-160, 1994. 27, 31, 32
William Peebles and Saining Xie. Scalable diffusion models with transformers, 2023. 2, 9, 28, 30, 31, 41
Jiahuan Pei. Cheng Wang, and György Szarvas. Transformer uncertainty estimation with hierarchical stochastic attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 11147-11155, 2022. 22
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, I. Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:3081 11130081-33000 1-30849, 2024. 29
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung. Matteo Grella. Kranthi Kiran GV, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305. 13048, 2023. 2, 34
A. Peters, B. McEwen, and Karl J. Friston. Uncertainty and stress: Why it causes diseases and how it is mastered by the brain. Progress in Neurobiology, 156: 2017. doi: 10. www.r/poeeroooo..... 201.....01.....000 37
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by
generative pre-training. 2018. 6
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18))99 2019, 6, 7,
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Perey Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250, 2016.
Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. Latent video
transformer. arXiv preprint arXiv:2006. 10704, 2020. 21
Pradeep Ravikumar, Alekh Agarwal, and Martin J Wainwright. Message-passing for graph-structured linear programs: Proximal methods and rounding schemes. Journal of Machine Learning Research, 11(3), 2010. 26
Ronald I. Rivest, Adi Shamir, and I Leonard Adleman. A method for obtaining digital signatures and public-key cryptosystems. Communications of the ACM, 2102):100-126, 1978. 3
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- reshrrcsolutin image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10684-10695, 2022. 21, 31, 32, 34,38
Nicolas Rougier, David C Noelle, Todd Braver, Jonathan D Cohen, and Randall C C''Reilly, Prefrontal cortex and flexible cognitive control: Rules without symbols. Proceedings of the National Academy of Sciences, 102 (20):7338-7343 2005. 2
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenct large scale visual recognition challenge. International journal of computer vision, 11 15::211-552, 2015. 9
Karthik Abinav Sankararaman, Sinong Wang, and Han Fang. Bayesformer: Transformer with uncertainty estimation. arXiv preprint arXiv:2206.0008 2022. 37
Issidoros C. Sarinopoulos, D. Grupe, Kristen L. Mackiewicz, J. Herrington, M. Lor, E. E. Steoge, and J. Nitschke. Uncertainty during anticipation modulates neural responses to aversion in human insula and amygdala. Cerebral cortex, 20 4:929-40, 2010. doi: 10. 10.0099/cerco//bp155.. 37
Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank J Reddi. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arxiv:2502. / 17416, 2025. 34, 35
16

---

Published as a conference paper at ICLR 2026
Rylan Schaeffer. Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? Advances Neural Information Processing Systems, 555585551111,, 2023, 29
Joan Serrà, David Álvarez, Vicenç Gómez, Olga Slizovskaia, José F Núñez, and Jordi Luque. Input complexity and out-of-distribution detection with likelihood--asee generative models. arXiv preprint arXiv:1999 1480, 2019. 37
Noam Shazeer. Glu variants improve transformer, 2020, 41
Raghay Singhal, 7achary Horvitz, Ryan Techan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference-time scaling and steering of diffusion models. arXiv preprint arXiv:2501.066 068448 2025. 34
Charlic Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling 1lm test time compute optimally can be more effective than scaling model parameters. arXiv preprint arXhh:2408.033// 2024. 24
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 2020. 37
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonsa, et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint 2rXiv:2206.046115 2022. 7,8
Stability. URL wttttt::://tgggggegeeeeee http://aw..........aaaaaaaaaaaaa wttttaaaiitttyyytiaaaaaaa
ft-mae. 21, 31,
Nisan Stiennon, Long Ouyang. Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in neural information processing systems, 33:3008-3021, 2020. 6
Gala Stojnić, Kanishk Gandhi, Shannon Yasuda, Brenden M Lake, and Moira R Dillon. Commonsense psychology in human infants and machines. Cognition, 235:105440,, 2023. 4
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position emhedding. Neumcompurinng 568:12706; 2024. 41
Yi Su, Dian Yu, Linfeng Song, Juntao Li, Haitao Mi, Zhaopeng Tu, Min Zhang, and Dong Yu. Expanding with verifiable rewards across diverse domains. arXiv preprint arXiv:2503.2 23822, 2025.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rons with expressive hidden states. arXiv preprint arXiv:2407.04622 046200 2024." 7
Gokul Swamy, Sanjiban Choudhury, Wen Sun, Zhiwei Steven Wu, and J Andrew Bagnell. All roads lead to likelihood: The value of reinforcement learning in fine-tuning, arXiv preprint arXiv:2503.01067 2025, 4
AbdeIrahman Tawfik. URL
co/ http://taaseett. http://gggggggaecocooooaaaaaaaaa httpttaaett//bbbb//////////
30
AlphaCode Team, Alphacode technical report. December 2023. 4
Christian Tomani, Kamalika Chaudhuri, Ivan Evtimov, Daniel Cremers, and Mark Ibrahim. Uncertainty-bart based abstention in Ilms improves safety and reduces hallucinations. arXiv preprint arXiv:2404. 10960, 2024. 36
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu. Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa. Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lce, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nic, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Erie Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu. Zheng Yan, Iliyan Zarov, Yuchen Zhang. Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. 6, 7, 9. 28, 33, 41
Miles Turpin, Julian Michael, Ethan Perez, and Samuel Bowman, Language models don't 't always say what they think: Unfaithful explanations in chain-of-though prompting. Advances in Neural Information Processing Systems, 36:74952-74966 2023. 34
17

---

Published as a conference paper at ICLR 2026
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017. 33, 36
Ashish Vaswani, Noam Shazeer, Niki Parmar. Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 6, 27, 33
I. Vilares, J. D. Howard, Hugo L. Fernandes, J. Gottfried, and Konrad Paul Kording. Differential representations of prior and likelihond uncertainty in the human brain. Current Biology, 22: 22:1641-1648, 2012. doi: 10.1016/j. sub.2012.07.0144 37
Pablo Villalobos, Jaime Sevilla, Lennart Heim, Tamay Besiroglu, Marius Hobbhahn, and Anson Ho. Will we run out of data? an analysis of the limits of scaling datasets in machine learning. arXiv preprint arXiv:2211.04325
2022. 21
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7): 1661-1674, 2011. 36, 39
Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori. Hierarchical reasoning model, arXiv preprint xxXiv:556.22/74 2025. 34
Po-Wei Wang. Priya Donti, Bryan Wilder, and Zico Kolter. Satnet: Bridging deep learning and logical reasoning using differentiable satisfiability solver. In International Conference on Machine Learning, pp. 6545-6554. PMLR, 2019, 10
Yezhen Wang. Tong Che, Bo Li, Kaitao Song, Hengzhi Pei, Yoshua Bengio, and Dongsheng Li. Your autoregres- sive generative model can be better if you treat as energy -ased one. arXiv preprint arXiv:2206.1288 12840. 2022. 36
Ze Wang, Jiang Wang, Zicheng Liu, and Qiang Qiu. Energy- self supervised pretraining for vision models. arXiv preprint arXiv:2302.011 384,, 2023. 4, 5, 36, 39, 42
Maurice Weber, Dan Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, et al. Redpajama: an open datasct for training large language models. Advances in neural information processing systems, 37:1 177:11662--116 11649,, 2024.7 7
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Fd Chi, Quoe V Le, Denny Zhou, al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information
Processing Systems, 35:24824-24837, 2022. 34
Dirk Weissenbom, Oscar Täckström, and Jakob Uszkoreit, Scaling autoregressive video models, arXiv preprint arXiv: 02034, 2019. 21
Peter West, Ximing Lu, Nouha Dziri, Faeze Brahman, Linjie Li, Jena D Hwang. Liwei Jiang, Jillian Fisher, Abhilasha Ravichander, Khyathi Chandu, et al. The generative ai paradox:" what it can create, it may not understand". arXiv preprint 3xkiv:331100005: 2023.
xAI. Grok 3 Beta- - The Age of Reasoning Agents, 2025. URL https: http://///xi///////////// grok- 3. Accessed: 2025-02-21. 34
Yang Yan, Yu Lu, Renjun Xu, and Zhenzhong Lan. Do phd-level Ilms truly grasp elementary addition? probing rule learning vs. memorization in large language models. arXiv preprint arXiv:2504.05260 2025. 1,
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023. 4
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade- school math and the hidden reasoning process. In The Thirteenth International Conference on Learning Representatiinss 2024. 6
Xi Ye and Guillaume Alexandre Bilodcau. Video prediction by efficient transformers. Image and Vision Computing, 130: 104612, 2023. 21
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does rein- forcement learning really incentivize reasoning capacity in Ilms beyond the basc model? arXiv preprint 2xXiv:2504.1387,, 2025.
Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. Dino-wm: World models on pre-trained visual features enable zero-shot planning. arXiv preprint arXiv22411.04988 2024. 19
18

---

Published as a conference paper at ICLR 2026
In this appendix, we provide additional insight and details on EBTs. First, we provide more insight on the broader impact/future work of EBTs in Section A. Next, in Section B, we include additional experiments. Then, we include additional approach details in Section C, as well as additional experimental details in Section D. After that, we include a comprehensive related work in Section E, additional facets of cognition in Section F, and discussion of counterarguments in Section G. Finally, in the hopes of making EBTs more accessible to general audiences, in Section H we contain general casier-oo-understand intro to EBMs, and in Section I we describe a tutorial for getting started with EBTs.
Large Language Model Usage Statement: We use Large Language Models (LLMs) to assist with grammar, formatting, and writing clarity/flow in the paper. All intellectual and experimental contributions were made by the authors.
A FUTURE WORKS AND BROADER IMPACT
A.I REVERSAL CURSE
Recently, phenomenon known as "The Reversal Curse" has been observed where LLMs fail to learn symmetric mapping Berglund et al. (2023) "B is A" despite learning "A is B". For example, LLMs trained on an example such as "Q: Who Tom Cruise's s mother? A: Mary Lee Pfeiffer" often fail to generalize to know the answer to the reverse question of "Who is Mary Lee Pfeiffer's son?" Remarkably, the Reversal Curse has manifested itself in LLMs regardless of the size or scale Berglund et al. limitations to traditional feed-forward LLMs. One predicted cause of The Reversal Curse is the nature of gradient updates, where backpropagation only updates tokens within context. That is, while learning the mapping "A is B", none of B's tokens are within context, meaning they do not receive gradient updates when updating A's tokens. We hypothesize that LLMs trained with EBTs rather than standard feed-forward Transformers could help reduce this phenomenon, as with EBTs the tokens of and B are within context during gradient updates due to predictions being made in the input space. Therefore, an exciting research direction would be investigating whether this hypothesis is correct in LLMs trained with EBT, allowing improved generalization.
researchers to investigate whether there are fundamental
A.2 IMPROVED STABLITY
In this work, we primarily trained EBT with either two or three optimization steps. While these parameters worked well, we suspect that increasing the number steps would improve the System 2 Thinking capabilities and the scaling during pretraining of EBTs, as more steps enable longer "thinking process" before needing to converge. However, because of challenges in stability when training with more steps, due larger gradient computation graph, we unable to successfully increase past two or three steps. Future work could focus more on extending the number of steps by studying ways to improve the training stability of EBMs.
A.3 WORLD MODELS
In this work, we focus on autoregressive and bidirectional models over just state information (no actions). EBTs offer high promise in modeling states and actions due to the nature of EBMs learning a distribution over all possible inputs. Particularly, given model trained to estimate the unnormalized joint distribution of the current context, future, as well as future actions, such world models could implicitly be used as policies to generate actions to achieve a specific state, similar to Chi et al. (2023); Janner et al. (2022); Zhou et al. (2024). This would involve holding the current context (past states) constant, and minimizing the energy by propagating the gradient back to the action inputs and future state predictions. Thus, world models trained in this manner become capable of more than just predicting the future, but also in decision making to achieve a specific goal state.
A.4 EBTS AS COMPLEMENTARY MODELS
As demonstrated in Bakhtin et al. (2021); Bhattacharyya et al. (2020), EBMs can be used to improve the quality of generated text from language models. It's possible EBTs could be used in similar
19

---

Published as a conference paper at ICLR 2026
manner for broad variety of tasks, serving the verifier of predictions initialized by standard feed-forward models. Therefore, although we do side-by-side comparison to existing models in this work, EBTs could be complementary to existing modeling approaches being used as the System Backbone for helping lighter models that perform System ttinking.
There exist several current real-world cases, such low-latency LLM serving, where doing single forward pass is sufficient, and where the added inference overhead of gradients with EBTs would be worth the extra computation. However, also envision world in which people use EBTs for long-term System 2 Thinking to solve challenging problems. How much computation would it be worth dedicating to prove long-standing mathematical conjecture, or figuring out cure to cancer?
A.5 RECURRENT ENERGY-BASE MODELS
While EBTs scale well, for latency-driven use cases, Transformers require significantly more memory than Recurrett Neural Networks. Additionally, there is strong evidence for recurrence in the human brain Douglas & Martin (2007). Therefore, we anticipate that recurrent Energy-Bass Based Models, possibly leveraging the Mamba architecture Gu Dao (2023) will eventually become common.
A.6 IMPROVED THINKING AL.GORITHMS
The EBM thinking algorithms described have strong connections
derived from Markov Chain Monte Carlo (MCMC) sampling. Therefore, we broadly expect known MCMC samplers with more advanced techniques for traversing the energy landscape be successful, such Hamiltonian Monte Carlo Betancourt (2017) or annealed Langevin dynamics Du & Mordatch (2019). Additionally, we did not explore more advanced search algorithms such as Monte Carlo Tree Search, which suspect could offer performance improvements and leave for future work.
A.7 MULTIMODAL ENERGY-BASED MODELS
We did not experiment with multimodal EBMs, however, EBMs offer several advantages for learning over multiple modalities. For example, multimodal EBMs would enable single energy scalar to represent the alignment between modalities, and would simplify joint training across modalities by providing unified objective that naturally captures inter-modal dependencies.
A.8 THINKING SCALABILITY
Due to a lack of computational resources, we were unable to train models with more than 102 FLOPs (<s 1300 A100 GPU Hours). Therefore, training and thinking with EBTs remains untested at larger foundation model scale. We leave to future work to scale with more GPUs and investigate the qualitative differences in training and thinking with EBTs.
A.9 LEARNING MULTIMODAL DISTRIBUTIONS
We found that EBTs, with the current training approach, struggle to capture distributions with many modes (e.g.. unconditional image generation or unconditional video generation). We believe this has to do with the current intimization-basee formulation used during training. Therefore, future work could explore approaches to improve the learning of distributions with many modes, via changing the training formulation support more modes. More info in Section B.3.
A.10 LEARNING BIDIRECTIONAL EBTs OVER ENTIRE SEQUENCES
Because of the described challenges with training EBTs for very multimodal distributions requiring uncertainty, we often trained EBTs combined with autoregression. Future work could focus on bidirectional EBT' operating entire sequences, resulting in verification being meaningful.
20

---

Published as a conference paper at ICLR 2026
A.11 UNDERSTANDING PREDICTIONS
We posit that there exists fundamental distinction between the internal representations associated with model inputs and outputs. Specifically, models generate internal representations of inputs, as these serve the foundation that dictate the model's behavior at any given point in time. Conversely, as outputs do not affect model's behavior at a given point in time, we contend that models construct representations for predicting (and not necessarily understanding) outputs. This distinction leads us to a consequential insight: models may not achieve a genuine understanding of outputs in the same way they understand imputs. This implies that while models may develop intricate understanding of input data, such understanding does not naturally extend to predictions that are made the output space. Therefore, existing feed-forward models primarily making predictions in the output space may not understand their predictions in the same way they understand their inputs. This intuition further supports the principles behind EBT, where predictions are made in the input space enabling representations of predictions to be developed. We leave the investigation of this hypothesis to future work.
A.12 SOCIETAL IMPACT
The ability to achieve human-like thinking with AI offers benefits in multiple domains. As such, EBT's offer several potential positive impacts, through the enabling of AI to potentially think more like humans. On the other hand, it's also possible that more intelligent AI models trained using EBT could be misused for harm by malicious actors.
B ADDITIONAL EXPERIMENTATION
B.1 AUTOREGRESSIVE VIDEO EXPERIMENTS
To assess how EBTs scale in continuous domains, train models to predict the next image in video conditioned on all previous frames common pretraining objective for video models Deng al. (2024b); Gu al. (2025); Rakhimov al. (2020); Weissenborn al. (2019); Ye Bilodeau (2023). Unlike in the NIP experiments, where models see each sample only once due to the dataset size, current popular video datasets are relatively small, requiring models to train repeatedly on the same data. As result, this setting probes a different question: "how well can models fit fixed dataset?", rather than how efficiently models scale under non data-bound regimes. This distinction is especially important given the recent scarcity of high-quality datasets Dat; Villalobos et al. (2022) and the perspective that data will increasingly become a bottleneck.
For experiments, encode all 224 x 224 images into 3136 dimensional features with the frozen SD-XL VAE Rombach et al. (2022); Stability. Then, all models are trained using a Smooth LI loss with Bc = 1.0 on the Something Something V2 dataset Goyal al. (2017), where we report the minimum validation loss achieved. In Figure B. we report scaling results for the embedding dimension and non-embedding parameter count, as we found these axes behaved the most linearly. The results demonstrate that, despite achieving higher initial loss, EBTs scale at a more than 33% faster rate than the Transformer++. This suggests that at larger scale, EBTs could achieve better performance than the Transformer++
We believe this large scaling rate gap can be linked to the fact that EBTs more seamlessly model continuous distributions than standard feed forward transformers due to being able to express uncertainty (Facet 3) through their energy scalar. To confirm this, we visualize results for different energies when predicting video frames in Figure B.2. The results demonstrate that EBT's successfully learn to capture uncertainty-where frames earlier on in the video have higher energy (higher uncertainty) due to no large objects being within the frame, and then the major object in the scene becomes revealed more EBT predicts lower energy (lower uncertainty). EBTs learn to exhibit this behavior without any supervision using Smooth L1 loss, whereas the standard forward Transformer++ would require discretization schemes such as Vector Quantization Islam et al. (2023) with categorical loss, or other tricks to achieve the same effect.
21

---

Published as a conference paper at ICLR 2026
Scales 33.66% Faster Than Transforeer+
EBT Scales 34.28% Faster Than Transformert+
1.9
Transfomner++
Transformert+
EBT
EBT
1.8
Loss (log seale)
1.7
Loss
1.6
1.5
1,5
400
500
600
700
8009000
1526
2000
10
100
Embodding Dimension scale)
Non-Embeediing Parameters log scale)
Scaling for the embedding dimension.
(b) Scaling for the Non-Em Con-Fmhedding Paramcters.
Figure B.I: Video Learning Wiith and Parameters. The minimum validation loss achieved on the Something Something V2 (SSV2) dataset. While EBTs achieve higher validation loss than the +ransformer++ smaller scales, the scaling rate is than 33% higher, suggesting that at larger scales EBTs could perform better than the Transformert++ Notably, scaling with respect the embedding dimension behaves more linearly than for the number of parameters, likely due to the embedding dimension serving as a bottleneck for the image representation.
Frame Energies Across Thinking Steps
1.0
0.8
Notmalizod Easayy
Frame
0.2
0.0
Iteration
Figure B.2: Learning Uncertainty on Video Results. In line with cognitive Facet 3. EBTs learn to express uncertainty across continuous video frames without supervision. At the start of the video, uncertainty high (high energy) because the frame mostly empty and the scene is highly unpredictable. As blue garment is placed into the frame, uncertainty decreases (low energy), reflecting the greater predictability of the scene. When the blue garment is removed from the scene, uncertainty increases again, indicating return to higher unpredictability. Such capability is significantly more difficult to achieve in continuous spaces with traditional feed-forward transformers without discretization schemes Pei al. (2022).
B.2 ADDITIONAL NATURAL LANGUAGE PROCESSING EXPERIMENTS
We conduct experiments to confirm hypotheses on thinking results obtained in the main paper. First, we confirm that EBTs become less adversarial with scale by comparing the performance of using Best-of-N (BoN) with 2 versus 10 samples. The results in Figure B. demonstrate that when models are trained on less tokens, there is little performance improvement by verifying 10 samples instead of just 2. In fact, verifying 10 samples occasionally leads to worse performance than verifying samples, likely because the EBT found an adversarial sample (a sample with low energy that is in fact not a good prediction). However, as data seale increases we observe that performance improvements from BoN-10 10 versus BoN-2 increase, and that these adversarial dynamics decrease. Together, these results suggest that with scale EBTs become less adversarial due to an improved energy landscape. In effort to understand the impacts of thinking the scale of modern foundation models,
22

---

Published as a conference paper at ICLR 2026
Token Energies Across Thinking Steps
Token Energies Across Thinking Steps
1.0
The
System
quick
brown
Thinking
0.8
Fox
Normalized Einegyy
is
Normalized Energy
0.6
0.6
Token
jumps
Token
challenging
over
0.4
but
0,
the
interesting
lazy
0.2
dog
problem
10 0.0
0.0
0
6
8
10
Iteration
Iteration
Figure B.3: Learning Uncertainty Text Results. EBTs learn to vary uncertainty across text tokens without any explicit supervision. As example, in both (a) and (b), simple tokens such "." "is", "a", "but", or "the" have lower energies across inference- time optimization (thinking) steps, indicating lower uncertainty. On the other hand, harder to predict tokens such as "quick". "brown", "research", and "problem" have higher energies across optimization steps, and more difficulty in achieving energy convergence, meaning the model more uncertain. Inspired by Geiping et al. (2025).
project results from Figure 7b the data scale of modern foundation models Grattafiori et al. (2024) to extrapolate a projected performance gain from self-verification based thinking. The results are visualized in Figure B.4b, where, while being highly speculative, they demonstrate that, because of the 1000x data scale modern foundation models, the performance improvement from verficccriination increases. While these results remain largely speculative, they demonstrate the potential for System 2 Thinking with EBT's that improves with scale.
Additionally, visualize results from EBT at representing uncertainty while predicting tokens in Figure B.3. The results demonstrate that for easier to predict tokens, such as "the" or "but", EBTs optimize to lower energies faster, whereas for harder to predict tokens, such "fox" or "problem" EBTs have higher energy that does not converge across steps. This suggests that during pretraining EBTs learn to capture uncertainty regarding which tokens are harder or easier to predict, achieving Facet 3. In addition to learning alcatoric uncertainty (uncertainty related to noise). in an effort to understand whether EBTs can capture epistemic uncertainty (uncertainty related to knowledge), we visualize the energies of different token sequences that are in versus nut-OO-Distributiin (OOD) in Figure B.5. The results demonstrate that, for a more in-distribution sequence, EBTs have lower energy (less uncertainty), than for an OOD sequence. This suggests that EBTs learn to know what they don't know, as they learn to have higher energy for OOD sequences signifying harder predictability. This is promising characteristic of EBTs, as it enables uncertainty estimation within continuous state spaces, enabling principled inference-time behavior adaptation (Facet 1) when models sense a more challenging problem.
B.3 EBT FAILURE CASES
During experimentation, along with image denoising experiments, also conducted small-scale experiments with text-to-image generation. We found that, in datasets such as COCO with many different modes for a single condition (e.g., hundreds images with a caption similar "giraffe with long neck"), ) that EBT's did not learn to generate high-quality novel images. Instead, EBT's often generated blurred images similar to the training distribution. We believe this is caused by the training approach pushing the energy landscape to have single local minimum surrounding the training examples. Therefore, when there are many different modes, these local minima within the energy landscapes "merge"
minima averaged between the different modes, resulting in blurriness. We believe that this is not a fundamental limitation of EBTs, and that future work could address this issue.
23

---

Published as a conference paper at ICLR 2026
% PPI. Improvemeeetttnooom Over Verifleation
Self- Velifieeriicntionn Capabilitiss Extrapolated
Self-Verification Capablities Scale During Training
1000
BoN-5 Samples
2,5
BoN-10 Sumples
Samples (extrapolated)
% Impooveneee PPL
2.0
1.5
100
1.0-
0.5
10
0.0
10'
10'
Tokeus Trained on (log scale)
Tokens Trined On (H)
(b) Results Fig. 7b projected to Llama3
(a) Self-verification with BoN-10 versus BoN-2.
scale Grattafiori et al. (2024).
Figure B.4: EBT Thinking Analysis for Data Scaling, (a) Self- verification BoN-2 compared to BoN- 10. EBTs become less adversarial and thus benefit more from verifying an increasing number of samples during training. (b) A projection of the results from Figure 7b to the data scale of Ilama3 Grattafiori et al. (2024), demonstrating that data scale increases, improvements from self-verification lead potentially massive performance increases from System Thinking.
Energy Comparison for Different Sequences Sequence
Sequence 2
The
NO
1.0
quick
0.8
brown
Normalized Energy
0.6
Token
jumps
over
0.4
the
lazy
0.2
dog
0.0
0
2
4
10
2
4
6
10
Iteration
Iteration
Figure B.5: Epistemic Uncertainty Comparison. EBTs learn express epistemic uncertainty (uncertainty related lack knowledge) on unseen data. Particularly, the sequence on left, which a text sequence likely seen during training, has consistently lower energy (uncertainty) for tokens than the sequence on the right, which is random text sequence not from the training distribution. This demonstrates that EBTs learn to "know what they don't know."
C ADDITIONAL EBT DETAILS
C.I FORMALIZING THINKING
Due to the recent surge of interest in scaling the performance of models during inference/test time, there are several common terms used to refer these ideas. These include scaling the thinking capabilities of models Jaech et al. (2024), inference time scaling Ma et al. (2025), inference time compute Manvi et al. (2024). Therefore, to reduce confusion stemming from a wide variety of terminol- ogy and unite community, in this work we broadly define these concepts as System Two Thinking or more simply Thinking, We formalize improvements made by Thinking as the following:
(2024), and test time compute Jaech et al. (2024); Snell
Definition C.1 (System 2 Thinking). Given a problem with data x. a model 8, and additional computational resources in the form of function evaluations F greater than the minimum number
24

---

Published as a conference paper at ICLR 2026
Performance as Forward Passes Increases
DIT
PSNR (1ligher is Better 1)
22
EBT
20
18
16
10 Number of Forward Passes (log scale)
100
Figure B.6: Image Denoising Thinking Scalability, A comparison between EBT and DiT on image denoising given a different number of forward passes. EBTs require only 1% of the forward passes used by DiT to achieve comparable better PSNR. Further, the scaling rate of PSNR improvement given more forward passes is much higher for EBTs than it for DiTs. These results suggest EBTs have superior thinking capabilities than DiTs on OOD data.
of function evaluations to get a valid prediction from the model Fo. System Two Thinking STT(.) quantifies the expected percentage improvement in performance as F increases. Let P(rr,00 0,, F) be the performance on input x when the model 9 uses F function evaluations:
P(x.0.F) P(x, 0, Fo)
STT(x,0,F) Ez
This formalization is compatible with any type of metric (e.g.. Accuracy, Perplexity. FID, etc), and uses psychology-alignes terminology Kahneman (2011). Further, avoiding terms such as "inference" or "test-time" makes the idea of Thinking more compatible with domains where the line between inference and training is blurry, such real-world continual learning, domain adaptation, or actual human term across machine learning representing several different ideas, we intend for thinking to similarly unify many diverse ideas under a common framework. For a greater justification on this perspective, please see Section G
processes Parr et al. (2022). Just as learning has become a flexible
C.2 EHERGY-BASED TRANSFORMER (EBT) THINKING TYPES
All experiments in the paper are conducted with two main variants EBTs, which we call System (S1) and System 2 (S2) EBTs S1-EBTs have hyperparameters specifically optimized for stability and learning convergence, whereas S2-EBT: Ts have hyperparameters optimized for System 2 Thinking capabilities. Many of the pretraining scaling experiments conducted in Section 4 are with S1 models as to reduce the computational resources required for experimentation. To confirm that these results hold for S2 models, we plot the scaling trends of S1 and S2 models side by side; in Figure C.I, we find that S2 models scale at the same or a higher rate than SI models during training, but have higher Y-intercept. This Y-intercept offset does not affect asymptotic scaling behavior (as asymptotically the scaling rate dominates), and hence, scaling trends that hold for SI models hold for S2 models (S2 models may even perform better than S1 models during pretraining asymptotically because of the higher scaling rate). Intuitively, switching from S1 to $2 models allows for compute trade-off between the model's pretraining performance and the model's System 2 Thinking capabilities.
S1 models have the gradient of predictions detached between optimization steps to increase training stability. Conversely, following Du et al. (2024), the S2 models truncate backpropagation and avoid detaching prediction tensors between optimization steps. Additionally, the S2 models have all the energy landscape regularization techniques described in Section 3.3, whereas the S1 models have none. We also find that S2 models require different value for the optimization step size, that the optimization step size not be learned, and to perform a minimum number steps (and not just any value from I to the max number of steps).
25

---

Published as a conference paper at ICLR 2026
EBT-S2 Scales 3.30% Faster Than EBT-S1
80
Validation Perplexity (log scale)
EBT-S1
EBT-S2
70
60
50
1.0
10,0
Non-Embedding Parameters Millions (log scale)
Figure C.I: EBT S1 S2 Scaling Comparison. Scaling rate System (S1) compared to System 2 (S2) models. System 2 models have a higher Y-intercept scale slightly faster than System models. Therefore, as the scaling rate ultimately dominates asymptotic scaling hehavior, and not the Y-intercept, scaling results in the main paper that hold for S1 models generally hold for S2 models.
C.3 ENERGY-BASED TRANSFORMER TOY EXAMPLE
Similar to standard neural network training, the optimization process for EBTs starts from initial random guess. For neural networks, this corresponds to randomly initialized weights; in our EBT formulation, it corresponds to initializing the prediction yo by sampling from Gaussian distribution. In neural networks, we compute loss and differentiate it to obtain a gradient with respect to the network parameters, which we then update reduce the loss.
In contrast, our method optimizes the prediction directly. Given a current prediction yt. we input Ve (and the conditioning I, if applicable) into the EBM to obtain scalar energy (2, (x).) We then compute the gradient of this energy with respect to the prediction, Eo,E, Wr), and update the prediction via gradient descent to decrease the energy. This iterative procedure for updating the prediction mirrors the way neural network parameters are updated during training to minimize the loss, but operates in prediction space instead of parameter space.
C.4 AVOIDING THE CURSE OF DIMENSIONALITY
This approach avoids the curse of dimensionality inherent in traditional contrastive Energy- Model (EBM) training by reframing learning as optimization problem (Algorithm 1). Contrastive methods must explicitly push up the energy of an exponentially increasing number of negative samples in high-dimensional space. Our iiriiiiaationbbseed approach, instead, implicitly shapes the energy landscape by training the network to successfully perform gradient descent from a random initial prediction to the ground truth solution. This implicitly regularizes the landscape to only have local energy minima at data points. Because regularized methods are known to suffer less from the curse of dimensionality, this successfully avoids such pitfall LeCun (2022). This aligns with more formal results on Ising models from theoretical computer science, where learning the energy function (Hamiltonian) is of much simpler complexity-wiss than accurately sampling learn the full resulting probability distribution Ravikumar et al. (2010).
C.5 COMPUTING SECOND-ORDER GRADIENTS
Because training EBTs involves supervising the gradient of the energy function, second-order gradients must be used. Using PyTorch Paszke et al. (2019), these second-order gradients are computed using torch. autograd. grad with create "reate_graph set to True. This function keeps the graph from backpropagation memory, such that the loss can be computed through this derivative.
26

---

Published as a conference paper at ICLR 2026
Because this involves derivative of derivative, it requires second-order derivatives. Unfortunately, formulating the entire Hessian naively is expensive. Thus, deep learning frameworks such as Py Torch often automatically resort to computing Hessian Vector Products (HVPs) when possible. These HVPs are linear in computational time and memory, enabling scalable second-order computation whose cost comparable normal backward pass Dagréou al. (2024); Pearlmutter (1994).
C.6 AUTOREGRESSIVE CAUSAL ENERGY- BASED TRANSFORMERS IMPLEMENTATION
In this section, we detail the implementation of decoder-on oorodeer-ollyyao autoregressive EBTs, and the complexity that arises due to the way EBMs work. Particularly, the implementation causal-attention decoder- only autoregressive transformers poses challenge due to making predictions in the output space rather than in the input space. To demonstrate why this poses a challenge, for the decoder-only autoregressive transformer consider the case of the attention scores matrix after the causal mask has been applied:
0221.21
0
0l22,02
Scores
Ot 22
where Oz,,. represents the attention score (probability mass) from state to state 2j. Now, in the case of an EBM, where predictions future states are made in the input space, the intended n x n 1 attention scores matrix would look like the following:
O/22,z1
Ol2z,22
Scores
O' An+1
This is challenging to compute because each & along the superdiagonal is unique for its row. Consequently, this matrix cannot be computed with a matrix multiplication (softmax regular attention, every value the superdiagonal is prediction and not past state.
[QKT)
vdd as in
Additionally, traditional transformer, if the context length is n, the size the passed in tensor will be bs d where bs is the batch size and d is the embedding dimension. However, since EBMs make predictions in the imput space, the input tensor needs to be different to allow for inputting future predictions. Therefore, for context length of n, we define the first n elements as 2,. or the original sequence representations, and the final n elements as Zp. or the predicted sequence representations.
The values for 201 or the given states, are computed the same as in the original transformer Vaswani al. (2017), as the attention scores of the original states do not depend on the predicted states. This can be formalized as the following:
Cak? Vdg.
Attention (Q., Ko, Vo)= = softmax
Vo.
Where Qo. Ko, and Vo the Query, Key, and Value matrices of the past states 2 In every block of the transformer, the representations all past states are updated in this manner, independent of the representations of the predicted future states.
For the computation of representations over future states, 3 matrices are also computed, but for the representations of predicted future states rather than past states. We call these Qp. K,, and Vp. First, we compute the self-attention scores of all future representations all past representations:
OpK7 Vik
unnormalized
Note, however, that the eeff-attention scores of each predicted future state with itself is not calculated- due to the key matrix being from the original states. Therefore, the superdiagonal needs to be replaced with the self-attention scores of each predicted future state with itself to achieve the attention score matrix shown in Equation 3.
27

---

Published as a conference paper at ICLR 2026
EBT involves two separate tensors one for past states and one for predicted future states. We denote these as 21 and 2f where & are known past states and & are predicted future states. The intended attention scores matrix is the following:
scores
a2
To compute Equation 3, first need append column to the right side of the unnormalized_scoree peoredd_soree__ppp_ matrix, as the size of the matrix is currently n - 1 X n - 1, but we need to have representations within context. After doing this, we first mask out the superdiagonal. to ensure that the probabilities in the score matrix only correspond to the values of the predicted future states with itself. This masking operation is done through elementwise multiplication of a matrix with I's everywhere except the superdiagonal, which has O's. Then, we compute the self-attention scores of each predicted future state with itself, using the following equation:
Z Z_p_self attention = sum p * Kp),
(6)
where the x indicates the Hadamard product and the sum across the fourth, attention head, dimension. Using superdiagonal mask again, set the diagonal of the unnormalized, scores to these values. Now, after applying the softmax:
scores_p softmax (unnormalized_scores..
have the intended scores matrix shown in Equation 3. However, more barrier towards finally extracting all updated 2/11 representations is the fact that we cannot simply multiply this resulting scores matrix by the values matrix, as each element of the superdiagonal corresponds to different predicted next future state. Thus, using similar techniques to before, we first clone and then extract the superdiagonal from this scores matrix using diagonal mask.
After extracting the superdiagonal, we can multiply the resulting scores matrix by the V, matrix to get all of the representations summed together of each predicted future state with all past states. This is represented as the following matrix multiplication:
2p1 scores V..
(8)
As we also need to add the representation of each predicted future state weighted with its attention score (what was extracted on the superdiagonal), we perform another Hadamard product of the V,, matrix with the cloned superdiagonal to get these values, and then add these element wise to the 2p1 representations. Now, we have computed the intended representations involving the scores matrix shown in Equation 3. Thus, 21 and 87 are updated using = and ep) respectively, by multiplying these tensors by the output weight matrix W.
C.7 AUTOREGRESSIVE ENERGY-BASED TRANSFORMERS SIMPLIFIED IMPLEMENTATION
A more simplified implementation involves the entire attention matrices and generalized causal mask as described in Deng et al. (2024a). However, because this implementation involves matrix multiplication with 2 times the sequence length, this results in 4 times the number of FLOPs normal attention, which is around double the number of FLOPs of our more efficient implementation.
D EXPERIMENTATION DETAILS
Tables D.3 and D.4 specify general model information and hyperparameters. We utilized the Llama 2 transformer implementation Touvron et al. (2023) for the Transformer++ and used this implementa- tion the backbone upon which built EBT. We seed all libraries using Py Torch Lightning Falcon (2019) for all experiments with a seed of 33. For the Diffusion Transformer, we use the implementa- tion from Peebles & Xie (2023) for the bidirectional EBT we build upon this implementation.
28

---

Published as a conference paper at ICLR 2026
D.1 AUTOREGRESSIV LANGUAGE MODELING EXPERIMENTAL DETAILS
For all scaling experiments, copy the popular Mamba paper pretraining settings (for the model configuration, not necessarily for data and learning rate configuration), shown in Table D.2. Because are compute-limited, we also include two extra model sizes in this table, extra extra small (xxs) and extra small (xs).
We manually created training and validation split of 66 million and 33 thousand samples for the red pajama dataset, respectively. Additionally, we focus on reporting perplexity as our relatively small models trained from scratch, when compared to current foundation models, do not achieve high accuracies on many of the benchmarks used. Furthermore, perplexity often functions as a more linear metric than accuracy Gu & Dao (2023); Schaeffer et al. (2023), enabling a more comparable analysis of downstream performance as we scale compute during inference.
D.1.1 LEARNING SCALABILITY EXPERIMENTAL. DETAILS
Conducting thorough scaling experiments is very challenging a recent survey on "scaling laws" Li et al. (2025a) showed just how fragile many of these "scaling laws" are to hyperparameters, data, etc. and how changing these parameters slightly can lead to different conclusions. Being bottlenecked by a limited set of computing resources further exacerbates this issue. Therefore, we sought out to conduct controlled experiments that revealed the most information possible regarding the scaling of EBTs compared to different models.
Most existing works studying scaling by changing several factors at the exact same time, including depth (number of transformer blocks) Kaplan et al. (2020), width (embedding dimension) Kaplan al. (2020), possibly batch size Chen et al. (2024); Hu et al. (2023), and the amount of data Kaplan et al. (2020). Therefore, to be more comprehensive determining when EBTs scale differently than the Transformer++, we decided conduct normal scaling experiments over all of these factors the exact same time (as is standard), as well as ablating over over.just changing just one of these parameters at time. Notably, conducting experiments in this manner allows for controlling a single independent variable at time (i.e., just changing the number of Transformer Blocks), which allows for stronger conclusions regarding what aspects of model scaling different models perform better over (i.e., EBTs scale better then the Transformer++ when increasing the number of Transformer Blocks). Scaling all factors at once does not allow for such insight, as there are many independent variables.
For our parameter and FLOP scaling law experiments in Figures 5 and b, we follow the popular and reliable Chinchilla scaling laws Hoffmann et al. (2022), where we scale the number of parameters and the data models are trained on proportionally. We use factor of 20x the amount of tokens as parameters." & We do this for XXS, XS, Small, Medium, and Large models in Table D.2. All model sizes use the same batch size of 64, and context length of 256 tokens, where scale the number of training steps to increase data as parameter scale increases. We use three random seeds (33 - 355) for the XXS, XS, and Small models in order to improve robustness and reliability of our scaling laws while also reducing variance.
In addition to plotting the results in Figures 5 and b, in order to increase transparency and repro- ducibility, we also report raw losses in Table D.1. This allows for fitting the scaling law exponents, which allows for computing the scaling rate difference in the two models used (EBTs and the Trans- former++). Following Chinchilla Hoffmann et al. (2022), we model the scaling behavior as a power law CC) = GC a + F, where the loss, C the compute budget (FI.OPs), and E is the irreducible entropy of natural text. Assuming F is constant across models, we approximate the scaling rate by fitting a linear regression in the log-log regime:
log(L- E' -a log(C) + log(33
Here, a represents the scaling exponent. By comparing the exponents for EBTs and Transformers, we determine which architecture improves loss more efficiently per unit of compute. Based on the fit linear regression, we also compute an R² coefficient to determine the reliability of these scaling laws for FLOPs, where we get value of 99.69% for the Transformer++ and a value of 98.40% for
Foor these experiments transitioned to using the FineWeh dataset Penedo et al. (2024), which is of high quality and has recently hecome popular.
29

---

Published as a conference paper at ICLR 2026
Table D.1: Pretraining Scaling Law Loss Across Model Sizes and Seeds. Cross-entrope loss for the Transformer++ EBT models five model sizes (XXS to Large) and three random seeds (33, 34, 35). Medium and Large models were trained with a single seed due to increased compute requirements.
Size
Model
Seed 33
Seed 34
Seed 35
Non-Embedding Params.
FLOPs
XXS
Transfomer++
4.7786
4.7638
4.7619
6.18
4.56B15
XXS
EBT
5.1210
5.1122
5,0392
6.18
3.04E16
XS
Transforner++
4.3728
4.3594
4.3546
12.4
1.82E16
XS
EBT
4.5119
4.5118
4,4906
12.4
1.22E17
Small
Transformer++
3.7333
3.7354
3,7395
48.8
2.87E17
Small
EBT
3,8262
3.8321
3,8189
48.8
1.92E18
Medium
Transformer++
3,2950
176
3.63E18
Medium
EBT
3.3922
176
2.42E19
Large
Transformer++
3.0443
395.5 395.5
1.87E19
Large
EBT
3.1079
1.24E20
Table D.2: Model sizes and hyperparameters for scaling experiments. For most model sizes follow Gu & Dao (2023).
Size
Non-Embeddiig Params
layers
embed. dim
# heads
XXS
6.18M
384
XS
12.4M
12
384
6
Small
48.8M
12
768
12
Medium
176M
24
1024
16
Large
396M
24
1536
16
XL.
708M
24
2048
32
EBTs for FLOPS, both of which are within well-accepted ranges for Chinchilla scaling laws Inbar & Sernau (2024),
D.1.2 THINKING SCALABILITY EXPERIMENTAL DETAILS
We train xxs models with the same setup as above, with the exception that models are trained with batch size of 128 for IM training steps. Increasing the data scale enables us to better understand how thinking scales during pretraining. It's worth noting that since we are training small language models, they could not benefit from modern techniques such Chain of Thought (CoT) in improving performance.
D.2 AUTOREGRESSIVE VIDEO EXPERIMENTAL DETAILS
We use the same model parameter scaling, shown Table D.2, as NLP experiments. We also used batch size of 256 for all models, as we found that it did not significantly affect the scaling performance due to models training for many epochs. We also processed videos with 0.25 seconds between frames. For the Transformer++ baseline, experiments. For EBT, we found that it was necessary to use a lower learning rate by factor of We use the standard SSV2 train and validation split for experiments. Other hyperparameters are shown in Table D.3 and Table D.4.
the same learning rates the NLP
D.3 BIDIRECTIONAL IMAGE DENOISING EXPERIMENTAL DETAILS
We use the COCO 2014 dataset Lin al. (2014); Tawfik with 128 by 128 images, its train/validation split, patch size of 16, and the Diffusion Transformer implementation from Peebles & Xie (2023).
Note that due not including all significant figures when reporting the losses and FLOOPs, the R² coefficient from this table may he slightly different from reported results.
30

---

Published as a conference paper at ICLR 2026
Table D.3: Hyperparameters for Transformer++.
Hyperparameter
CV
NLP
Optimizer
Adam
Optimizer Momentum
31, 82 = 0.9, 0.999
LR Schedule
Linear warm up cosine decay
Warmup steps
le4
Minimum LR Scale
10
Gradient Clip Value
1
Weight Decay
0.01
Context Length
16 SD-XL VAE Rombach et al. (2022); Stability 224x224
256
Encoder
Image Dimension
Tokenizer
EleutherAA/ggttneox-2222 Black al. (2022a)
Vocab Size
50277
All models were trained using the large model size described in Table D.2, with learning rate of le for 100, 000 steps. For the DiT baseline, we used the same hyperparameters from Peebles & Xie (2023), changing only the batch size to 128 from 256. We based our bidirectional EBT implementation on the code from this repository. We experimented with several different diffusion inference strategies to ensure fair comparison, including DDPM, DDIM, increasing the number of diffusion steps at inference, as well as recursing the diffusion model on its own denoised output. Ultimately, found that the combination of DDIM recursed
its own output performed best, hence we used this as the baseline all experiments. We used the default denoising schedule from the DiT codebase Peebles & Xie (2023). As the noise level 3 was to 0.1 during training, and the default number of diffusion denoising steps was 1, 000, the number of denoising steps the diffusion model was trained
100.
To make the denoising experiments compatible with diffusion models, we deviate from the original noising schemes performed in the denoising works mentioned in the main paper Du et (2022), and use a scheme based on the noising schedule from diffusion models. Specifically, we follow Peebles & Xie (2023), and use linear variance schedule ranging from 10 4 to 10-2, To control the noise level, we use hyperparameter denoted 33 representing the percentage of the diffusion schedule to noise samples; was set to 0.1 during training.
For both DiTs and EBTs, we found that models performed best on OOD noise levels when denoising their own outputs twice, that applying the model to denoise the same image three times recursively, This is how we are able to get the results in Figure B.6 demonstrating the performance for 300 forward passes from DiTs and forward passes from EBTs. We found that for image denoising, was not necessary to train EBTs with the S2 hyperparameters for System 2 Capabilities to emerge, although its possible these would further improve performance. Additionally, for image classification, for both models, take the average of all the final patch tokens, and for DiT's feed in T = 0.
D.4 COMPUTATIONAL RESOURCES
All experiments were conducted on either Nvidia A100s or GH200s, with the largest scale experiment requiring approximately x 1300 A100 GPU Hours. The runtime for each experiment was dependent on the model sizes used as well as the amount of data trained on.
D.5 THEORETICAL FLOP AND MEMORY CALCULATIONS
We adopt the standard estimate of 6N FLOPs per token for the AR Transformer++ Casson (2023), where N denotes the number of non-embedding parameters. For AR EBTs, however, the per-token cost varies with the number of training optimization steps and chosen hyperparameters.
To derive the FLOPs for EBT training, follow Dagréou al. (2024) for the Hessian-vector product (HVP), which EBTs require to backpropagate through a first-order derivative. Since an HVP has the same theoretical complexity a gradient computation Pearlmutter (1994), we express the per-step
31

---

Published as a conference paper at ICLR 2026
Table D.4: Hyperparameters for EBT experiments.
Hyperparameter
CV
NLP
Optimizer
Adam
Optimizer Momentum
B1, 82 0.9, 0.999
LR Schedule
Linear warm up cosine decay le4
Warmup Steps
Minimum LR Scale
10
Gradient Clip Value
Weight Decay
0.01
Context Length
16
256
Encoder
SD-XL VAE Rombach et al. (2022); Stability 224x224
Image Dimension
Tokenizer
EleutherAU/gpt-noo-22200 Black et al. (2022a)
Vocab Size
50277
Optimization Steps
2
2
Optimization Step Size
30, 000
500 1, 500
Optimization Step Size LR Multiplier
90, 000
Learnable Optimization Step Size
FLOPs
= BB
Based on Casson (2023), the forward and backward passes require approximately 2.V and 4N FLOPs token, respectively, where N denotes the count of non-embedding parameters in the Transformer. In the autoregressive EBT implementation, the effective sequence length becomes twice that of the original Transformer (formally 2S 2 for original sequence length S). Owing to the efficient scheme of Section C.6, this doubling of sequence length translates roughly into a two-fold increase in FLOPs, rather than four-fold increase in FLOPs. Hence, each second-order optimization step demands roughly
(F + B + B) 2 (2N + 4N 4N) 2 10N 2,
making it = 3.33x more expensive than a standard feed-forward Transformer step.
The overall FLOP count also depends on one's choice of hyperparameters. For S2 models, where the loss is evaluated at every iteration without gradient truncation, the total FLOPss simply multiply by the number steps. Therefore, for our pretraining experiments using two steps, we get that EBTs used 6.66x the FI.OPs of a comparable Transformer++ training. In contrast, for S1 models, random number of optimization steps are used, the gradient is truncated, the loss is only calculated at the last step following Du et al. (2022), and a Replay Buffer is used. Therefore, the FLOP count varies and can both decrease (as truncating uses less FLOPs for earlier steps) as well as increase (as using more steps and a replay buffer both use more FLOPs). These numbers also vary during inference, where the full EBT implementation parallelizing all predictions at once is not necessary.
Given the scarcity of published methods for computing higher-order derivative FLOPs and inability to leverage existing libraries for Hessian-vector products, these estimates remain approximate. We welcome corrections or additional insights from readers familiar with FLOP calculations for second-order methods.
The theoretical memory utilization of EBTs compared normal feed-forward Transformers follows the same trend and calculations, where optimal HVP computation involves linear time and memory the same as standard backpropagation Pearlmutter (1994).
D.6 EMPIRICAL HARDWARE BENCHMARKING
In addition to discussing theoretical FLOPs, to help understand performance more practically with modern hardware and software, we discuss empirical hardware performance and benchmarks for EBTs versus the more standard Transformer++. We conduct experiments measuring throughput, memory usage, and latency of EBTs and Transformer++ Autoregressive Language Models. We measure performance three times for each metric, and report average. All experiments with
32

---

Published as a conference paper at ICLR 2026
Table D.5: Training Efficiency of XXS Models. Throughput (iterations/s after 5 minutes), memory usage (GB.
batch size 32), and latency of the first training step (ms) over three runs and their average.
Model
Trial 1
Trial 2
Trial 3
Avg
Throughput (it/s, higher is better)
EBT MCMC (xxs)
4.35
4.28
4.54
4.39
EBT 2 MCMC (xxs)
2.18
2.16
2.30
2.21
Base (xxs)
16.74
16.75
18.16
17.2
Memory (GB with batch size 32, lower is better)
EBT MCMC (xxs)
39.6
39.6
35.7
38.3
EBT MCMC (xxs)
62.6
62.6
59.7
61.6
Base (XXs)
22.4
22.4
19.9
21.6
Lateney for first training step (ms, lower is better)
EBT MCMC (xxs)
1020
1020
1020
1020
EBT MCMC (xxs)
1860
1840
1900
1870
Base (XXs)
587
592
552
577
XXS models using a batch size of 32, and memory/throughput being reported after minutes of training. We measure EBTs with both 1 optimization step and 2 optimization steps to determine the impact of the number of steps on performance The results are shown in Table D.5, where the empirical performance of EBTs compared Transformers very roughly matches theoretical FLOPS and memory. It's possible that a large part of the mismatch stems from PyTorch implementing HVPs suboptimally, with Jax implementing them much more efficiently Dagréou et al. (2024). Future algorithmic progress of EBTs could drive further Py PyTorch software development with better support for efficient HVPs.
D.7
EXPERIMENTS
All models are trained for up to 100k training steps, where we often start to observe overfitting within the first 10k training steps. We use learning rate of 1e for all models, a batch size of 64, and weight decay of 0.01. For RNNs we experiment with the number of iterations being in the range of 20, and for fair comparison, we experiment with EBTs having the number of MCMC steps in the range of 2 to 7 (which is actually less compute per training step than the RNNs). Both RNNs and EBTs were only trained at the XXS scale. Because feed-forward Transformers have no notion of inference-time computation, we also ablated training larger Transformer models to favor them (small, instead of just XXS), where performance improved from 0.0% to 0.03% moving from XXS to Small. Small transformers use more FLOPs per forward pass than the XXS EBTs and RNNs that achieved much better test-set performance, revealing how even when mmtthedd Transformers perform worse than EBTs/RNNs due to overfitting data rather than learning generalizable reasoning.
E RELATED WORKS
E. TRADITIONAL TRANSFORMERS
The Transformer architecture Vaswani et al. (2017) has become ubiquitous across various do- mains Latif et al. (2023); Oquab et al. (2023); Radford et al. (2019); Touvron et al. (2023). The most commonly used transformer variant of today makes predictions directly in the output space with a single forward pass, demonstrated in Figure 1a. Because these models have finite depth and width, and make predictions putation to each prediction being made. Furthermore, they cannot model uncertainty in continuous state spaces in the same way they can in discrete state spaces because the normalization process for continuous state spaces is not well-defined as it is for discrete spaces using softmax Dawid LeCun (2024). Rather, training these models to express uncertainty relies on tricks such as Vector Quantization Van Den Oord al. (2017) pseudo losses/objectives (e.g., ELBQ Kingma
single forward pass, they are unable to dynamically allocate more com-
33

---

Published as a conference paper at ICLR 2026
et al. (2013)). Finally, because these models are not trained to explicitly verify samples, improving inference-ting tineieeee-meeeee performance at a per-prediction level often requires external models Lightman et al. (2023).
E.2 RNNS
Recently several RNN variants (Figure 1b) have emerged alleviate memory bottlenecks and achieve faster inference Gu & Dao (2023); Peng et al. (2023). These approaches have scaled similarly to Transformers in autoregressive sequence modeling and achieve better memory efficiency and reduced latency. However, traditional RNNs updating internal state based only on new information/data Gu & Dao (2023); Peng et al. (2023) are not capable of allocating additional computation during inference, and thus suffer from the same flaws as traditional transformers in achieving human-like System 2 Thinking.
To resolve these issues, people have equipped RNNs with the ability to allocate computation dynami- cally, with architectures such as the Universal Transformer Dehghani et al. (2018), Recently, this type of RNN has also been applied to LLMs Geiping et al. (2025); Saunshi et al. (2025), allowing LLMs to reason using additional computation in a continuous latent space through the depth of an unrolled RNN. Parallel work has developed RNNs for algorithmic reasoning
Sericceur-Maattnge (2025); Wang et al. (2025). However, like Diffusion models, these models learn to amortize gradient prediction of the energy function Geiping et al. (2025), rather than learning to explicitly verify predic- tions, meaning they cannot model uncertainty or explicitly verify predictions. Consequently, EBMs generalize these RNN-based architectures by offering explicit prediction verification capabilities Ma et al. (2025). Further discussion on this relationship is provided in Section E.
E,3 DYNAMIC COMPUTATION (THINKING) WITH LLMS
The ability to leverage dynamic amount of computation in LLMs has been emulated using chain-of- thought prompting Wei et al. (2022) and continuous latent space reasoning Hao et al. (2024), While these approaches can improve performance, they don't 't seamlessly transfer to continuous modalities, and LLM chain-of-the of-thught has been shown to be unreliable for reasoning Agarwal et al. (2024); Lin et al. (2025); Turpin et al. (2023). More recently, models have been explicitly trained to perform reasoning using Reinforcement Learning Anthropic (2025); Guo et al. (2025); Jaech et al. (2024); XAI (2025). These approaches allow LLMS simulate additional computational depth based on the number of tokens decoded before making prediction, and as result, significantly improve performance Guo al. (2025); Jaech et al. (2024). The main limitations of these approaches are that they currently apply only to discrete domains (i.e., LLMs), are effective on narrow set of problems that are easily verifiable (e.g., math and coding), and require additional supervision, typically in the form of reward signals, making them incompatible with purely unsupervised pretraining Guo et al. (2025).
E.4 DYNAMIC COMPUTATION (THINKING) WITH DIFFUSION
The most common instance of a model architecture specifically created to leverage dynamic computa- tion is diffusion models (Figure Ic), where using multiple forward passes to generate a prediction is core aspect of both training and inference Höppe et al. (2022); Rombach et al. (2022). Although dif- fusion models implicitly define likelihood through the reverse process Ho et al. (2020), which could theoretically be used to verify predictions, in practice an external verifier is necessary to improve performance inference time beyond increasing denoising steps Liu et al. (2025): Ma al. (2025); Singhal et al. (2025). This requirement limits the generalizability and scalability of diffusion models as an approach for System Thinking, as they do not have two of the cognitive facets described: the ability to model uncertainty in continuous state spaces or explicitly verify predictions without additional models Liu et al. (2025); Ma et al. (2025); Singhal et al. (2025). Furthermore, diffusion models rely on fixed denoising schedule, which restricts their ability to adaptively halt or extend computation unlike EBMs. Additionally. diffusion models can be seen as predicting the gradient of the data density/energy function Du al. (2023), and therefore that EBMs are generalization of diffusion models that learn to explicitly verify predictions. More on this connection is in Section E.6, and side-by-side comparison of diffusion models and EBMs is in Figure E.I.
34

---

Published as a conference paper at ICLR 2026
Diffusion Models
Energy Based Models
Learned Denoising Path
Learned Energy Landscape
1 1-44 at
ft-
22.,2.2 () tote.e-N(0.)
9++1 - (Eo,,
vac
/ - 8,
80 Initial
Initial
Target y
80
Supervision
y
Supervision
Target
(a) Diffusion Model
(b) Energy- -Baed Model (EBM)
Figure E.1: EBM and Diffusion Comparison. Diffusion models receive supervision at each step of the denoising process (e.g.. for one thousand steps), whereas EBMs only receive supervision at the end the optimization process. This training procedure allows EBMs to learn an entire Energy Landscape over predictions, associating a scalar energy for every prediction according to its likelihood. Learning landscapes in this manner reduce "error" accumulation throughout the denoising process Du et al. (2024) and makes EBMs more flexible by allowing unnormalized likelihood estimation at each step of the denoising process. Additionally, diffusion models update predictions by predicting the noise at cach timestep, meaning they must follow a set denoising schedulc. On the other hand, EBMs update predictions by performing gradient descent with respect to the energy scalar, allowing for flexible inference where this optimization process can be performed for any number of steps. here refers some condition (e.g., class text) whereas y is the generated prediction.
E.5 ENERGY- -ASED MODELS (EBMs)
The perspective of energy minimization as thinking/reasoning has been known for some time LeCun et al. (2006). Therefore, the most similar approaches to EBTs also train EBMs to do reasoning/think ing Du et al. (2022; 2024). While these works achieved impressive generalization results, they only focus on small-scale problems, and did not scale EBMs to high-dimensional real-world problems such as language or video. Additionally, these works did not perform an in-depth analysis on the types of System 2 Thinking that emerge with EBMs, more complex inference time procedures beyond just increasing the number of gradient descent steps, approaches towards improving EBM scalability, and required techniques for enhancing System 2 Thinking in EBMs.
E.6 ENERGY-BASED MODELS (EBMS) AS A GENERALIZATION OF DIFFUSION MODELS AND RECURRENT DEPTH MODELS
Both diffusion models Du al. (2023) and RNNs Saunshi et al. (2025) can seen as predicting the score, the gradient of the energy function/data density, Eg(),, where diffusion models do this with an additional time condition Liu et al. (2022). Thus, the largest benefit of explicit EBMs over these approaches, which can be seen as implicit EBMs (due only implicitly defining an energy function), is that using an explicit EBM allows for explicit verifieation/ikelittodd modeling. We show improve predictions, whereas with diffusion models
that this enables the use of self-verification and RNNs an additional verifier model is necessary to achieve this capability Ma et al. (2025).
It's also worth noting that RNN, diffusion models, and EBMs need not be incompatible with one another. For example, Du et al. (2024) combines EBMs and diffusion to reason over challenging problems. This can increase stability of the learned energy landscape by adding explicit score supervision.
E.7 EBM AND DIFFUSION MODEL COMPARISON
Because of the similarity of diffusion models and EBMs, we present a side-by-side comparison of the training and inference approach for both in Figure E. where the primary difference lie in the supervision they receive during training and the update rule for predictions. Additionally, we provide more information to compare these approaches.
35

---

Published as a conference paper at ICLR 2026
Under the assumption that the energy landscape is well formed and that optimization is well behaved, EBMs offer several distinct advantages over diffusion models. In EBMs, the energy function is trained to represent meaningful landscape where the energy value of a sample directly corresponds to its relative unnormalized likelihood. Consequently, two samples can be directly compared to determine which is more likely, in single forward pass. On the other hand, diffusion models require running samples through the entire reverse diffusion process to get likelihoods, which often requires hundreds to thousands of steps, and rely on likelihood approximations such as ELBOs or numerical solvers for SDEs/ODEs. In practice, these result in incomparable likelihoods, as ELBOs only give likelihood lower bounds and numerical solvers result in high approximation error Ma et al. (2025).
Furthermore, the learning of an energy landscape means that any approximation errors at each individual step of the Markov Chain (optimization process) do not result in cumulative error, as the minimum of the energy landscape can still he reached. This differs from diffusion models, where any approximation error at each step will result in increasing accumulated error across the entire Markov Chain Du et al. (2024) (demonstrated in Figure E.
Lastly, EBMs, giving an unnormalized likelihood estimate at each step, are in practice much more flexible for generation than diffusion models. While diffusion models require running the entire reverse diffusion process with a specific denoising schedule to generate a sample, EBMs can be trained to directly predict the next sample in single step, and giving an unnormalized likelihood at each step can indicate how likely they think this sample is. This better approximates human-like System 2 thinking, where humans naturally evaluate the strength of current predictions, and on the basis of knowing how good their predictions are, decide to dynamically allocate more or less computational resources.
E.8 ADDITIONAL ENERGY-BASEL MODELS RELATED WORKS
One contribution of this work was the design of a custom architecture for EBM's 's called the Energy- Based Transformer (EBT). Roughly similar is the work of the Energy Transformer Hoover et al. (2024). Despite strong similarity in the names of these architectures, however, they are very different- with the primary similarity in architectures being the usage of attention mechanisms as well as a global energy function. The existing work integrated ideas from Modern Hopfield Networks, including RNNs, whereas in our work the architecture is non-recurrent and does not use associative memories. Additionally, EBTs differs with its focus on System 2 Thinking, which this previous work did not experiment with.
Other somewhat similar approaches to EBTs involve autoregressive Energy-Based Models, including E-ARM Wang et al. (2022), EBR Bhattacharyya et al. (2020), and Residual EBMs Bakhtin et al. (2021). E-ARM involves adding an objective the learning process to turn traditional autoregressive model into an EBM, and as such does not achieve two of the cognitive facets discussed. EBR and Residual EBMs involve the training of an EBM on top of an already pretrained autoregressive language model. Both works, however, leverage a contrastive objective, which suffers from the curse of dimensionality.
The optimization procedure used to train EBTs can be seen as form of denoising score match- ing Vincent (2011); Wang et al. (2023). Particularly, predictions being initialized at some Gaussian, and then being optimized using the gradient of the energy function can be seen as training the EBM to denoise by learning the score of the data. However, we find the optimization perspective is more intuitive, and this denoising score matching perspective is more similar to the diffusion model training procedure than it is EBMs, involving multiple levels of noise rather than just one.
F ADDITIONAL COGNITIVE FACETS
Facet 3: Modeling Uncertainty in Continuous State Spaces. While thinking longer is important for improving performance, humans also weigh how uncertain they are before committing to a decision. In language, LLMs can simulate this through token-level probabilities Tomani et al. (2024). In the context of continuous state spaces, such as in vision, without the usage of discretization schemes such as Vector Quantization Van Den Oord et al. (2017) or pseudo losses/objectives (such ELBO Kingma et al. (2013)), standard implementations of the most successfully used approaches with Transformers, RNNs, Diffusion models generally do not provide strong or reliable uncertainty estimates Heng
36

---

Published as a conference paper at ICLR 2026
et al. (2024); Nalisnick al. (2018); Sankararaman al. (2022); Serrà et al. (2019). EBMs can naturally model uncertainty without having to model exact likelihoods Dawid & LeCun (2024) by modeling the relative unnormalized likelihoods of predictions, demonstrated in Figure 3. As the real world often contains many inherently unpredictable elements, for instance, when a pedestrian might emerge from behind parked vehicle, the ability express uncertainty in predictions is essential to being cautious, and is natural capability of humans Peters et al. (2017); Sarinopoulos al. (2010); Vilares et al. (2012).
Facet 4: Compositional Reasoning and Systematicity. Humans routinely solve novel tasks by recombining familiar primitives (e.g.. verbs with new arguments or visual parts into unseen objects). hallmark of compositional generalization well-documented in neuroscience Friederici & Weissenborn (2007). In contrast, state-of-the-art Transformers and diffusion models often fall short when evaluated on compositional generalization Huang et al. (2023); Kobayashi et al. (2024). Emergy-Based Models (EBMs) seamlessly address these limitations: energies for individual factors are composable in several different manners Du et al. (2023), enabling zero-shot generation of novel combinations without retraining, where gradient-based sampling provides an intrinsic mechanism to verify and iteratively correct compositions Du et al. (2023). Thus, EBMs offer promising path toward human-like like systematicity that remains elusive for existing approaches.
G COUNTERARGUMENT
G.1 1 SYSTEM 2 THINKING
In this paper, strong claims were made regarding the capabilities of current models and their ability to perform System 2 Thinking. However, there are common counterarguments to our claims, which we address here, in hopes of clarifying why we believe this is not currently possible.
G.1.1 SYSTEM 2 THINKING AND INFERENCE-TING COMPUTE TERM USAGE
Whether the computational effort spent at inference time fully captures what psychologists term System 2 Thinking is still actively debated. In Section C.1 we outline three reasons for preferring the broader label System Thinking: (i) it naturally extends to settings such as continual learning where terms such "inference-tine time compute" becomes ambiguous, (ii) it connects our discussion to a substantial body of cognitive-sciencr work, and (iii) it offers a conceptually straightforward entry point for readers beyond the machine-learnin community.
is widely acknowledged, and we agree with the idea, that human System Thinking encompasses far greater depth and complexity than the specific approaches explored in this paper, such as "Thinking Longer" and "Self- solif-eerficatioo.. We wish to emphasize that our work does not claim current models replicate the full spectrum of human System 2 Thinking. Rather, we view the methods presented here as foundational steps toward that more ambitious long-term goal.
We propose that "System Thinking" offers useful umbrella term that can effectively encompass and generalize other existing terminologies, including "inference-time compute," "test-time compute." " or "reasoning." A parallel can be drawn with the term "learning" in our field. "Learning" itself has evolved to describe wide array of processes, some of which, such k-Nearest Neighbors (KNNs) or the specific mechanisms of weight matrix updates in Artificial Neural Networks (ANNs), represent distinct facets rather than the entirety of human-like learning, meaning these approaches may not encompass the complexity of true human-like learning. However, despite this breadth and these simplicities, "learning" has become cornerstone term, upon which our entire field of machine learning is named upon.
In similar vein, while the "thinking" exhibited by the models discussed in this paper may not yet capture the full nuance and intricacy of human cogmition, we believe the term "System 2 Thinking" can still serve valuable role. It offers a generalizing framework for existing vocabulary and can contribute to making complex concepts within the field more accessible and understandable to
We acknowledge that there are approaches to achieve uncertainty with the models discussed, such as Mixture Density Networks Bishop (1994) as well as score-based diffusion models Song al. (2020). However, these approaches have seen less widespread success and scalability than the current dominant approaches.
37

---

Published as a conference paper at ICLR 2026
newcomers experts from other fields. Our intention is contribute constructive and unifying dialogue of intelligent systems.
G.1.2 Is CHAIINOFF THOUGITT (CoT) SUFFICIENT FOR SYSTEM THINKING?
CoT is commonly thought to be sufficient for advanced reasoning to emerge in LLMs. However, in this paper we argue there are several flaws with CoT preventing advanced thinking capablities. First, Chain-of-Thouge
(CoT) involves reasoning over a discrete state space, which limits the granularity of "thoughts." " Second, CoT is not an intrinsic architectural capability but an external procedure applied to token sequences. Ideally, such reasoning should be embedded within the model and learned during training. Third, each token is produced with fixed computational budget, restricting the depth of reasoning per step. In contrast, humans allocate variable effort across steps when reasoning "step by step". Similarly, models should be able to spend a variable amount of computation per token, as enabled by EBTs. This aligns with the intuition behind the saying: "a chain is only as strong its weakest -ack- step in the chain should receive sufficient computation to avoid failure points that result in bad reasoning chains.
H ENERGY- BASED MODELS (EBMs) INTRODUCTION
H.1 1 SIMPLIFIED ENERGY- MODEL (EBM) INTRODUCTION
Feed-forward neural networks generally take the form of: given an predict y (Fig. H.),), Energy- Based Models (EBMs) are family of models that learn the compatibility (unnormalized probability) over all possible combinations of x and (Fig. H.lb), Intuitively, this can be seen as learning to verify the strength of y as prediction with æ as the input. Training models in this manner allows for representing multiple plausible versions of compatible with a given z. The differences between these models is visualized in Fig. H.I.
Formulating models in this manner ultimately brings about two primary questions:
First question: Assuming we still care about ultimately predicting ý how do we use such an EBM to predict d? With feed forward models, generally we can just input x and get the output the model as ÿ. but we can't do this with EBMs?
With EBMs, what happens conceptually similar to diffusion models Rombach et al. (2022), where we (commonly) initialize I as random noise. Then, we input x and ý into the model, and get single scalar energy output (our initial energy output) from the model. Now, because our entire model is differentiable, we can get the gradient from this energy scalar to ý and perform gradient descent along the energy landscape (energy landscapes are surfaces resulting from mapping all possible predictions scalar values, visualized in Figure 3) using this gradient (this is the key)! This process is visualized in Figures 3, 2. This gradient can be seen as the opposite of the noise (e.g.. denoising) and therefore EBMs have strong relations with Diffusion models predicting the noise. EBMs can be seen as generalization of diffusion models, where diffusion models are predicting the gradient of the energy function/scalar (more on this in Section E).
Second question: How do we train an EBM? Generally, models are trained over a dataset of a and pairs, but now there are several different possible y values that can be associated with any given value- so how does that work?
turns out that all the training techniques for EBMs boil down to two main categories: contrastive and regularizing approaches Dawid & LeCun (2024).
Contrastive approaches are more common for EBMs and are easier to rationalize about due to their similarity to GAN discriminators. The idea behind contrastive approaches is to push down on the energy of positive samples (i.e., the true data), and to push up on the energy of negative samples. While these positive samples are easy to rationalize about, they are just the true data, the difficulty of contrastive EBMs is finding negative samples. Several approach exist, such as GANs, which use a generator to amortize negative sample generation, running MCMC (similar to optimization) for some time. However, as discussed in Section 3.1, such approaches don't 't scale well due to the curse of dimensionality.
38

---

Published as a conference paper at ICLR 2026
Therefore, to achieve scalable EBM approach, train EBMs through an optimization procedure (which has strong resemblance to Langevin Dynamics). That is, EBMs are trained to, starting from an initial prediction, optimize predictions to the ground truth solution (shown in Figures 3, 2). This pushes the energy landscape to be smooth with single local minimum at the ground truth solution, thereby regularizing the energy landscape have low energy only the true data. As mentioned in Wang et al. (2023), this can be seen as being similar to denoising score matching Vincent (2011).
Commonly, when people learn about EBMs and the iterative denoising/optimizatinn procedure performed during training, they think of diffusion models, so we include a more in depth comparison between the two Section E.7.
x
x
Feed Forward Model
Energy Based
Model
Energy(x, y)
Feed Forward Model
(h) Eineryy-Based Model
Figure H.I: Feed-Forward and Energy-Bass Model Comparison. Feed-forward models (a), given an input directly try to predict ğ. Instead of just getting the compatibilliy of all possible values of with by outputting scalar energy value for each combination. Low energy corresponds to high probability, high energy low probability. In practice, ý is often initialized as random.
imput, EBMs receive both and as input and learn
H.2 ENERGY-BASED MODEL TYPES
Because EBMs are broad modeling framework, and generalize many existing approaches, aim to provide precise language to distinguish EBM types. We broadly classify the EBMs described throughout this paper as explicit EBMs, meaning they explicitly define energy function over inputs as the entire function being learned. In other words, explicit EBMs directly map all variables (inputs) single scalar energy as the output of the neural network. We define these in contrast to implicit EBMs, or EBMs where the energy function is not the learned model but rather some implicit definition of the learned model. For example, with diffusion models, the energy function implicitly defined by the learned score network (sp (x,
t)) as the following:
veE(x,t) s0(x,t)
E(x, = / so(u, t) du + C(t).
Other notable examples of implicit EBMs include Hopfield Networks Hopfield (1982), RNNs Geiping et al. (2025), and Boltzmann machines.
H.3 ENERGY-BASED MODEL FREQUENTLY ASKED QUESTIONS (FAQ)
What is energy/oompatibiliyy. what does represent, and how is it learned? What energy corresponds to what probability?
Energy is just a learned compatibility score between x and (lower means more likely). The EBMs described in this paper learn it implicitly deseribed in Section 3.1 such that the true data (good pairs) have low energy and bad pairs (non-data) have higher energy. Probabilities follow:
po(x)) 6 (a.9)
Z(0)
(a v) Eo(x.)
39

---

Published as a conference paper at ICLR 2026
the energy E' is essentially the (unnormalized) negative log- likelihood up to additive constant. The term compatibility is just a term used for intuition.
Does training EBMs require a full Hessian calculation?
No-the approach described in the paper only requires Hessian-vector products. That makes training only about a constant 1. 1.66x as expensive as a vanilla feed-forward model given everything else remains constant and you use single step.
Why low energy good (and high energy bad)? Why not just use probability? Low energy is good because of the negative exponential. The reason we don't use prob- abilities avoiding normalized probabilities makes the problem much more tractable in real-world high-ding nigh-dimensional continuous state spaces by removing the focus on explicit normalization via regularizing the partition function. EBMs come from a long line of work in statistical physics.
Is it okay that energies are unnormalized probabilities? Yes, for most real-world applications, you only ever need sample relative likelihood comparison: it's significantly less common to need the exact likelihood of samples. An example of this is reward models, which can be seen as EBMs (just multiplying the reward by -1), where all that really matters is the relative reward for choosing which sample to or which behavior to perform.
Is it fine to not do Maximum Likelihood Training?
Contrary to what your intuition may say, the answer yes! For most real-world distributions, data lies completely concentrated on a very thin manifold with no defined distribution outside of this manifold. Thus, directly doing Maximum Likelihood Estimation (MLE) training would push EBMs to have low energy on the true data manifold and then infinite energy off that manifold (as the probability of such samples 0). We don't 't want this as would make the score (gradient of the energy function) undefined and the energy landscape untraversable- so not doing MLE makes the problem tractable.
I ENERGY-BASED TRANSFORMERS (EBTs) TUTORIAL
1.1 IMPROVING STABILITY AND SCALABILITY
Energy-Basee -Base Models are notorious for instability during training Arbel al. (2020); Du & Mor- datch (2019); Du et al. (2020); Li et al. (2023). Therefore, we experiment with several different hyperparameters to increase the stability and scalability of EBTs and EBMs in general.
1.1.1 OPTIMIZATION STEP SIZE AND STABILITY
We found that the step size for gradient descent updates of predictions (ax) was one of the primary factors affecting the stability of EBTs Thus, for S1 models, make the step size learnable parameter (this is not the case for S2 models). We calculate its learning rate by multiplying the model's learning rate by the step size learning rate multiplier. We found that the values for the step size have a large effect on the magnitude of gradients generated for the optimization of predictions. This is because the step size is directly multiplied by the prediction gradients. Particularly, smaller step size results in larger generated gradients, whereas a larger step size results in smaller gradients. Therefore, the step size needs be tuned per modality, as the update required for predictions depends on data. It's also worth noting that we do not weight decay the step size in any of the models.
We also found that a relatively high optimization step size was necessary (30, 000 for video and between 5 and 500 for text). Without a high optimization step size, gradient magnitudes continued to increase throughout training, resulting unstable training dynamics.
1.1.2 ARCHITECTURE STABILITY
For autoregressive models, we found that simply prepending a learnable "step" embedding to sequences significantly improved scalability and stability, especially for S1 models. This step embedding mapped discrete step index (i.c., step 0, 1, etc.) the current optimization step to an embedding the same dimension as the model's embedding. We believe this helped improve stability
40

---

Published as a conference paper at ICLR 2026
by enabling the accumulation of attention mass, well enabling less steep energy landscapes conditioned on the optimization step.
Additionally, experimented with adaptive layer normalization from the DiT architecture, but we found that the timestep embedding worked better. We also experimented with several different normalization and initialization approaches, where found that the standard Llama2 Touvron et al. (2023) architecture and initialization worked best. This involves using RMSNorm, Xavier init Glorot & Bengio (2010), SwiGLU MLP Shazeer (2020), and RoPE Su et al. (2024).
1.1.3 SYSTEM THINKING HYPERPARAMETER STABILITY
For S2 models, we found certain hyperparameters to be essential for the stability and scalability of models. First, we found that using lower number of optimization steps resulted in more stability, as using more optimization steps necessitates longer gradient chains. Additionally, we found that the strategy used for the randomization of a to be very important for stability. Particularly, we found that randomizing with a single value for an entire batch resulted in issues with training convergence. We believe this is because using the same value for every element in the batch resulted in high-variance gradients. Thus, when randomizing differently for every batch and sequence element, we found training convergence to be more stable. It's possible that randomizing the number of optimization steps would yield similar results reducing gradient variance, however, did not experiment with such a configuration.
1.1.4 CLAMPING OPTIMIZATION GRADIENTS FOR STABILITY
We found that clamping gradients of the energy function with respect to predictions (or prediction update gradients) could help improve training stability at the cost of some slight reductions in convergence speed. We do not conduct any experiments the paper with clamped prediction gradients as we found that the other hyperparameters were sufficient for stable training, however, it's a potentially useful trick worth discussing.
1.1.5 NORMALIZING DATA FOR STABILITY
We found that normalizing/saandddiinng input data crucial for the stability of EBTs. An example of this was in our NLP experiments, where we ran experiments with and without normalizing proba- bility distributions. The experiments with unnormalized distributions often had extreme activations as well as large loss spikes, whereas the experiments with normalized distributions (by applying softmax) were stable.
1.2 How TO TRAIN YOUR EBT
The hyperparameters used for training EBT's are extremely important and can often be highly sensitive towards performance, so here we offer a guide toward hyperparameter tuning. First, we recommend starting off with training S1-EBTs, which are the easiest and most stable variant of EBTs not designed for System 2 Thinking. Then, once $1-EBT's can be trained successfully and scaled, we recommend changing the hyperparameters gradually towards the hyperparameters used for S2 models (we say gradually here as occasionally these parameters can cause instability and require additional tuning).
When training S1 EBTs, recommend tuning hyperparameters in the following order:
First, tune standard hyperparameters such Learning Rate (LR), batch size, etc. Having high batch size helps with stability by making gradients less noisy (because you initialize predictions from random noise this makes gradients noisier).
Second, start tuning S1 S1-specific hyperparameters primarily alpha and its LR multiplier (we recommend keeping its LR multiplier around 3z the value of alpha) and then tuning the number of optimization steps.
Third, potentially tune whether the step size is learnable and try other EBT architectures (inspired by DiT Peebles & Xie (2023) we tried time embedding as well as adaptive layer normalization).
41

---

$ Porform gradient descent w.r.t. the energy function where self. alpha is the optimization step size prediiceed_distribittt deecleteed.isttibautt.o Published as a conference paper at ICLR 2026
Once you have tuned these, the model should be stable and fine for most cases. At which point, if you are desiring System capabilities, you can proceed to the S2 models L.3. Some potential metrics to monitor and look out for include the gradient magnitudes (if these increase too much spike lot that's a bad sign) and the gap between the initial and final energy after optimization (if this is too high low it could be sign your model's alpha value needs be adjusted).
Following Wang et al. (2023), we give pseudocode for training EBTs in natural language for language modeling (Listing 1) well as computer vision for autoregressive video modeling (Listing 2). The pseudocode is primarily for S2 models without any energy landscape regularization techniques. The first primary design decision the presented pseudocode whether or not detach predictions in between steps. Not detaching predictions in between steps allows for more "Thinking Time" before making predictions, but makes the gradient computation graph longer and therefore increases the likelihood of stability issues with gradients. Similarly, calculating the loss at every step versus solely the last step enables model to "think for longer" before needing make accurate predictions, and therefore affects the "funnel-like structure! 11 and smoothness of the energy landscape. For S2 models, we found that not detaching between steps was best, and similarly that calculating the loss only at the last step was best. For S1 models, we found the opposite to be most stable. Generally, if one is calculating the loss only at the last step, then one should not detach between steps as it's i's best the gradient propagates previous steps in this reader to the source code as well as Section I.3.
For more details on these techniques, we refer the
make sure to enable gradient tracking with torch. wrrch.set_grad_eee
(True): loss in nn. CrssssrrtrtpopLcsssss (weight=None, ignore_ index= tokenizer sade_tkee___
context_embeddinggs
faf.fembbedingg (input_ids [:, :-11) 8,
next_tokens input_ids[; 1:1
next next_embeddings = self.embeddingg... (next_tokens) # B, S, Vy are just
used for shaping next tensor
= terch.randn, like (eeext__mbeddiins) B, 8, initialize predictions as random
beciictedddistriuuttee
for in range (num :
$ Can optionally detach predicted distributions that no gradient flows through past ateps predicted_di stttiiutiitiiiitionn predicted stri detach ()
rredicttd_embeddingg = vocab to embed (softmax predicted distributions)) # B, S, D; need to proj. to embed space for transformer to work in, linear layer, weighted sum, etc
aall_embeddings = torch.cat
predicted_ embeddings)) dim=l) # B, 28, predicted_energeee
transforeefrrmmr (al1 fall_embeddiings B, this returns only energies for the
# Compute the gradient et predicted energies w.r.t. predicted distributions
rredicceesdistrrrbitiooo...o credeeegeeeeeerriees sum torch torch.autograd..raal
predicted_distribitt
create giitte_grpphTrree
) [0] # B, S,
prediiteed self.alpha
"'hile one may he eager to classify these landscapes as convex, such properties unlikely in high- dimensional spaces spanned by highly nonlinear neural networks.
42

---

Published as a conference paper at ICLR 2026
Calculate loss based predicted and ground truth distributions, optionally at each optimization step or only the end
ecc_loss =
ions, next_tokens)
Listing I: Autoregressive Language Model Training Pseudocode in PyTorch
make sure to enable gradient tracking
with
(True)
loss_fn torch, nn. Smoothl.1 less (beta=1. 0) use wh chever loss function desi red
context_ embeddinnm embeddings 1;, :-1] $ B, S, next embeddings embeddings :, 1: B, S, predicted embeddings = torch. randn like (next__eeeddeeeee embeddings) B, S, D; initialize predictions as random
for range (num_steps):
# Can optionally detach embeddings so that no gradient flows through past steps
# predictted_embeddingg predictted_embaedddng..cco detachh
all_ _i__embedd torch. cat (ccontext_embeddian
predicted_embeete
dim=1) # B, 2s, predicted ereeiited_energies returns only energies for the predi eted_enbeddingg B, S,
transtrrasformerr (all embeddings) # this
# Compute the gradient of predicted energies w.r.t. predicted embeddings
rredicted_emmedddinns grad corch.autograd. g
predicted_energiess.co snurgies.sum))
predictee__mmeedddane
create Sreate_graph-Trie ) [0] # B, S, D
Perform gradient descent w.r.t. the energy function where self. al alpha is the optimization step size
srddiicted_embeddee = serelictedembedddinnn self.alpha predicted embeddinng _gad
Calculate reconstruction less based on predicted and ground truth embeddings, optionally at each optimization step or only at the end reconstruction_J0 = loss_fn
next_ eext_embedd:ngs)
Listing Autoregressive Video Model Training Pseudocode in PyTorch
1.3 How TO THINK USING YOUR EBT
Once an SI EBT has been trained, we recommend tuning the System 2 hyperparameters in the following manner:
Remove detaching tensors between optimization steps and add loss truncation the loss is only calculated at the final step.
Then, tune alpha, as it's by far the most important EBT-specific hyperparameter. But, do not make it learnable.
Next, tune the number of optimization steps, including potentially a minimum and maximum number when randomizing the number of steps.
Then add a replay buffer, Langevin Dynamics, and eventually a randomized alpha (step size). Tune all of these in tandem while tweaking the earlier parameters (particularly alpha).
It's possible that randomizing the number of steps for each sample within a batch would work better (similar randomizing the alpha value within batch). Additionally, it's worth mentioning that all
43

---

Published as a conference paper at ICLR 2026
optimization steps performed along the same energy landscape (same time embedding condition), unlike with S1-EBTs using multiple time steps.
44


Reviews:
[{"review_id": "3K0YLg3RXV", "reviewer": "Reviewer_k52K", "summary": "This paper proposes a novel framework of an energy-based transformer model. This model integrates a new method that allows efficient and scalable training of EBMs, and thus allows inference-time verification.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- I think the core idea, integrating self-verification in inference time, is promising.\n- The empirical results verify the effectiveness of the proposed model.", "weaknesses": "- Although this paper is overall clear, I think the authors have spent too much space introducing related concepts, instead of the technical details of the proposed method. It is unclear how the model specifically works. It would be good if there could be an extra section to explain the details of the model, perhaps through a toy example.\n- It is unclear why \"frame EBM as an optimization problem\" can avoid the curse of dimensionality. All evidences provided in the paper are just vague discussions, and I don't see any formal analysis. Also, it is unclear how this \"frame EBM as an optimization problem\" actually works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "df2KXzh3cm", "reviewer": "Reviewer_nhmE", "summary": "This paper introduces Energy-Based Transformers (EBTs), a new framework that integrates energy-based modeling with Transformer architectures to enable what the authors describe as “System 2 Thinking.” Rather than directly predicting outputs, EBTs assign an energy value to input–output pairs, representing their compatibility, and perform inference through energy minimization. This allows the model to iteratively refine predictions and dynamically allocate computation based on task difficulty. The authors frame this as a form of unsupervised reasoning or self-verification. To make this scalable, they introduce several key techniques for stabilizing EBM training, including energy landscape regularization, Langevin noise, and replay buffers. EBTs are evaluated across language and vision domains, showing improved scaling behavior over standard Transformers, better generalization on out-of-distribution data, and significantly higher efficiency than Diffusion Transformers for denoising tasks.", "questions": "1. You report that EBTs require 3.33-6.66× more FLOPs than standard Transformers during training. However, the paper doesn't compare what standard Transformers could achieve with an equivalent computational budget (e.g., training with 3× more parameters, 2× depth, or simply training longer). Could you provide an ablation where Transformer++ is given the same total compute budget as EBT? This would help disentangle whether improvements come from the EBM formulation itself or simply from using more computation. Additionally, what are the wall-clock training times and memory requirements compared to Transformer++ on identical hardware?\n2. Section B.3 says that EBTs fail on text-to-image generation with multiple modes, producing blurred images. This seems fundamental, yet the paper broadly claims superior generalization for EBTs. Can you clarify: (a) What classes of problems will inherently fail with EBTs due to such limitations? (b) Is there a fundamental trade-off between the strong OOD generalization you demonstrate and the ability to handle multimodal output distributions? (c) For language modeling specifically, how does this limitation manifest (e.g. does it affect creative generation tasks that benefit from diversity)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The core contribution is genuinely novel and conceptually elegant. The insight that verification is easier than generation (grounded in complexity theory) is well-motivated, and coupling the verifier and generator through energy gradients avoids adversarial dynamics. The cross-modal validation across language, video, and images demonstrates generality beyond domain-specific tricks. Most compellingly, Figure 6 showing that thinking gains increase with distributional shift mirrors human cognition and suggests the approach addresses fundamental generalization challenges. The experimental analysis is thorough with good ablations and comprehensive scaling experiments across six axes.The technical execution is solid, with practical solutions to known EBM training challenges (replay buffers, Langevin dynamics, path randomization) and honest discussion of limitations including text-to-image generation failures. The OOD generalization results showing better downstream performance despite worse pretraining perplexity are particularly interesting. The extensive implementation details and promised code support reproducibility, and the work opens important research directions around reasoning without external supervision and uncertainty quantification in continuous spaces.", "weaknesses": "The scale limitations severely undermine the paper's claims. All experiments max out at 800M parameters while making assertions about foundation model behavior through extrapolation. The paper uses extrapolation to larger sized for many claims, which is speculative given that scaling laws often break at different regimes. Without validation the central claims about foundation model potential remain unsubstantiated.\nThe computational costs (3.33-6.66× FLOPs for training, gradient computation overhead at inference) are dismissed too casually as a temporary inconvenience rather than a fundamental barrier. At foundation model scale, this multiplier translates to enormous real-world costs. The paper lacks wall-clock time comparisons, memory analysis, or ablations showing what standard transformers could achieve with the same computational budget. The \"System 2 Thinking\" here is essentially inference-time optimization and the analogy to human reasoning is somewhat overstated, as it doesn't incorporate concrete verbose reasoning.\nCritical baselines are missing, particularly comparisons to recent inference-time scaling methods. The Transformer++ baseline showing \"no improvement\" with more compute is somewhat unfair since it wasn't designed for that.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NK5s2qHgP5", "reviewer": "Reviewer_8YzV", "summary": "This paper proposes EBTs that learn to verify the compatibility between inputs and candidate predictions through energy minimization. EBTs address limitations of existing System 2 Thinking approaches by enabling dynamic computation allocation and prediction verification across modalities without requiring additional supervision. The work introduces techniques for stable and parallelizable training of EBMs, demonstrating improved performance over Transformer++ and DiTs on language modeling and image denoising tasks while using fewer forward passes.", "questions": "Please refer to the weaknesses section for specific technical questions. This work presents a novel approach to System 2 Thinking through Energy-Based Transformers, which provides a very novel perspective and contributes to the community. The work meets the acceptance bar for ICLR, but based on the noted issues, cannot receive an absolutely high score.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- This paper is well written with clear motivation analysis, comprehensive experimental design, and informative figures that effectively support the main arguments. \n- The core insight of building generation models based on the principle that verification is easier than generation is novel and insightful. This approach provides a fresh perspective on how to design generative models by leveraging verification capabilities.\n- The dynamic computation allocation based on problem difficulty is an interesting contribution that offers an alternative technical pathway for adaptive computation. This differs from common approaches that rely on layer skipping and provides a new way to handle varying computational requirements.\n- The experiments are comprehensive, particularly in the language modeling domain where multiple scaling dimensions are explored to validate the proposed method. The extensive evaluation across different model sizes, data scales, and computational budgets demonstrates the robustness and effectiveness of the approach.", "weaknesses": "- The verification approach in this work does not align with intuitive expectations. For language tasks, meaningful verification should operate on complete sentences or full responses to questions. However, EBT uses autoregressive generation where each token is verified independently during generation. This raises questions about the validity of verification at the token level rather than at the semantic level of complete thoughts or responses.\n\n- This work lacks sufficient analysis of computational efficiency. The paper primarily uses forward pass counts as an indirect measure, but absolute metrics such as latency or other time-based measurements would provide a more objective assessment of EBT's computational performance compared to Transformer++. Such analysis is crucial for understanding the practical viability of EBT.\n\n- The implementation details of EBT are not clearly explained. The training process involves gradient descent optimization for token generation, which likely requires second-order gradient computation. However, this paper does not provide sufficient explanation of how second-order gradients are computed, their computational cost, or their impact on training efficiency.\n\n- The attention mechanism in EBT may face efficiency challenges with second-order gradient computation. Modern LLMs and DiTs rely heavily on optimized implementations like Flash Attention for CUDA acceleration. However, these optimizations may not be compatible with second-order gradients, potentially limiting the computational efficiency of EBT's modified attention mechanism.\n\n- The DiT experiments focus on image denoising, which is not representative of typical industrial applications. Most DiT usage in practice involves image generation or video generation tasks where denoising acceleration is highly demanded. This work would benefit from experiments on these more relevant applications to demonstrate EBT's effectiveness in real-world scenarios.\n\n- Figure 5(a) and (b) show concerning patterns where EBT curves appear above the baseline in the upper right region. Given that the x-axis represents model size or compute and the y-axis represents perplexity, this suggests EBT achieves higher perplexity under similar conditions, which is not a positive trend. This may indicate a plotting issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gaTXkSTDdD", "reviewer": "Reviewer_KJ9J", "summary": "This paper presents Energy-Based Transformers (EBTs), a framework that reconceptualizes prediction as an optimization process rather than direct generation. Instead of emitting outputs (tokens, pixels, etc.) in a single forward pass, EBTs learn an energy function that evaluates the compatibility between inputs and candidate predictions. Inference then proceeds through gradient-based refinement toward lower-energy (more consistent) states.", "questions": "------\nSome additional questions:\n- Thinking longer presents an interesting argument. I know that some prior ideas like \"HRM/TRM\", and \"Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach\", have explored similar ideas (thinking longer, although they don't have explicit verification mechanism).-\n- It would be great to see where EBTs perform well in CONSTRAINED data settings, for eg, solving extreme-sized sudoku puzzles, mazes, and unsupervised object discovery.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The approach enables two appealing properties often associated with \"System 2\" reasoning: (i) adaptive computation, where more challenging inputs naturally require more optimization steps, and (ii) explicit self-evaluation, since energy levels can serve as confidence or verification signals.\nEmpirically, the paper demonstrates competitive performance across discrete and continuous domains and reports improved robustness under distribution shifts. The idea of learning to verify rather than learning to generate feels conceptually elegant, and the result ,especially the scaling behavior and cross-domain generality, suggest that EBTs may offer a promising direction for inference-time reasoning emerging directly from unsupervised pretraining.", "weaknesses": "Weaknesses\n\n- Due to computational limitations, experiments are limited to models of up to 800M parameters, making claims about scaling advantages at the billion-parameter or foundation scale speculative. We can discount the requirement for experiments in billion/trillion parameter range, due to constraint of academic resources. However,  the projected breakeven point at which EBTs would become FLOP-efficient relative to Transformer++ is estimated at roughly 10⁵² total FLOPs, a scale that is effectively beyond practical reach. For context, even a state-of-the-art exascale system (≈10¹⁸ FLOPS) would require approximately 10³⁴ seconds of continuous computation to reach this point. Hypothetically distributing the workload across all 8 billion people on Earth, each operating a dedicated exascale supercomputer 24/7 (≈8×10²⁷ FLOPS combined), the breakeven computation would still demand about 1.25×10²⁴ seconds, equivalent to 3.96×10¹⁶ years, or nearly 400 trillion human lifetimes (assuming 100 years per lifetime) .Such an extreme requirement implies that the nominal scaling crossover occurs far beyond any conceivable computational regime. Consequently, the reported scaling advantage, while theoretically positive, lacks practical significance: in realistic operating regimes, EBTs remain orders of magnitude less compute-efficient than Transformers, making it difficult to justify claims of superior scalability.\n\n\n- EBTs require roughly 10× more compute to achieve the same validation perplexity, calling into question their current efficiency. The suggestion that EBTs could surpass advanced Transformer variants remains speculative without stronger empirical backing. The argument that the slope \"suggests\" that EBTs would surpass transformer++ at high levels of compute, while correct, would benefit from more backing. \n- The scaling law analysis lacks transparency: the procedure for estimating exponents and the dataset used for fitting are not described, making reproducibility difficult. Including a detailed table of fitted data points would improve credibility.\n- The iterative inference procedure introduces multiple hyperparameters (e.g., step count, learning rate, noise levels) that require tuning. This makes deployment and fair comparison more involved than with standard feedforward Transformers.\n\n\n- The method incurs substantially higher computational cost per token, yet the comparisons are not normalized for equal FLOPs. Without such baselines, it is unclear whether the observed gains stem from architectural benefits or simply increased compute. For text generation in particular, where predictions are locally conditioned, the incremental benefit appears modest relative to diffusion-style alternatives.\n\n\n- The reported “2.91% faster scaling” claim (highlighted in Figure 5) lacks methodological clarity, no explanation is given for how this value was computed or whether variance in scaling fits was accounted for.\n\n\n- It is not specified whether models in Table 3 were compared under matched computational budgets. Clarifying this would improve interpretability of the downstream comparisons.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)