Learning Iterative Reasoning through Energy Diffusion

A wide variety of reasoning and decision making problems can be formulated as an optimization problem. Traditionally, researchers have been focused on designing various domain-specific algorithms for solving different problems, typically with search, gradient-based optimization, or other forms of iterative computation, and also integrating machine learning to help. In this work, we take a different approach of formulating various kinds of reasoning and decision-making problems as an optimization process over a learned energy-based model (EBM): $E_{\theta}({\bm{x}},{\bm{y}}):\mathbb{R}^{O}\times\mathbb{R}^{M}\rightarrow\mathbb{R}$. Under this formulation, the final prediction problem can be cast as finding the solution ${\bm{y}}$ according to:

One can use gradient descent to find such solutions:

where $\lambda$ is the step size for optimization and the initial prediction ${\bm{y}}^{0}$ is initialized from a fixed noise distribution (i.e., Gaussian throughout the paper). The final output of ${\bm{y}}^{T}$ is obtained after $T$ steps of optimization.

In earlier work using a similar formulation Du et al. (2022), such EBM $E_{\theta}({\bm{x}},{\bm{y}})$ is trained by differentiating through the $T$ steps of optimization and minimizing the MSE with the ground truth label ${\bm{y}}$ $\mathcal{L}_{\text{Opt}}(\theta)=\|{\bm{y}}_{i}^{T}-{\bm{y}}_{i}\|^{2}.$ This approach requires the forward and backward computation of K steps of optimization at training, which makes it slow and unstable. Furthermore, because the EBM $E_{\theta}$ may have a complex optimization landscape¹¹1For example, the 3-SAT problem exhibits steep energy minima surrounded by flat energy landscapes., robustly finding solutions to Equation 2 is fundamentally hard.

As a general solution to stable training and better test-time optimization, instead of directly learning $E_{\theta}({\bm{x}},{\bm{y}})$, at a high-level, we propose to learn a sequence of annealed energy functions $E^{k}_{\theta}$ ($k=0,1,\cdots,K$), and supervise the EBM learning with the gradient of the energy function:

During training time, we obtain the ground truth for ${\bm{y}}$ by generating a noise-corrupted label ${\bm{y}}+\epsilon$, following a schedule of noise corruptions. By supervising on the gradient, our approach is substantially faster and more stable than earlier works using plain EBMs (Du et al., 2022) as it only supervises training of a single step of the optimization.

Our key idea to mitigate the hard optimization problem of Equation 2 is to use simulated annealing — where smoother energy landscapes are first optimized before optimizing sharper ones afterward, as illustrated in Figure 2.

Similar to diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020), we propose to optimize and learn an annealed sequence of energy landscapes, with earlier energy landscapes being smoother to optimize and the latter ones more difficult. Given a ground truth label ${\bm{y}}$, we learn a sequence of $K$ energy functions²²2Empirically, in our experiments, we found that setting $K=10$ was sufficient across all the domains we considered. With a total of 10 energy landscapes, we can smoothly transition from a Gaussian-like landscape (with $K=10$) to a sharp and discontinuous landscape (with $K=1$). $E^{k}_{\theta}({\bm{x}},{\bm{y}})$ over the ground truth label distribution $p({\bm{y}}^{*}|x)$, where each energy function is learned to represent an EBM distribution

over a sequence of noise scales $\sigma_{k}$. Here, $\mathcal{N}(\cdot|\mu,\sigma)$ is the Gaussian density function. Larger values of $\sigma_{k}$ correspond to smoother energy landscapes while smaller values lead to sharper landscapes, with the energy minima of landscape $k$ corresponding to $\sqrt{1-\sigma_{k}^{2}}{\bm{y}}^{*}$ (which can be scaled by $\tfrac{1}{\sqrt{1-\sigma_{k}^{2}}}$ to obtain the ground truth prediction ${\bm{y}}^{*}$).

We can directly learn each energy landscape by supervising the gradient of energy function to denoise the corrupted ground truth label ${\bm{y}}^{*}$ from the dataset

