Energy-Based Models

Energy-Based Models assign a numerical score, called energy, to each candidate output. An output might be an image, a scene interpretation, a reasoning solution, or a robot trajectory. Lower energy indicates a better fit to the input or constraints. The energy values across possible outputs form an energy landscape. 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 and searching for outputs with low total energy.

Our research develops energy landscapes as reusable representations. After training, we combine and search these landscapes to solve new generation, reasoning, and planning problems.

Browse publications ↓Implicit generation with EBMs (2019) ↗

What is an Energy-Based Model?

An Energy-Based Model (EBM) evaluates a proposed output for a given input and returns a single number, called its energy. Lower energy indicates a better fit to the input. In a neural EBM, a network learns this scoring function from data. For example, a robot-planning model can score an entire proposed sequence of states and actions, called a trajectory, given a starting state and a goal.

To make a prediction, we start with an initial guess and repeatedly adjust the proposed output to reduce its energy. For the robot, we change the proposed states and actions while keeping the starting state, goal, and learned network weights fixed. Gradients tell us how small changes to the trajectory affect its energy, so we can use them to guide these updates. This search for an output with the lowest energy is called energy minimization.

The same energy function can also define a probability distribution over possible outputs. Exponentiating the negative energy gives each output a positive weight, so lower-energy outputs receive larger weights. Dividing by the sum or integral of these weights, called the partition function, normalizes the distribution. A sampling procedure can then generate different plausible outputs according to this distribution.

An EBM represents a probability distribution through a single energy function over complete outputs. Autoregressive models factorize the distribution into a product of conditional probabilities, each conditioned on the variables generated so far. Diffusion models introduce noisy intermediate states and factorize the generative process into conditional transitions that progressively remove noise. An EBM does not require either a variable ordering or a sequence of noise levels to define its distribution. Because the energy scores all output variables together, it provides a direct way to represent global dependencies among them. This simple representation comes at a computational cost. Maximum-likelihood training typically requires repeated sampling from the current model, which can be slow and difficult in high dimensions.

We can also evaluate the same candidate with several energy functions. In robot planning, one energy can assess whether a trajectory follows the robot’s dynamics, meaning how states change under actions, and another can score how close its final state is to a goal. Adding the energies gives a single score for both requirements. We then adjust the trajectory to reduce this total energy. Weights control the balance between the requirements. Changing the goal energy creates a new planning problem without retraining the dynamics model.

My PhD thesis, Learning Generalizable Systems by Learning Composable Energy Landscapes, develops this approach: learning reusable energy landscapes from data, then composing and optimizing them at inference time to solve tasks not represented directly in the training set.

Flexible representationOne function scores complete candidate outputs and can define a probability distribution over them.
Multiple inference proceduresChange a proposed output to reduce its energy, or sample several plausible outputs from the distribution.
CompositionAdd scores for different requirements and search for an output with low total energy.

Representation and inference

We write the scoring function as Eθ(x, y), where x is the input, y is a candidate output, and θ denotes the learned parameters. For robot planning, x specifies a starting state and a goal, and y is a candidate trajectory, a sequence of states and actions. The energy scores how well the trajectory follows the dynamics and reaches the goal. For unconditional generation, we omit x and score the output y alone.

A neural EBM implements the energy as a network that takes the input and candidate output and returns one scalar. Its learned weights are θ. Training changes θ to shape the energy landscape. During inference, we hold x and θ fixed and change the candidate y. The learned energy tells us how well each candidate fits, and an inference procedure uses it to search for a prediction or generate samples.

Representation

An energy function represents compatibility between an input and its possible outputs, with lower energy indicating a better match. We can use this representation directly to search for a prediction, or convert its scores into a probability distribution over outputs:

pθ(y | x) = exp(−Eθ(x, y)) / Zθ(x)

Here, exp(−E) converts energies into positive relative weights. Lower-energy outputs receive larger weights. Dividing by Zθ(x) normalizes these weights over y for the given input x. For continuous outputs, p is a probability density: integrating it over a region gives the probability of an output in that region.

The partition function. Zθ(x) sums exp(−Eθ(x, y)) over all possible discrete outputs y, or integrates it over continuous outputs. When this quantity is finite, dividing by Z makes the probabilities sum or integrate to one. Computing Z exactly is usually too costly in high dimensions. For a fixed input, however, Z is the same for every candidate, so we can compare outputs and follow energy gradients with respect to y without evaluating it.

