Generalizable Reasoning through Compositional Energy Minimization

TL;DR

This work tackles the generalization gap in reasoning by reframing problems as energy minimization, where a learnable energy $E_{ heta}(\bm{x},\bm{y})$ assigns lower energy to valid solutions and higher energy to invalid ones. By decomposing complex problems into tractable subproblems and composing their energy landscapes, the method enables solving larger, harder tasks than those seen during training. A key contribution is Parallel Energy Minimization (PEM), a particle-based sampling mechanism that resamples and perturbs multiple candidate solutions to effectively explore composed energy landscapes and avoid local minima. The paper also introduces a diffusion-based training objective for EBMs, paired with a contrastive loss to shape the energy landscape, and demonstrates superior generalization across N-Queens, SAT, graph coloring, and crosswords compared with strong baselines, with tunable computational budgets enabling performance trade-offs. Overall, the approach promises improved generalization for complex reasoning tasks by leveraging compositional energy landscapes and parallelized inference.

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 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. Project website can be found at: https://alexoarga.github.io/compositional_reasoning/

Figures (23)

• Figure 1: Compositional Generalizable Reasoning. We formulate reasoning as optimization problem with inputs $\bm{x}$ and solutions $\bm{y}$. By combining multiple optimization objectives, we can generalize to larger problem instances (bottom) than those seen during training (top). This enables us to solve a complex instance of N-queens (left) or a more complex instance of graph coloring (right).
• Figure 1: 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. We sampled 100 8-queens solutions. IR: Iterative Refinement, BP: Belief Propagation, TS: Tree Search, S: Sampling, GS: Guided Sampling,
• Figure 2: Parallel Energy Minimization (PEM)
• Figure 2: Number of Particles vs Correct Instances. We sampled 100 solutions from the 8-queens problem. Increasing the number of particles with PEM significantly improves the number of correct instances.
• Figure 3: N-Queens Problem Composition. To compose a row model to solve the N-queens problem, we add the energy of each row $i$ ($E_{\theta}^{ri}$), each column $j$ ($E_{\theta}^{cj}$), each diagonal $k$ ($E_{\theta}^{dk}$) of the chessboard. We then sample from the resulting energy function $E_{\theta}^{total}$ to generate valid solutions.