where $\epsilon\sim\mathcal{N}(0,1)$. Given a set of $K$ different learned energy landscapes, we can initialize a data sample from Gaussian noise and sequentially run $T$ steps of optimization following Equation 2 over each energy landscape $k$ (starting with high noise levels and progressing to lower noise levels). The optimization result in the previous energy landscape is used to initialize optimization in the next landscape, after scaled by the appropriate scaling factor $\tfrac{\sqrt{1-\sigma_{k}^{2}}}{\sqrt{1-\sigma_{k-1}^{2}}}$.

In the denoising training objective Equation 5, while the gradient of the energy landscape is locally trained to restore the ground truth label ${\bm{y}}$, it is not necessarily the case that the overall global energy minima $\operatorname*{arg\,min}_{{\bm{y}}}E_{\theta}({\bm{x}},{\bm{y}},k)$ corresponds to the ground truth label $\sqrt{1-\sigma_{k}^{2}}{\bm{y}}^{*}$.

To enforce that the global energy minima of each of the $k$ energy landscapes corresponds to the ground truth energy minima, we further propose a contrastive loss, where we construct a set of negative label ${\bm{y}}^{-}$ (formed by noise corrupting the ground truth label ${\bm{y}}^{*}$). Given an energy $E^{+}=E_{\theta}({\bm{x}},\sqrt{1-\sigma_{k}^{2}}{\bm{y}}^{*}+\sigma_{k}\epsilon;k)$ of the ground truth label ${\bm{y}}^{*}$ and an energy $E_{-}=E_{\theta}({\bm{x}},\sqrt{1-\sigma_{k}^{2}}{\bm{y}}^{-}+\sigma_{k}\epsilon;k)$ of the negative label ${\bm{y}}^{-}$, $\mathcal{L}_{\text{Contrast}}(\theta)=-\log\left(\frac{e^{-E^{+}}}{e^{-E^{+}}+e^{-E^{-}}}\right)$. To reduce the variance of the contrastive loss, we use the same sampled noise value $\epsilon$ for both ${\bm{y}}$ and ${\bm{y}}^{-}$.

We provide the overall pseudocode for training IRED in Algorithm 1 and executing algorithmic reasoning with IRED in Algorithm 2. We use a cosine beta schedule to train annealed energy landscapes and use a total of 10 energy landscapes (we empirically found that more landscapes did not lead to improved performance). At inference time, we can vary the number of optimization steps $T$ for each energy landscape to make trade-offs between performances and inference speed.

In principle, when the solution is not well-defined, it is possible to use IRED to model multi-modal distributions, similar to how diffusion models have been proven effective in modeling multi-modal image distributions. Depending on the particular use case, one may also add additional inference-time constraints (e.g., by composing the learned IRED energy function with other energy functions) to select favorable solutions.

Setup.  We first evaluate IRED on a set of continuous algorithmic reasoning tasks from  Du et al. (2022). We consider three matrix operations on $20\times 20$ matrices, which are encoded 400-dimensional vectors:

1. 1.

   Addition: We first evaluate neural networks in their ability to add matrices together (element-wise). We also evaluate neural network on harder variants of the addition problems at test time by feeding input vectors with larger magnitudes.

2. 2.

   Matrix Completion: Next, we evaluate neural networks on their ability to do low-rank matrix completion. We mask out $50\%$ of the entries of a low-rank input matrix constructed two separate rank 10 matrices $U$ and $V$, and train networks to reconstruct the original input matrix. We construct harder variants of the matrix completion problem at test time by increasing the magnitude of values in $U$ and $V$.

3. 3.

   Matrix Inverse: Finally, we evaluate neural networks on their ability to compute matrix inverses. We construct harder matrix inverse problems by considering less well-conditioned input matrices.

We report the underlying mean-squared error (MSE) between the predictions and the associated ground truth outputs on test problem instances. To more effectively generate negative samples for IRED in this domain, we first noise-corrupt ground truth labels and then run two steps of gradient optimization on the energy landscape to form negative samples. Details can be found in Appendix A.