Learning the energy. Maximum-likelihood training adjusts θ to increase the probability assigned to observed outputs. Although Z is constant when comparing candidates under a fixed model, it changes when we update θ. A common way to account for this change is to sample outputs from the current model, then update the network to lower the energy of observed input-output pairs and raise the energy of these model samples. These updates shift probability toward the training data. Denoising objectives provide another route: add noise to training outputs and learn energy gradients that guide the noisy outputs back toward the data.

When the goal is a prediction, we can also train through optimization, adjusting θ so that minimizing the energy recovers the desired output. This trains the energy to support accurate predictions without requiring it to fit a probability distribution.

Once the energy is learned, optimization searches for a low-energy output. In the probabilistic formulation, the global energy minimum is also the most likely output. Sampling targets the full distribution, representing the range of plausible outputs.

Optimization

ŷ = arg miny Eθ(x, y)

The argmin selects the output with the lowest energy. Under the probability model above, this is the maximum a posteriori (MAP) solution. For continuous outputs, gradient descent searches for it by repeatedly moving y toward lower energy. This search can settle at a local minimum, where nearby changes cannot reduce energy even though a lower minimum may exist elsewhere.

Sampling (Langevin)

yt+1 = yt − η∇yEθ(x, yt) + √(2η) εt

Langevin dynamics is a Markov chain Monte Carlo (MCMC) method for drawing samples through repeated random updates. Each update moves a candidate toward lower energy and adds noise, producing a sequence of candidates called a chain. The noise helps candidates explore beyond a single minimum. With sufficient exploration and a suitable step size, these candidates approximate samples from the probability distribution defined by the energy.

Here, t indexes inference steps and η controls the step size. The gradient ∇yE describes how energy changes as we vary the candidate. Subtracting it moves the candidate downhill, while εt supplies independent standard Gaussian noise. These gradient-based procedures assume continuous outputs and a differentiable energy. For discrete outputs, we can use a discrete search method or represent the choices with continuous variables during inference.

Figure 1. Gradient descent and Langevin sampling

To isolate inference, we use a simple, hand-specified energy over one continuous output z. The default energy has two valleys, representing two plausible outcomes. Starting from the same point, gradient descent moves toward a minimum, while Langevin dynamics explores the distribution. The completed comparison is shown below.

Energy and candidate positions
-1.00.62.23.8-2-1012E(z)zglobal min.
Gradient descentLangevin chainsShared start
Target density and sampled positions
0.00.51.11.6-2-1012Densityz
Target densityCurrent chain positions
3,000 steps
Completed comparison

Descent is at z = 1.06, near a global minimum. 75% of the chains are on the right, compared with 74% of the target probability.

In the default run, gradient descent (magenta) reaches the global minimum on the right, giving the most likely output. The current positions of 128 independent Langevin chains (teal) occupy both valleys, with more chains in the lower-energy valley. Their histogram approximates the normalized target density (black), illustrating how sampling represents uncertainty across plausible outputs. The target density uses a temperature that controls how strongly lower energies are favored.
Explore parameters and numerical details

Increasing b raises the energy hill between the valleys, making crossing harder. Negative a favors the right valley, and positive a favors the left. Temperature T controls how strongly energy differences affect probabilities. Lower T concentrates probability near energy minima, while higher T spreads it more broadly and increases the Langevin noise. Changing a parameter resets the run. Press “Run comparison” to see the result.

E(z) = b(z² − 1)² + az,   pT(z) = exp(−E(z)/T) / ZT

ZT is the partition function at temperature T. The default T is 0.8; the earlier Langevin formula uses T = 1. Both procedures start at z₀ = 1.6. We show 128 independent chains after 3,000 Langevin updates with η = 0.003 and independent standard Gaussian noise: zk+1 = zk − ηE′(zk) + √(2ηT) εk. Here, E′ is the one-dimensional energy gradient. Each update is used without an additional accept-or-reject correction. The random sequence is fixed so that replay follows the same updates. The default starting point lets gradient descent reach the global minimum, but this is not guaranteed in general. For example, changing the tilt to favor the left valley can leave descent in the right local minimum.

