This video discusses a research paper titled "Energy-based transformers are scalable learners and thinkers," which proposes combining energy-based models (EBMs) with transformers to develop scalable, unsupervised "system two thinking" capabilities in AI. The paper redefines "thinking" as inference-time computation and argues that EBMs inherently support this, along with dynamic allocation of compute, uncertainty modeling, and prediction verification, unlike traditional deep learning models.
The core of the approach involves training an energy function (a transformer) that assigns low energy to compatible input-output pairs and high energy to incompatible ones. Inference is then performed by optimizing (e.g., via gradient descent) this energy function to find the best output, requiring multiple forward passes. To ensure scalability and smoothness of the energy landscape, training incorporates techniques like backpropagation through optimization steps (using second-order derivatives), adding noise, and randomizing optimization steps. While EBMs have a higher initial computational cost than classic transformers, experimental results on smaller datasets show promising scaling trends for both learning efficiency and the ability to improve performance with increased inference-time computation ("thinking").
Summarizing "Energy-based Transformers are Scalable Learners and Thinkers"
0:00 Introduction to Energy-Based Transformers (EBTs): The paper explores combining energy-based models (EBMs) with transformers to create scalable learners and "thinkers," aiming for unsupervised "system two thinking."
1:19 Defining "System Two Thinking": System two thinking, in cognitive science, refers to slow, logical, explicit thought (vs. rapid, intuitive system one thinking). In the paper's context, it primarily refers to inference-time computation where a model performs multiple steps to arrive at a conclusion, going beyond a single forward pass.
6:07 "Thinking" as a Metric: The paper views "thinking" not as an act but as a metric, quantifying how much performance improves by investing more computation (multiple forward passes) at inference time.
7:28 Energy-Based Models (EBMs) Overview: EBMs use an "energy function" that assigns a low numerical energy to compatible input-output pairs and a high energy to incompatible ones. Unlike GANs, the energy function itself (a transformer in this case) is the only learned parameter.
13:58 Inference in EBMs: Instead of a direct prediction, inference starts with a random distribution and iteratively refines it through an optimization procedure (e.g., gradient descent) on the energy function. This allows for dynamic computation allocation – more steps can lead to a more accurate output.
15:00 "Thinking" as Optimization: The paper proposes viewing thinking as an optimization procedure with respect to a learned verifier (the energy function) that evaluates compatibility between inputs and candidate predictions.
26:01 Gradient Descent at Inference Time: The process involves calculating the gradient of the energy function with respect to the output and backpropagating it through the transformer to adjust the output iteratively. This allows for modeling uncertainty and identifying out-of-distribution data.
27:42 Challenges in Training EBMs: Naive training can lead to jagged, non-smooth energy landscapes. Overcoming this requires "energy landscape regularization techniques."
29:48 Training with Inference in Mind: EBMs are trained by backpropagating the loss through the entire multi-step optimization process, which requires second-order derivatives (gradients of gradients) computed efficiently via Hessen vector products.
32:58 Regularization Techniques: Key techniques for smooth, scalable energy landscapes include:
33:09 Adding Noise: Injecting noise into each gradient descent step during training helps generalize and smooth the energy landscape by "washing out" thin paths.
35:28 Randomizing Step Size & Steps: Varying the gradient step size and the number of optimization steps during training improves generalization and makes the model more robust to different inference compute budgets.
36:40 Combining EBMs and Transformers: Transformers address traditional EBM challenges like parallelizability, stability, and scalability, making them a natural fit.
38:12 Experimental Focus: The paper investigates "learning scalability" (how quickly models fit pre-training data) and "thinking scalability" (how performance improves with increased inference-time computation).
39:03 Learning Scalability Results: While EBTs start at a slight disadvantage, they show promising trends, scaling better than classic transformers with increased training data, batch size, and depth, despite a higher initial FLOP cost.
40:50 Higher Initial FLOP Cost: EBTs require significantly more training FLOPs per step than classic transformers due to the embedded inference optimization. However, their steeper scaling slope suggests potential long-term efficiency gains at very large scales.
42:09 Thinking Scalability Results: EBTs demonstrate significant improvement in perplexity as the number of inference-time forward passes increases, outperforming classic transformers which do not benefit from additional passes for the same input.
43:13 Dynamic Inference Potential: The model's energy values can indicate the "easiness" or uncertainty of a prediction, potentially allowing for dynamic allocation of compute – spending less "thinking" on easier predictions.
45:41 Promising Trends, Limited Scale: The scaling trends are promising, but the authors acknowledge resource constraints limit testing at larger scales, so the extrapolated benefits are not yet fully verified.
46:18 Speaker's Critique: The presenter notes the paper's heavy philosophical framing and questions the direct comparability of EBMs with existing "reasoning" models like Chain-of-Thought, suggesting EBM inference is a distinct mechanism that could then be combined with such techniques.