Energy-Based Models assign a numerical score, called energy, to each candidate output. A candidate output might be an image, a scene interpretation, a reasoning solution, or a robot trajectory. Lower energy means that a candidate better fits the input or constraints. Rather than producing an answer in one fixed way, an EBM learns a landscape. Once learned, we can search this landscape for one low-energy solution or sample from it to obtain several plausible solutions. We can also compose models by adding their energies, then search for outputs that score well under their combination.
Our research uses learned energy landscapes both to model which outcomes are plausible and as reusable problem representations that can be combined and searched after training to solve new generation, reasoning, and planning problems.
Browse publications ↓ Implicit generation with EBMs (2019) ↗
An EBM learns a landscape over possible outputs. The landscape tells us which outputs fit the input, but it does not specify how outputs should be generated. Below, x is the input, y is a candidate output, and θ denotes the learned model parameters. During training, we learn θ; during inference, θ remains fixed while we search over y for one low-energy output or sample several plausible outputs. Because the energy function remains available during inference, we can add new energy terms representing additional goals or constraints. If the EBM is trained as a joint energy over several variables, we can also fix any observed variables and infer the others.
Representation
pθ(y | x) = exp(−Eθ(x, y)) / Zθ(x)
The energy induces a probability distribution in which lower-energy outputs receive higher probability.
How the energy is learned. One common approach is maximum-likelihood training. It lowers the energy of observed input-output pairs and raises the energy of alternatives sampled from the current model. Repeating this process shapes the model distribution toward the training data.
What Z does. Zθ(x) is the partition function: the sum or integral of exp(−Eθ(x, y)) over every possible output y. Dividing by Z makes the probabilities add up to one. For large or continuous output spaces, this calculation is usually too expensive to perform exactly, so training uses samples or objectives that avoid evaluating Z directly.
One solution (MAP)
ŷ = arg miny Eθ(x, y)
Search for the lowest-energy output. In probabilistic terms, this is the maximum a posteriori (MAP) solution.
Multiple solutions (Langevin)
yt+1 = yt − η∇yEθ(x, yt) + √(2η) εt
Langevin dynamics is a gradient-based Markov chain Monte Carlo (MCMC) procedure: it repeatedly moves a candidate toward lower energy while adding noise. Different initializations and noise can produce different plausible outputs.
Both procedures use the same learned energy function. MAP is useful when one best solution is needed; Langevin dynamics approximately explores the distribution and can produce diverse samples. In the Langevin equation, η is the step size and εt is Gaussian noise. This form assumes that y is continuous and differentiable; discrete outputs require a different search or relaxation procedure.
An EBM assigns an energy Eθ(x, y) to every candidate output y given an input x. This energy defines the conditional probability distribution pθ(y | x), where the partition function Zθ(x) normalizes the probabilities over all possible outputs.
Representing a distribution with an EBM is simple: a scalar energy function defines an unnormalized density—relative likelihoods before they are normalized into a probability distribution. Autoregressive models represent an output as a fixed sequence of predictions, while diffusion models learn a multi-step denoising process. An EBM does not require either structure. These categories can overlap: a diffusion model can be parameterized so that its denoising direction is derived from an energy at each noise level. The EBM representation is simple and flexible, but the partition function is generally intractable and both training and inference can be difficult.
Many conventional predictors produce an answer through one fixed forward computation. An EBM instead retains a landscape that an inference procedure can search, sample, evaluate, or combine with new requirements.
Composition has a concrete meaning for Energy-Based Models: each energy function can represent one piece of knowledge, such as dynamics, a goal, or a constraint. We add the energies and then search for an output with low total energy. Intuitively, low total energy corresponds to satisfying the represented goals and constraints together. A robot planner, for example, can combine a learned dynamics energy with an energy for a new goal, then find a trajectory that follows the dynamics and reaches the goal. Because the component models remain separate, we can change the goal or add another constraint without retraining the full system.
My PhD thesis, Learning Generalizable Systems by Learning Composable Energy Landscapes, discusses some of the advantages of EBMs: we can learn reusable energy landscapes from data, then compose and optimize them at inference time to solve tasks not represented directly in the training set.
An EBM separates the representation of a problem from the procedure used to solve it. This separation is useful in several different ways.
An energy function only needs to score a complete configuration. Unlike an autoregressive model, it does not need to generate the variables one at a time in a fixed order. Unlike a diffusion model, it does not require a prescribed process that adds noise and then learns to reverse it.
When concepts, relations, dynamics, goals, and constraints are encoded as compatible energies over shared variables, their weighted sum creates a new landscape that balances these components. We can therefore form a new inference problem without retraining the models jointly.
Each model can score a small part of a problem, while the summed energy defines the full structured problem. A relation learned between two objects can be reused in a scene with many objects, and a local dynamics model can be applied repeatedly across a longer plan. The same models can therefore support different numbers of objects, longer horizons, or new task layouts.
An energy function does not fix a single way of producing an answer. We can minimize it for one solution, sample from it for several possibilities, or combine it with a utility or constraint to make a decision. Harder problems can use more optimization steps or maintain more candidate solutions at inference time.
A model is useful only if inference can reach its low-energy regions. We can therefore shape the learned landscape together with the solver, so that optimization or sampling can find good solutions rather than merely assign them low energy. For example, inference can begin on a smoother landscape and gradually sharpen it, optimize several candidates in parallel, or train the landscape specifically so that the chosen solver can reach good solutions.
If one energy function scores x and y together, we can hold x fixed and infer y, or hold y fixed and infer x. This requires a joint model; the conditional formula above only models y given x. For example, the same relational model can generate a scene that satisfies several relations or infer the relations present in an observed scene, as in Learning to Compose Visual Relations.
EBMs connect probabilistic modeling with optimization, sampling, and inference-time computation.
When used as a generative model, an EBM assigns greater relative probability to lower-energy data. This defines an unnormalized density: the model specifies relative likelihoods without computing the constant needed to turn them into a probability distribution. These relative likelihoods let us rank candidates directly and identify the most likely one. MCMC methods such as Langevin dynamics generate samples by repeatedly moving a candidate toward lower energy while adding noise.
Diffusion models learn how to move a noisy sample toward more likely data at many noise levels. In an energy-based parameterization, each such direction is given by the negative gradient of a scalar energy function. This connection lets us use energy-based sampling and composition within diffusion models. In our work, we study energy-based diffusion samplers, as well as how to blend diffusion objectives to build equilibrium energy landscapes.
To compose EBMs, we add their energies and run inference in the resulting landscape. When component scales are compatible or explicitly weighted, low-total-energy candidates tend to score well across them. Probabilistically, adding energies multiplies the models’ relative likelihoods. The result can nevertheless fail if one model’s scores dominate or the sampler cannot reach a region that scores well across the components.
An energy function can score candidate answers or trajectories without generating them in one pass. Inference searches across candidates and refines the low-energy ones. Harder problems can use more optimization steps or maintain several candidates in parallel, while new goals or constraints can enter as additional energy terms.
Composition is most reliable when each model is used on local situations it understands and the models agree on the meaning and scale of their shared variables. For example, combining a dynamics model with a new goal is safer when the resulting trajectory remains within states covered by the dynamics model. The global combination may be new, but its individual components should not be arbitrary out-of-distribution inputs.
These advantages depend on effective inference. Computing exact probabilities often requires a partition function—a sum or integral over all possible outputs—that is too expensive to evaluate. Optimization and MCMC can also be slow or sensitive to initialization. When several energies are composed, one component may dominate, the combined landscape may contain new local minima, or the component models may be unreliable in different regions.
Energy-Based Models borrow a simple idea from statistical physics: represent each configuration with a scalar energy, so that lower-energy configurations are more probable or more compatible. Hopfield networks used an energy landscape to describe attractor dynamics and associative memory. Boltzmann machines made these dynamics stochastic and showed how an energy function could be learned from data.
Later work developed more practical and general ways to learn energy functions. Hinton’s Products of Experts introduced contrastive divergence for training product-form latent-variable models. Score matching and noise-contrastive estimation provided other ways to learn unnormalized densities without directly evaluating the partition function. LeCun et al. unified these ideas as energy-based prediction: score candidate outputs, then optimize or search for a low-energy one. This work emphasized prediction and structured outputs, while later neural EBM research placed greater emphasis on MCMC sampling for generation.
In our 2019 paper, we showed that MCMC-based training of continuous neural EBMs could scale to high-dimensional images and robotic trajectories, using Langevin dynamics for both learning and generation. We also studied compositional generation, reconstruction, robustness, continual learning, and trajectory prediction. The later work on this page develops one part of this picture further: energy functions as reusable representations that can be optimized, conditioned, and composed at inference time.
We introduce Equilibrium Matching (EqM), a generative modeling framework built from an equilibrium dynamics perspective. EqM discards the non-equilibrium, time-conditional dynamics in traditional diffusion and flow-based generative models and instead learns the equilibrium gradient of an implicit energy landscape. At inference time, EqM initializes candidates from noise and generates samples by optimizing the learned landscape with gradient descent. Different initializations can produce different outputs, and the optimization can use adjustable step sizes, adaptive optimizers, and adaptive compute. EqM surpasses the generation performance of diffusion/flow models empirically, achieving an FID of 1.90 on ImageNet 256×256. EqM is also theoretically justified to learn and sample from the data manifold. Beyond generation, EqM is a flexible framework that naturally handles tasks including partially noised image denoising, OOD detection, and image composition. By replacing time-conditional velocities with a unified equilibrium landscape, EqM offers a tighter bridge between flow and Energy-Based Models and a simple route to optimization-driven inference.
Generalization is a key challenge in machine learning, specifically in reasoning tasks, where models are expected to solve problems more complex than those encountered during training. Existing approaches typically train reasoning models in an end-to-end fashion, directly mapping input instances to solutions. While this allows models to learn useful heuristics from data, it often results in limited generalization beyond the training distribution. In this work, we propose a novel approach to reasoning generalization by learning energy landscapes over the solution spaces of smaller, more tractable subproblems. At test time, we construct a global energy landscape for a given problem by combining the energy functions of multiple subproblems. This compositional approach enables the incorporation of additional constraints during inference, allowing the construction of energy landscapes for problems of increasing difficulty. To improve the sample quality from this newly constructed energy landscape, we introduce Parallel Energy Minimization (PEM). We evaluate our approach on a wide set of reasoning problems. Our method outperforms existing state-of-the-art methods, demonstrating its ability to generalize to larger and more complex problems.
Inference-time computation techniques, analogous to human System 2 Thinking, have recently become popular for improving model performances. 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?” Interestingly, 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 to every input and candidate-prediction pair, with lower energy indicating greater compatibility, and obtain predictions through gradient-based energy minimization. This formulation enables System 2 Thinking to emerge from unsupervised learning, making it modality and problem agnostic. Across both discrete (text) and continuous (visual) modalities, we find EBTs scale faster than the dominant Transformer++ approach during training, achieving an up to 35% higher scaling rate with respect to data, batch size, parameters, FLOPs, and depth. During inference, EBTs improve performance with System 2 Thinking (i.e., extra computation) by 29% more than the Transformer++ on language tasks, and EBTs outperform Diffusion Transformers on image denoising while using fewer forward passes. Further, we find that System 2 Thinking with EBTs yields larger performance improvements on data that is farther out-of-distribution, and that EBTs achieve better results than existing models on most downstream tasks given the same or worse pretraining performance, suggesting that EBTs generalize better than existing approaches. Consequently, EBTs are a promising new paradigm for scaling both the learning and thinking capabilities of models.
We explore how generative models can be used not only to synthesize visual content but also to understand the properties of a scene given a natural image. We formulate scene understanding as an inverse generative modeling problem, where we seek to find conditional parameters of a visual generative model to best fit a given natural image. To enable this procedure to infer scene structure from images substantially different than those seen during training, we further propose to build this visual generative model compositionally from smaller models over pieces of a scene. We illustrate how this procedure enables us to infer the set of objects in a scene, enabling robust generalization to new test scenes with an increased number of objects of new shapes. We further illustrate how this enables us to infer global scene factors, likewise enabling robust generalization to new scenes. Finally, we illustrate how this approach can be directly applied to existing pretrained text-to-image generative models for zero-shot multi-object perception.
This thesis develops energy landscapes as a representation for prediction problems. Test-time search can incorporate new constraints, while compositions of learned landscapes construct models for unseen combinations of factors. It develops methods for probabilistic, deterministic, and annealed energies; an algebra for logical, probabilistic, graphical-model, and hierarchical composition; and applications across vision, robotics, foundation models, and scientific design.
We introduce iterative reasoning through energy diffusion (IRED), a novel framework for learning to reason for a variety of tasks by formulating reasoning and decision-making problems with energy-based optimization. IRED learns energy functions to represent the constraints between input conditions and desired outputs. After training, IRED adapts the number of optimization steps during inference based on problem difficulty, enabling it to solve problems outside its training distribution -- such as more complex Sudoku puzzles, matrix completion with large value magnitudes, and pathfinding in larger graphs. Key to our method's success is two novel techniques: learning a sequence of annealed energy landscapes for easier inference and a combination of score function and energy landscape supervision for faster and more stable training. Our experiments show that IRED outperforms existing methods in continuous-space reasoning, discrete-space reasoning, and planning tasks, particularly in more challenging scenarios.
Given an image of a natural scene, we are able to quickly decompose it into a set of components such as objects, lighting, shadows, and foreground. We can then envision a scene where we combine certain components with those from other images, for instance a set of objects from our bedroom and animals from a zoo under the lighting conditions of a forest, even if we have never encountered such a scene before. In this paper, we present a method to decompose an image into such compositional components. Our approach, Decomp Diffusion, is an unsupervised method which, when given a single image, infers a set of different components in the image, each represented by a diffusion model. We demonstrate how components can capture different factors of the scene, ranging from global scene descriptors like shadows or facial expression to local scene descriptors like constituent objects. We further illustrate how inferred factors can be flexibly composed, even with factors inferred from other models, to generate a variety of scenes sharply different than those seen in training time.
Effective motion planning in high-dimensional spaces is a long-standing problem in robotics. Potential-based planners are naturally compositional: different motion constraints can be combined by adding their potentials. However, finding a path requires global optimization over the configuration-space landscape and is often vulnerable to local minima. We learn an easily optimized potential over motion trajectories. The resulting planner outperforms classical and learned alternatives, avoids many local-minimum failures, and composes across a wide range of motion constraints.
Large monolithic generative models trained on massive amounts of data have become an increasingly dominant approach in AI research. In this paper, we argue that we should instead construct large generative systems by composing smaller generative models together. We show how such a compositional generative approach enables us to learn distributions in a more data-efficient manner, enabling generalization to parts of the data distribution unseen at training time. We further show how this enables us to program and construct new generative models for tasks completely unseen at training. Finally, we show that in many cases, we can discover separate compositional components from data.
Text-to-image generative models have enabled high-resolution image synthesis across different domains, but require users to specify the content they wish to generate. In this paper, we consider the inverse problem -- given a collection of different images, can we discover the generative concepts that represent each image? We present an unsupervised approach to discover generative concepts from a collection of images, disentangling different art styles in paintings, objects, and lighting from kitchen scenes, and discovering image classes given ImageNet images. We show how such generative concepts can accurately represent the content of images, be recombined and composed to generate new artistic and hybrid images, and be further used as a representation for downstream classification tasks.
Systems consisting of interacting agents are prevalent in the world, ranging from dynamical systems in physics to complex biological networks. To build systems which can interact robustly in the real world, it is thus important to be able to infer the precise interactions governing such systems. Existing approaches typically discover such interactions by explicitly modeling the feed-forward dynamics of the trajectories. In this work, we propose Neural Interaction Inference with Potentials (NIIP) as an alternative approach to discover such interactions that enables greater flexibility in trajectory modeling: it discovers a set of relational potentials, represented as energy functions, which when minimized reconstruct the original trajectory. NIIP assigns low energy to the subset of trajectories which respect the relational constraints observed. We illustrate that with these representations NIIP displays unique capabilities in test-time. First, it allows trajectory manipulation, such as interchanging interaction types across separately trained models, as well as trajectory forecasting. Additionally, it allows adding external hand-crafted potentials at test-time. Finally, NIIP enables the detection of out-of-distribution samples and anomalies without explicit training.
Since their introduction, diffusion models have quickly become the prevailing approach to generative modeling in many domains. They can be interpreted as learning the gradients of a time-varying sequence of log-probability density functions. This interpretation has motivated classifier-based and classifier-free guidance as methods for post-hoc control of diffusion models. In this work, we build upon these ideas using the score-based interpretation of diffusion models, and explore alternative ways to condition, modify, and reuse diffusion models for tasks involving compositional generation and guidance. In particular, we investigate why certain types of composition fail using current techniques and present a number of solutions. We conclude that the sampler (not the model) is responsible for this failure and propose new samplers, inspired by MCMC, which enable successful compositional generation. Further, we propose an energy-based parameterization of diffusion models which enables the use of new compositional operators and more sophisticated, Metropolis-corrected samplers. Intriguingly we find these samplers lead to notable improvements in compositional generation across a wide set of problems such as classifier-guided ImageNet modeling and compositional text-to-image generation.
Large pre-trained models exhibit distinct and complementary capabilities dependent on the data they are trained on. Language models such as GPT-3 are capable of textual reasoning but cannot understand visual information, while vision models such as DALL-E can generate photorealistic photos but fail to understand complex language descriptions. In this work, we propose a unified framework for composing ensembles of different pre-trained models -- combining the strengths of each individual model to solve various multimodal problems in a zero-shot manner. We use pre-trained models as "generators" or "scorers" and compose them via closed-loop iterative consensus optimization. The generator constructs proposals and the scorers iteratively provide feedback to refine the generated result. Such closed-loop communication enables models to correct errors caused by other models, significantly boosting performance on downstream tasks, e.g. improving accuracy on grade school math problems by 7.5%, without requiring any model finetuning. We demonstrate that consensus achieved by an ensemble of scorers outperforms the feedback of a single scorer, by leveraging the strengths of each expert model. Results show that the proposed method can be used as a general purpose framework for a wide range of zero-shot multimodal tasks, such as image generation, video question answering, mathematical reasoning, and robotic manipulation.
Large text-guided diffusion models, such as DALLE-2, are able to generate stunning photorealistic images given natural language descriptions. While such models are highly flexible, they struggle to understand the composition of certain concepts, such as confusing the attributes of different objects or relations between objects. In this paper, we propose an alternative structured approach for compositional generation using diffusion models. An image is generated by composing a set of diffusion models, with each of them modeling a certain component of the image. To do this, we interpret diffusion models as Energy-Based Models in which the data distributions defined by the energy functions may be explicitly combined. The proposed method can generate scenes at test time that are substantially more complex than those seen in training, composing sentence descriptions, object relations, human facial attributes, and even generalizing to new combinations that are rarely seen in the real world. We further illustrate how our approach may be used to compose pre-trained text-guided diffusion models and generate photorealistic images containing all the details described in the input descriptions, including the binding of certain object attributes that have been shown difficult for DALLE-2. These results point to the effectiveness of the proposed method in promoting structured generalization for visual generation.
We motivate Energy-Based Models (EBMs) as a promising model class for continual learning. Instead of relying on external memory, growing models, or regularization, EBMs change the training objective to reduce interference with previously learned information. Our approach is simple, efficient, and outperforms baselines by a large margin on several benchmarks. Its contrastive-divergence objective can also be combined with other continual-learning methods, producing substantial improvements.
Deep learning has excelled on complex pattern recognition tasks such as image classification and object recognition. However, it struggles with tasks requiring nontrivial reasoning, such as algorithmic computation. Humans are able to solve such tasks through iterative reasoning -- spending more time thinking about harder tasks. Most existing neural networks, however, exhibit a fixed computational budget controlled by the neural network architecture, preventing additional computational processing on harder tasks. In this work, we present a new framework for iterative reasoning with neural networks. We train a neural network to parameterize an energy landscape over all outputs, and implement each step of the iterative reasoning as an energy minimization step to find a minimal energy solution. By formulating reasoning as an energy minimization problem, for harder problems that lead to more complex energy landscapes, we may then adjust our underlying computational budget by running a more complex optimization procedure. We empirically illustrate that our iterative reasoning approach can solve more accurate and generalizable algorithmic reasoning tasks in both graph and continuous domains. Finally, we illustrate that our approach can recursively solve algorithmic problems requiring nested reasoning.
We introduce an approach to decompose images, in an unsupervised manner, into separate component energy functions. These energy functions can represent both global factors of variation, such as facial expression and hair color, and local factors of variation, such as the objects in a scene. Decomposed energy functions generalize well and may be recombined with energy functions discovered by training a separate instance of the approach on another dataset, enabling the recombination of objects and lighting conditions across datasets.
The visual world around us can be described as a structured set of objects and their associated relations. In this work, we propose to represent each relation as an unnormalized density (an energy-based model), enabling us to compose separate relations in a factorized manner. We show that such a factorized decomposition allows the model to both generate and edit scenes that have multiple sets of relations more faithfully. We further show that decomposition enables our model to effectively understand the underlying relational scene structure.
We present tools to improve the underlying contrastive divergence objective for training EBMs. First we illustrate a neglected term in contrastive divergence training of EBMs, and present a loss function to mitigate this term. We further propose to utilize data augmentation to aid the mixing of MCMC chains when training EBMs and propose to use a multiscale architecture to further improve the underlying generative performance. We illustrate how our tricks improve the underlying generative performance of EBMs, and further show improved out-of-distribution detection.
A vital aspect of human intelligence is the ability to compose increasingly complex concepts out of simpler ideas, enabling both rapid learning and adaptation of knowledge. In this paper we show that Energy-Based Models can exhibit this ability by directly combining probability distributions. Samples from the combined distribution correspond to compositions of concepts. For example, given one distribution for smiling face images, and another for male faces, we can combine them to generate smiling male faces. This allows us to generate natural images that simultaneously satisfy conjunctions, disjunctions, and negations of concepts. We evaluate compositional generation abilities of our model on the CelebA dataset of natural faces and synthetic 3D scene images. We showcase the breadth of unique capabilities of our model, such as the ability to continually learn and incorporate new concepts, or infer compositions of concept properties underlying an image.
We propose an energy-based model (EBM) of protein conformations that operates at atomic scale. The model is trained solely on crystallized protein data. By contrast, existing approaches for scoring conformations use energy functions that incorporate knowledge of physical principles and features that are the complex product of several decades of research and tuning. To evaluate the model, we benchmark on the rotamer recovery task, the problem of predicting the conformation of a side chain from its context within a protein structure, which has been used to evaluate energy functions for protein design. The model achieves performance close to that of the Rosetta energy function, a state-of-the-art method widely used in protein structure prediction and design. An investigation of the model’s outputs and hidden representations finds that it captures physicochemical properties relevant to protein energy.
Model-based planning holds great promise for improving both sample efficiency and generalization in reinforcement learning (RL). We show that Energy-Based Models (EBMs) are a promising class of models to use for model-based planning. EBMs naturally support inference of intermediate states given start and goal state distributions. We provide an online algorithm to train EBMs while interacting with the environment, and show that EBMs allow for significantly better online learning than corresponding feed-forward networks. We further show that EBMs support maximum entropy state inference and are able to generate diverse state space plans. We show that inference purely in state space - without planning actions - allows for better generalization to previously unseen obstacles in the environment and prevents the planner from exploiting the dynamics model by applying uncharacteristic action sequences.
Energy-Based Models (EBMs) are an appealing class of models due to their generality and simplicity in likelihood modeling. However, EBMs have traditionally been difficult to train. We present techniques to scale MCMC-based EBM training on continuous neural networks in high-dimensional domains such as ImageNet and robotic hand trajectories. We highlight unique capabilities of implicit generation. Finally, we illustrate how EBMs are useful across a wide variety of tasks, including out-of-distribution classification, adversarially robust classification, online continual learning, and compositionality.
No publications match this search. Try a broader term or select another topic.