All calculations use [−2.4, 2.4], reflecting sampling steps back into this interval and clipping descent steps to its boundaries. A 601-point grid approximates the target normalization and global minimum. The histogram uses 24 bins and a vertical scale fixed across the displayed replay frames. Energy values above the display range are clipped. Incomplete exploration, random variation from finitely many chains, and bias from the nonzero step size can keep the histogram from matching the target. Langevin is not guaranteed to cross the energy hill in a given run.

Why Energy-Based Models?

An EBM separates the representation of a problem from the procedure used to solve it. The same learned model can therefore support different tasks and inference procedures.

Represent complex distributions directly

An energy function scores 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.

Compose different kinds of knowledge

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. This also lets us compose pretrained models by turning their likelihoods or prediction errors into energy terms.

Reuse structure across scale

Energy functions can score overlapping subsets of variables. A relation model can be reused across pairs of objects, while a dynamics model connects successive states in a trajectory. Summing these terms defines a joint energy over the full scene or plan.

Flexible inference-time computation

An energy function leaves the amount of inference computation open. Harder problems can use more optimization steps or search over several candidates in parallel. The learned parameters remain fixed while the search adapts to the problem.

Co-design the landscape and the solver

Training can shape both the solutions an energy prefers and how easily inference reaches them. We can train the energy so that running optimization produces accurate predictions. We can also learn a sequence of landscapes that first guides candidates toward broadly plausible solutions and then refines their details.

Flexible conditional sampling

A joint energy scores a scene together with its relations, such as one object being left of another. We can fix some of these variables and optimize or sample the rest. Changing which variables we fix lets the same model switch between generation and understanding. For example, we can specify relations to generate a scene, or fix the scene and infer its relations.

Properties and connections

EBMs connect probabilistic modeling with optimization, sampling, and inference-time computation.

Generation by sampling or optimization

For probabilistic EBMs, MCMC methods generate samples from the distribution defined by the energy. Energy landscapes can also be trained for generation through optimization. Equilibrium Matching learns a single landscape that guides different random initializations toward different outputs. Its output distribution is determined by the distribution of starting points and the optimization procedure.

Relationship to diffusion models

Diffusion models can be viewed as a sequence of EBMs at different noise levels, with denoising functions providing estimates of their energy gradients. The diffusion score is the gradient of log density with respect to the candidate, a vector pointing toward higher density. An explicit energy-based formulation learns a scalar energy and uses its negative gradient as this score to guide denoising. This connection lets us use energy-based sampling and composition within diffusion models. Our work develops energy-based diffusion samplers and adapts diffusion objectives to learn a single energy landscape for generation through optimization.

Inference for reasoning and planning

Checking a candidate solution can be simpler than constructing one. An energy learns to assess candidates, while inference searches for solutions with low energy. For structured reasoning, separate energies assess different constraints. In planning, dynamics, goals, and constraints score shared trajectory variables, allowing a goal at the end of a plan to shape earlier states and actions.

Composing Energy-Based Models

We can compose EBMs through different operations on their energy functions. Adding energies forms a distribution proportional to the product of the component densities, favoring outputs that satisfy multiple requirements together. For two energies EA and EB, their soft minimum is E = −log(exp(−EA) + exp(−EB)). Exponentiating −E adds their unnormalized densities. After normalization, this gives a mixture, or weighted average, of the component distributions. It preserves outputs favored by either model. These operations let us construct new distributions from learned components and run inference in the resulting landscape, without retraining the component models.

Figure 2. Composing a model with a goal energy

To see what adding energies does, suppose z is a gripper’s final position. The original model favors positions near −1 and +1. A quadratic goal energy penalizes distance from a desired position, changing the distribution when added to the original energy.

Adding a goal energy
03710-2-1012Energyzgoal
Original energy EAAdded goal energy λEBSum EA + λEB
Original and composed densities
0.00.71.42.1-2-1012Densityzgoal
Original densityComposed density

The original model assigns 50% probability to each side. With the right goal, 98% of the composed probability lies on the requested side; the original energy stays fixed.

The goal energy (teal) adds to the original energy (magenta) to form the total energy (black). The density plot shows the resulting shift toward the goal. Select “No added goal” to recover the original distribution. The goal is a soft preference: it penalizes distance from the target while allowing a trade-off with the original model.
E = EA + λEB   ⟹   p(z) ∝ pA(z) pB(z)λ
Explore the goal and its weight

Here, pA and pB are the densities defined by the original and goal energies. Adding the weighted energies multiplies pA by pB raised to λ, favoring positions with low energy under both components. Increasing λ gives the goal more influence. Move the goal between the original density peaks to see how the composed model trades off its original preferences against the added goal.