Baselines.  We compare our approach to a set of iterative reasoning baselines found in  (Du et al., 2022): (Feedforward): an iterative reasoning approach where the same MLP is repeatedly applied, (Recurrent): an iterative reasoning approach where the recurrent network is repeatedly applied, and (Programmatic): an iterative reasoning approach which repeatedly applies a learned programmatic module (Banino et al., 2021). We further compare with the IREM method  (Du et al., 2022), as well as using a denoising diffusion model directly to solve continuous tasks. All methods use identical architectures (with small differences due to recurrent layers or timestep conditioning).

Quantitative Results.  We compare IRED with baselines across settings in Table 1. Similar to IREM, IRED is able to nearly perfectly solve the task of the addition, as well as generalize to larger addition matrices. On other tasks, IRED outperforms IREM and also generalizes better to harder problems. Furthermore, our approach substantially outperforms directly using a diffusion process to predict solutions, which lacks the iterative energy minimization procedure that explicitly learns the task constraints.

Qualitative Visualization.  We provide a qualitative visualization of the error map of the optimized solution on the matrix inverse task at each different learned energy landscape in Figure 3. The error of optimized solutions at different energy landscapes decreases over time.

Energy Landscape.  We visualize the learned energy landscape in Figure 4 as a function of the distance of an input label from the ground truth label. In early energy landscapes, the difference between energy values of solutions close and far from the ground truth solution is low, and therefore the energy landscape is relatively flat. At later landscapes, the energy value increases substantially as the input solution deviates from the ground truth solution.

Performance with Increased Computation.  We analyze the performance on the matrix inverse task as a factor of an increased number of computational steps in Table 2. We find that running additional steps of optimization slightly improves performance on in-distribution tasks and substantially improves performance on harder problems.

Ablation.  We ablate each component of IRED in Table 3. In the first two rows of Table 3, we compare our gradient-descent-based optimization with a noisy optimization procedure corresponding to the diffusion reverse process for each energy landscape. In the third row, we compare the difference between running multiple steps of optimization as opposed to a single energy optimization step. Finally, We then consider the effect of contrastively shaping the energy landscape. All components lead to improved performance.

Setup.  The second group of tasks evaluates IRED on its reasoning in discrete spaces (i.e., values are all binary or one-hot categorical). We run evaluations on two tasks: Sudoku solving and graph connectivity reasoning.

1. 1.

   Sudoku: In the Sudoku game, the model is given a partially filled Sudoku board, with 0’s filled-in entries that are currently unknown. The task is to predict a valid solution that jointly satisfies the Sodoku rules and that is consistent with the given numbers. We use the dataset from SAT-Net (Wang et al., 2019) as the training and standard test dataset. In SAT-Net, the number of given numbers is within the range of $[31,42]$. Our harder dataset is from RRN (Palm et al., 2018), which is a different Sudoku dataset where the number of given numbers is within $[17,34]$. We will show that our system, being trained on simpler Sudoku games with fewer blank entries, generalizes to harder instances.

2. 2.

   Connectivity: In the graph connectivity task, the model is given the adjacency matrix of a graph (1 if there is an edge directly connecting two nodes). The task is to predict the connectivity matrix of the graph (1 if there is a path connecting two nodes). In literature (Dong et al., 2019), this task is an example that requires a dynamic number of reasoning “steps” (depending on the treewidth of the graph). Therefore, prior papers primarily focus on computing connectivity between nodes within $k$-steps away. Our training and standard test sets contain graphs with at most 12 nodes and our harder dataset contains graphs with 18 nodes. Since we do not limit the path lengths between nodes (in contrast to Dong et al. (2019)), on the training set, the maximum distance between two connected nodes is 9, and in the harder test set, the maximum distance is 16.

