Generalizable Reasoning through Compositional
Energy Minimization

Alexandru Oarga, Yilun Du
PDF Code

TL;DR

Compositional energy-based approach outperforms existing methods by combining subproblem energy landscapes, enabling better generalization to larger and more complex reasoning problems. Parallel Energy Minimization is applied to improve solution quality when sampling from these constructed energy landscapes.

Train
Train example 1
Train example 2
Train example 2
Test
Test example 1

N-Queens

Test example 2

Graph Coloring

Test example 2

Crosswords

Abstract

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 solu- tion 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 Mini- mization (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.

Compositional Reasoning

We propose a framework that generalizes to complex problems by decomposing them into simpler, tractable subproblems with their own energy-based models (EBMs). A global solution is then obtained by jointly minimizing the sum of the subproblem energy functions, effectively combining their solutions to solve much larger problems than those seen during training.


N-Queens
Composition of Energy Landscapes. In this example, the energy function for a single edge (left) assigns lower energy to valid color assignments and higher energy to invalid ones. By composing multiple such edge-level energy functions, we can solve larger graph coloring problems (right). The composed energy landscape of the right graph is obtained by summing the individual edge energies, and the minima of the composed landscape correspond to a valid solution.

Parallel Energy Minimization

To sample from complex energy landscapes, obtained after combining multiple subproblem energy functions, we propose a novel optimization framework called Parallel Energy Minimization (PEM). PEM maintains and resamples multiple particles to effectively explore composed energy landscapes. By using a larger number of particles, PEM enhances exploration of the composed energy landscape.


N-Queens

N-Queens Problem. We illustrate the generation process for the N-Queens problem with eight particles. By using multiple particles with PEM, we effectively generate valid solutions to the problem.

N-Queens

Graph Coloring Problem. We illustrate the generation process of a graph coloring problem with eight particles. By using multiple particles with PEM, we are able to find a valid coloring of the graph.

N-Queens

Crosswords. We illustrate the generation process of a solution to a given crossword puzzle using eight particles. Hints for each word are omitted for clarity.

Performance

Our approch can generate correct solutions to the 8-queens problem with a higher accuracy than existing state-of-the-art methods. We report below the performance of our method compared with baselines and the effect of varying the number of particles.


8-Queens Problem Evaluation. We compare the performance against state-of-the-art combinatorial optimization models on the 8-queens solution generation task. All the models were trained with 1 single instance of the 8-queens problem.

Number of Particles vs Correct Instances Across Problem Difficulties. We sampled 100 solutions each from N-queens problems of increasing difficulty (7 to 10 queens). In all cases, a higher number of particles results in more correct instances.


Our method outperforms the previous state-of-the-art neural combinatorial and SAT solvers, and is able to find a larger number of correct instances of the problem in the 3-SAT solving. It consistently outperforms other approaches both within the training distribution and when applied to larger, out-of-distribution instances. We report below the comparison in both distributions.


3-SAT Problem Evaluation. We compare the performance against the state-of-the-art combinatorial optimization models and neural SAT solvers on the 3-SAT task. Models are evaluated on a distribution similar to the training distribution and a larger distribution. Similar distribution has 100 instances with 20 variables and 91 clauses, while larger distribution has 100 instances with 50 variables and 218 clauses. Our approach outperforms existing methods.


We compare our approach with GNN-based methods for graph coloring and show that it is able to generalize to larger graphs and different graph distributions while maintaining strong performance across different scales. On average, our method produces solutions with fewer conflicting edges. The comparison with other methods is presented below.


Graph Coloring Problem Evaluation. We compare the performance against canonical GNNs and GNN-based methods for graph coloring on different random graph distributions. Performance is measured as the number of conflicting edges, with lower indicating better. For each distribution, we report the average over five instances. Our approach outperforms existing methods on most instances and generalizes better to larger and denser graphs.

Related Projects

Check out a list of our related papers on compositional generation and energy based models. A full list can be found here!

Learning Iterative Reasoning through Energy Diffusion
We propose energy optimization as an approach to add iterative reasoning into neural network. We illustrate how this procedure enables generalization to harder instances of problems unseen at training time on both continuous, discrete and image processing tasks.

Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC
We propose new samplers, inspired by MCMC, to 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.

BibTeX

@article{oarga2025generalizable,
  title={Generalizable Reasoning through Compositional Energy Minimization},
  author={Oarga, Alexandru and Du, Yilun},
  journal={arXiv preprint arXiv:XXXX},
  year={2025}
}