At T = 1, EA(z) = 1.3(z² − 1)² and EB(z) = (z − g)² / (2 × 0.55²), where g is the goal location. The teal curve is the weighted term λEB that enters the sum. Both plotted densities are normalized on [−2.4, 2.4]. The energy scale is fixed, with high values clipped. Both energies score the same position z, and λ sets their relative influence. A soft goal does not guarantee satisfaction of a constraint. The example does not model dynamics or collision avoidance.

Compositional generalization

A new task can combine familiar relations, motions, and constraints in arrangements never observed during training. Composing models over these components lets inference construct solutions that are globally new while remaining locally familiar. For example, models of short trajectories or pairwise interactions can support longer horizons or more interacting objects.

Computational and compositional limitations

These advantages depend on effective inference. Optimization and MCMC can be slow or sensitive to the initial candidate. Composed models need to agree on the meaning of shared variables, and energy weights must reflect the intended balance between their scores. Their sum may introduce local minima, and a component may become unreliable outside the situations represented in training. Exact normalized probabilities also require the often-intractable partition function.

Historical context

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 retrieve stored patterns from incomplete or corrupted inputs. Boltzmann machines extended this approach to probabilistic models and learned energies from data.

Later work developed more practical ways to learn energy functions. Hinton’s contrastive divergence approximates likelihood training using short sampling runs initialized from observed data. Score matching matches the model’s log-density gradients to those of the data through an objective computable from samples. Noise-contrastive estimation learns by distinguishing observed data from samples drawn from a known noise distribution. Both avoid directly evaluating the partition function. LeCun et al.’s 2006 tutorial presented a general framework for 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. Our subsequent work develops energy functions as reusable representations that can be optimized, conditioned, and composed at inference time.

Publications

Equilibrium Matching generative samples

Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

Runqian Wang, Yilun Du
arXiv 2025
[Project] [Paper] [Code]

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 an equilibrium landscape, EqM connects flow and Energy-Based Models through optimization-based inference.


Compositional energy minimization for reasoning

Generalizable Reasoning through Compositional Energy Minimization

Alexandru Oarga, Yilun Du
NeurIPS 2025
[Project] [Paper] [Code]

Generalization is a key challenge 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. We approach 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 composition lets us add constraints during inference and construct energy landscapes for increasingly difficult problems. To improve sample quality from the composed 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.


Energy-Based Transformer architecture and inference

Inference-time computation techniques, analogous to human System 2 Thinking, have 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). We ask whether these System 2 Thinking approaches can generalize across tasks and emerge solely from unsupervised learning. We find that they can, by learning to verify compatibility between inputs and candidate predictions, then framing prediction 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 up to a 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.


Compositional inverse generative modeling

Compositional Scene Understanding through Inverse Generative Modeling

Yanbo Wang, Justin Dauwels, Yilun Du
ICML 2025
[Project] [Paper] [Code]

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 infer the conditioning parameters of a visual generative model that best fit a given natural image. To enable this procedure to infer scene structure from images substantially different from those seen during training, we further propose to build this visual generative model compositionally from smaller models over pieces of a scene. This procedure infers the objects in a scene and generalizes robustly to test scenes with more objects and new shapes. It also infers global scene factors and generalizes robustly 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.


PhD thesis on composable energy landscapes

Learning Generalizable Systems by Learning Composable Energy Landscapes

Yilun Du
MIT PhD Thesis 2024
[Thesis] [Defense]

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.


Iterative reasoning through energy diffusion

Learning Iterative Reasoning through Energy Diffusion

Yilun Du*, Jiayuan Mao*, Joshua Tenenbaum
ICML 2024
[Project] [Paper] [Code]

We introduce iterative reasoning through energy diffusion (IRED), a framework that formulates reasoning and decision-making tasks as energy-based optimization problems. 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. Two techniques are central to the method: 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.


Compositional image decomposition with diffusion models

Compositional Image Decomposition with Diffusion Models

Jocelin Su*, Nan Liu*, Yanbo Wang*, Joshua B. Tenenbaum, Yilun Du
ICML 2024
[Project] [Paper] [Code]

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 substantially different from those seen during training.


Potential-based diffusion motion planning

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.


Compositional generative modeling framework

Compositional Generative Modeling: A Single Model is Not All You Need