Quantitative Results.  We compare IRED with both IREM and diffusion baselines. On the Sudoku task, we further compare with the domain-specific SAT-Net method. In Table 4, we find that our approach substantially outperforms all baselines across both evaluated discrete-space reasoning settings. In Sudoku, our approach generalizes substantially better than the SAT-Net model and the RRN model to the harder dataset consisting of fewer given Sudoku elements and is capable of obtaining an accuracy of roughly 62.1% compared to an accuracy of 3.2% obtained by SAT-Net and 28.6% by RRN. For cases in which our approach fails, we found that our approach sometimes erroneously assigns low energy to partially accurate answers. For example, there can be a Sudoku board that is not fully valid, but the model assigns lower energy to it than to the ground truth board.

Qualitative Results.  We qualitatively visualize intermediate optimized samples across energy landscapes on Sudoku in Figure 6. Optimized boards are increasingly accurate in later landscapes.

Performance with Computation.  In Figure 6, we assess the impact of the number of optimization steps at each energy landscape on the performance in Sudoku on both the test and harder datasets. We find that performance substantially improves on the harder dataset with an increased number of optimization steps with more modest improvement on the test datasets. By formulating reasoning as energy optimization, we can adaptively change the number of optimization steps dependent on difficulty of task, enabling us to generalize substantially better on harder tasks.

Extension to Visual Sudoku.  IRED can also be extended to deal with other input formats, such as images. To illustrate this, we conducte a new experiment on the Visual Sudoku dataset (Wang et al., 2019), where the board is not represented by one-hot vectors but now consists of MNIST digits written on a grid. We use a CNN to encode the image and fuse the image embeddings with the noisy answer to predict energy values. Shown in Table 4, we observed a similar performance advantage of our model compared to the baseline.

Ablation.  In Table 5, we ablate each component of IRED on the Sudoku task. Similar to the continuous setting, we find that each component of IRED, gradient based optimization, multiple steps of optimization, and contrastive energy shaping all lead to improved performance. While performance is modestly improved on the test dataset, it is substantially improved on the harder generalization dataset with each added component.

Setup.  In this section, we evaluate IRED on a basic decision-making problem of finding the shortest path in a graph. In this task, the input to the model is the adjacency matrix of a directed graph, together with two additional node embeddings indicating the start and the goal node of the path-finding problem. The task is to predict a sequence of actions corresponding to the plan. Concretely, the output is a matrix of size $[T,N]$, where $T$ is the number of planning steps and $N$ is the number of nodes in the graph. Each entry $(t,i)$ has a value of 1 if the $t$-th step of the shortest path is at node $i$ and has a value of 0 otherwise. For all models, including ours, we use a spatial-temporal graph convolution network (STGCN; Yan et al., 2018) to encode the adjacency matrix and the prediction and predict energy function values or score functions. In short, the STGCN has multiple layers. At each layer, for each node, it fuses all features of nodes that are connected to it in the graph and from the current timestep or adjacent timesteps. Just like standard graph convolutional networks, it uses a sum-pooling mechanism to aggregate all embeddings from adjacent nodes. Since, in practice, such planning models are generally evaluated with closed-loop execution, here we only evaluate the success rate that the first action produced by the model shortens the distance between the current node and the goal node. This is the same planning and execution strategy as DiffusionPolicy (Chi et al., 2023).

Quantitative Results.  We compare IRED with all baselines across settings in Table 6. IRED outperforms both baselines, especially the diffusion model, by a large margin. Both methods based on energy based formulation (IREM and IRED) perform well on this task, validating the hypothesis that learning energy functions is an effective method for encoding planning problems such as the edge constraints in the path-finding task. Finally, since all methods uses the same STGCN encoder, there is no significant performance drop on generalization to the harder dataset.

Qualitative Results.  Similarly to other tasks, we can also visualize the generated solutions by the model across different landscapes. Figure 7 visualizes the prediction of the first node to move to by the planning model. We normalize the prediction scores to 0 to 1. The darker the color, the higher the score. As can be seen in the figure, our IRED model is capable of gradually finding the immediate next action to take: at step 3 it excluded node 7, and the score gradually concentrates on node 8.

Ablations.  We ablate each component of IRED in Table 7. We found that running multiple steps of optimization and shaping the energy landscape both lead to improved performance on both the same-difficulty test cases and harder difficulty ones.