Yilun Du, Leslie Kaelbling
ICML 2024
[Paper]

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. This compositional approach learns distributions more efficiently from data and generalizes to parts of the distribution unseen during training. It also lets us construct new generative models for tasks completely unseen during training. Finally, we show that in many cases, we can discover separate compositional components from data.


Unsupervised Compositional Concepts Discovery with Text-to-Image Generative Models

Nan Liu*, Yilun Du*, Shuang Li*, Joshua B. Tenenbaum, Antonio Torralba
ICCV 2023
[Project] [Paper] [Code]

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. We consider the inverse problem: discovering the generative concepts that represent each image in a collection. 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. These representations give NIIP several capabilities at 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. We find that 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.


Composing Ensembles of Pre-trained Models via Iterative Consensus

Shuang Li*, Yilun Du*, Joshua B. Tenenbaum, Antonio Torralba, Igor Mordatch
(*equal contribution. Shuang Li did experiments on image generation, video question answering, and mathematical reasoning. Yilun Du did all the experiments on robot manipulation.)
ICLR 2023
[Project] [Paper]

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 combining 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.


Compositional Visual Generation with Composable Diffusion Models

Nan Liu*, Shuang Li*, Yilun Du*, Antonio Torralba, and Joshua B. Tenenbaum
(*equal contribution)
ECCV 2022
[Project] [Paper] [Code] [Colab] [HuggingFace Demo]
Press coverage: MIT News, MIT CSAIL News

Large text-guided diffusion models, such as DALLE-2, are able to generate 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 proved difficult for DALLE-2. These results point to the effectiveness of the proposed method in promoting structured generalization for visual generation.


Energy-Based Models for continual learning

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.


Learning Iterative Reasoning through Energy Minimization

Yilun Du, Shuang Li, Joshua B. Tenenbaum, and Igor Mordatch
ICML 2022
[Project] [Paper] [Code]

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 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. This formulation lets us allocate more computation to harder problems with more complex energy landscapes by using a more complex optimization procedure. Our experiments show that this approach solves algorithmic reasoning tasks more accurately and generalizes better in both graph and continuous domains. Finally, we illustrate that our approach can recursively solve algorithmic problems requiring nested reasoning.


Unsupervised compositional energy concepts

Unsupervised Learning of Compositional Energy Concepts

Yilun Du, Shuang Li, Yash Sharma, Joshua B. Tenenbaum, and Igor Mordatch
NeurIPS 2021
[Project] [Paper] [Code]

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.


Learning to Compose Visual Relations

Nan Liu*, Shuang Li*, Yilun Du*, Joshua B. Tenenbaum, and Antonio Torralba
(*equal contribution)
NeurIPS 2021, Spotlight
NeurIPS Workshop on Controllable Generative Modeling 2021, Outstanding Paper Award
Press coverage: MIT News, MIT CSAIL News
[Project] [Paper] [Code]

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.


Improved contrastive divergence training for Energy-Based Models

Improved Contrastive Divergence Training of Energy Based Models

Yilun Du, Shuang Li, Joshua B. Tenenbaum, and Igor Mordatch
ICML 2021
ICLR EBM Workshop 2021, Oral
[Project] [Paper] [Code]

We present tools to improve contrastive divergence training of EBMs. First, we identify a neglected term in the objective and present a loss function to mitigate it. We use data augmentation to improve MCMC mixing during training and a multiscale architecture to improve generative performance. These techniques improve both generation and out-of-distribution detection.


Compositional visual generation with Energy-Based Models

Compositional Visual Generation and Inference with Energy Based Models

Yilun Du, Shuang Li, and Igor Mordatch
NeurIPS 2020, Spotlight
[Project] [Paper] [Code]

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.


Atomic-level protein conformation energy model

Energy-Based Models for Atomic-Resolution Protein Conformations

Yilun Du, Joshua Meier, Jerry Ma, Rob Fergus, and Alexander Rives
ICLR 2020, Spotlight
[Paper] [Code]

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 with Energy-Based Models

Model Based Planning with Energy Based Models

Yilun Du, Toru Lin, and Igor Mordatch
CORL 2019
ICML MBRL Workshop 2019, Oral
[Paper] [Code]

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.


Implicit Generation and Generalization in Energy-Based Models

Yilun Du and Igor Mordatch
NeurIPS 2019, Spotlight
[OpenAI Blog] [Paper] [Code]

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.