Generalizable Reasoning through Compositional Energy Minimization
Authors: Alexandru Oarga
Yilun Du
Venue: NeurIPS 2025 spotlight ()
Keywords: Reasoning Generalization
Compositional Inference
Test-Time Computation
Test-Time Inference
Combinatorical Optimization
Primary area: deep_learning
TL;DR: We propose a compositional energy-based approach with a new sampling method (PEM) that improves generalization on complex reasoning tasks.

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/

Generalizable Reasoning through Compositional Energy Minimization
Alexandru Oarga University of Barcelona
Yilun Du Harvard University
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 learn useful heuristics from data, 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 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 approach on wide set of reasoning problems. Our method outperforms existing state-of-the of-te-of-theaar methods, demonstrating its ability generalize larger and more complex problems. Project website be found at: https::/ http/////ea//aaaa////////////aaa github. io/ /to/ooopsii ti onall reasoning/
Introduction
Being able to solve complex reasoning problems, such as logical reasoning, combinatorial puzzles and symbolic manipulation, is one of the key challenges in machine learning. This is particularly interesting because requires models to go beyond pattern recognition. For a model to successfully perform reasoning tasks, it is expected to be able to generalize unseen distributions during test time. That is, they are expected to solve not only problems similar to those encountered during training, but also to be able to generalize novel conditions and distributions [9].
The standard paradigm in machine learning for solving reasoning tasks is to train models end-to-end to map inputs outputs. During training, models are exposed to large number of solutions and learn statistical heuristics that allow them to solve similar problems. This contrasts with human reasoning, where we first learn the rules and constraints governing a problem, and then apply them in compositional manner to arrive at a solution. Notably, humans are able to solve such problems exact solution before. Moreover, when prompted with harder tasks, can
without having invest more time, effectively searching for a solution, rather than relying on heuristics [1. 37].
In this work, we present an approach to reasoning where the overall process is cast as an optimization problem [22]. Specifically, we learn an energy function Ea y) across all possible solutions y of the problem, where are the given conditions of the problem. This energy function is learned such that valid solutions are assigned lower energy, while invalid solutions receive higher energy.
Corresponding author: ydudseas... waduseaas.harvard.coo hww.ard..dcc....... edu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Train
Evaluation
Figure I: Compositional Generalizable Reasoning. We formulate reasoning as optimization problem with inputs æ and solutions 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 complex instance of N-queens (left) or more complex instance of graph coloring (right).
Reasoning corresponds to minimizing the energy function to find low-energy solutions. Harder problems can be solved by spending more time reasoning and minimizing the optimization objective.
To solve more complex problems than those seen during training (see Figure 1), we can jointly reason and minimize the sum of several optimization objectives at the same time. For instance, in logical reasoning, we can learn the energy landscape of individual clauses, and then optimize over multiple clauses simultaneouslyy. This composition of energy objectives enables the model to find assignments that satisfy all clauses simultaneously. enabling to solve larger, more complex problems. However, combining multiple objectives makes the landscape increasingly complex and introduces local minima, making hard to optimize. To address this, we propose a parallel strategy where we system of particles for optimization. In this setup, leverage the energy function as a resampling mechanism to improve the quality of the samples and avoid local minima that attract particles. This approach improves exploration and ultimately makes optimization more effective.
We illustrate the applicability of our approach across a set difficult reasoning problems, including the N-Queens, 3-SAT and the Graph Coloring. We compare against domain-specifie state of the art combinatorial optimization models, and show that our approach outperforms them in terms of solution quality and generalization larger and more complex problems. We further show that, by adjusting the computational budget, we can adapt the model to solve more complex problems. Finally, ablation studies show that training strategy leads improved results and that our sampling strategy able to produce better quality solutions than existing samplers.
Overall, the contributions of this work are threefold. First, propose compositional approach to reasoning generalization, where we combine energy landscapes during inference to solve more complex problems. Second, we introduce new sampling strategy, Parallel Energy Minimization (PEM), a particle-baser optimization strategy that enables us to effectively optimize composed energy functions to solve hard reasoning tasks. Finally, we illustrate the efficacy of our approach empirically across a wide set of reasoning tasks, outperforming many existing combinatorial optimization approaches in generalization.
2 Related work
Reasoning as Optimization. Reasoning includes multiple cognitive processes, such as logical infer- ence, decision-making, planning, and scheduling. Many of these can be formulated as optimization problems, where the goal is find variable assignments that minimize an objective function under certain constraints. Prior works have integrated logical reasoning into neural networks through differentiable relaxations of SAT [26] and MAXSAT solvers [64], differentiable theorem proving [53, 47], probabilistic logic [45], diffecrentiblle logic rules [59], and gradient- based methods for discrete distributions [49]. Other approaches incorporate continuous optimization directly into models via differrentialee convex [2], quadratic [5] or integer solvers [61]. These methods, however, often target specific domains or rely on strong assumptions.
Another line formulates reasoning using general- purpose optimization frameworks. For instance, [54, 15] simulate physical dynamics using energy minimization. Latent space optimization has been used in variational methods for molecule generation [28] and the Traveling Salesman Problem (TSP) [35]. More related to our work, [22] learns energy functions backpropagating through optimization steps with diffusion based losses [23]. Nonethcless, most of these approaches adopt end-to-

---

end methods to learn reasoning tasks, limiting their generalization ability. We instead propose compositional strategy, combining energy landscapes learned on subproblems to tackle larger tasks.
Finally, recent approaches focus on learning-baa based methods for Combinatorial Optimization (CO), which aim reduce the computational cost by generating near-optimal solutions. Graph Neural Networks (GNNs) are the standard in this domain due to their ability to represent variable-constraint relations. Recent works have employed GNNs directly predict solutions [10, 36, 55J, learning discrete diffusion over graphs [60, 42], using reinforcement learning [8, 3] or learning Markov processes [72]. However, it is well known that GNNs struggle out of distribution [67, 27] and require large, diverse datasets. Our approach leverages the compositional nature of reasoning problems by producing more generalizable energy landscapes combining multiple energy objectives.
Reasoning as Iterative Computation. Some strategies use neural networks to iteratively refine solutions to reasoning problems. This motivation is drawn from optimization solvers, which operate with iterative updates. Within this category, we can identify three main directions: (1) methods that incorporate explicit program representations [30, 52, 12, 70, 48], (2) works based on recurrent neural networks [29, 38, 13, 69, 18, 56, 70], and (3) techniques that approximate solutions via iterative refinement [58, 7, 43, 42, 64, 44]. In our work, we cast reasoning problems as optimization problems, hence we use optimization algorithms as refinement steps for solution search.
Energy- -Baed Models and Diffusion Models. Our work is closely related to Energy -Bagee Models (EBMs) [32, 40, 24]. Most of the work in this field has focused on learning probabilistic models over data [24, 50, 21, 6, 66, 17]. In contrast, train EBM for solving reasoning tasks by performing optimization over the learned energy landscape.
3 Method
3.1 Reasoning as Energy Minimization
Let D = {X, } be dataset of reasoning problems with inputs ORO and solutions y E RR We wish to find an operator f'()) that can generalize test problems ((xx' where x' RRO', is potentially larger and more complex than æ. Let Ex y) TROO x RM + R, be an EBM defined across all possible solutions y given æ. such that ground-truth solutions y are assigned lower energy. Finding solution to the reasoning problem corresponds to finding an assignment ý such that:
ÿ arg min Es y)
(1)
To find the solution ÿ. one use gradient descent:
y' = --==-"--,VV,aa Eo,,,
(2)
where Ả is the step size, and yd is the initial solution drawn from fixed noise distribution (c.g. Gaussian). The resulting solution y is found after T iterations of the above update.
Diffusion Energy-Based Models. The effective training of EBMs is a challenging task, and currently. many approaches exist in the literature for this purpose [24, 11, 20]. Previous works trained EBMs by backpropagating the gradient through T generative steps [22]. However, this could lead to instabilities in the training and high computational cost of backpropagation.
In this work, we propose instead to use the denoising diffusion training objective introduced in [20]. Specifically, we train the gradient of the EBM to match the noise distribution at each timestep t. Formally, given a truth label y from the dataset, and a gaussian corrupted label y", where 3² = V1 o.y + orE and (0, 1) we can define the diffusion objective as:
LMSE (0) Ey ((0.))) [||e V, Eoy",, (1)
with Ea(y"," .t) being explicitly defined scalar function and ar a sequence of fixed noise schedules.
This formulation allows us supervise the gradient of the energy function at each optimization step f, avoiding the need to backpropagate through sequence of I' steps. As a result, we learn an energy gradient that transforms noisy input into the target distribution, through a series of optimization steps. To generate outputs, we can then use, for example, the update rule given in (2).
Empirically, in this work we use the same function as in the original paper, this is, (e t) = [1s0 2) where so (Iz, t) is vector-output neural network.
3

---

Shaping the Energy landscape. The training objective presented in Eq. (7) does not guarantee that the target label y is assigned the energy minima of the energy landscape. In this work, enforce that the energy minima align to the ground-truth label, and to enhance regions of the landscape not covered by the liffuusin-based training, we follow the approach of [23], and introduce an additional contrastive loss function to shape the energy landscape.
This contrastive loss guides the energy function by comparing noise-corrupted labels of given pairs of positive and negative samples. Formally, the objective at step & is formulated as:
exs
Lcu(0) = - log
c + E
where E+ = Eo
Eo .t), with ØJ and ÿ being positive and negative noise corrupted samples respectively, this is, ÿ v1 o,y+ ore and ÿ ==11 o:y + dyE.
t) and E =
3.2 Compositional Reasoning
We wish to construct a reasoning framework that can generalize to complex problems that are much harder than those seen at training time, consisting of a significantly greater number of constraints. To construct an effective energy function to tackle such problems, we propose to decompose the energy function into smaller ones that are defined over tractable subproblems that have been seen before. These subproblems are then simpler to handle and represent with energy functions compared to trying to solve the original problem. In particular, we propose decomposing a full reasoning problem into simpler subproblems x {.1...., INY. such that finding solution yı to each subproblem 2; solves the original problem æ.
Given this decomposition, let E (æ: y) be an EBM of the k-th subproblem, where yı is assigned the lowest energy when is solution subproblem æg. A complete solution ý the original problem æe is obtained by solving all subproblems simultaneously, formally optimizing the composition of each energy function:
ÿ arg min
Σ
F (at ys)
(5)
where ý can be found as in (2). We illustrate how to effectively optimize these objectives next.
3.3 Improving Sampling with Parallel Energy Minimization
Optimization over EBMs can be done through approximate methods such as Markov Chain Monte Carlo (MCMC). MCMC simulates a Markov chain, starting from initial state yo, drawn from a noise distribution, with subsequent samples generated from a transition distribution. A common approach to MCMC sampling in EBMs is Unadjusted Langevin Dynamics (ULA) [24, 50], which is defined as:
y'= XV (o, y v24g, {~N((011)
(6)
where is the step size of the optimization method. This essentially corresponds to performing gradient descent on the energy function with some added noise.
However, such a noisy optimization procedure will often become stuck in local minima, which are especially prevalent in composed energy landscapes such as Equation 5. In MCMC, a common technique more effectively sample from such difficult probability distributions is Sequential Monte Carlo (SMC) [19], where a set of parallel particles is maintained and evolved over iterations of sampling to help prevent premature convergence to local minima.
Inspired by this insight, we propose Parallel Energy Minimization (PEM), a parallel optimization procedure for optimizing composed energy landscapes. We initialize and optimize parallel set of P particles across the T steps of optimization as presented in Algorithm and illustrated in Figure 2. At each optimization step, PEM resamples particles based on their energy values, allowing local minima to be discarded. To further help particles cover the entire landscape of solutions, we add predefined amount of noise each particle each time they are resampled.

---

Algorithm Parallel Energy Minimization (PEM)
Input: T optimization steps, P particles
Given set of particles [v]
[v]")) ) " ())
oo (5)
for timestep t in ,
Importance evaluation
D
g
(o(t-1)
w <- (oftmax(-Eo(s)
Selection
Resample (C)) based on weights (*)
((-1))
D Resampling
7d)) y8th + orE E~ N~0,0
Figure PEM Sampling. At timestep t, particles y( first resampled using weights w(t) derived from E f). Next, scheduled Gaussian noise is added to obtain new set of resampled particles. Finally. the par- ticles of the next timestep t- are generated optimizing gradient the energy function time t.
Optimize solutions with gradient
(1-1) ((t) Eg
end for return (0)
3.4 Refinement of the Energy Landscape
Composing multiple energy objectives together can effectively generate complex energy landscapes suitable for solving larger problems. However, the number of objectives increases, the energy landscape becomes increasingly complex, which can lead to inaccuracies in the overall energy function. In particular, minima might appear in the function that incorrectly assigns lower energy to invalid solutions.
To mitigate this issue, we propose refinement strategy for the composed landscape. Given set of N energy functions {Eğ (ay, yx)) > each trained on subproblem æ", we refine the composed energy function using ground-truth solutions y. The resulting training objective is defined as:
MSSE (0) We.. V.V Σ Ey(,",
having 3" vI oiy ore and AN 1). This refinement helps align energy minima with valid solutions, correcting inaccuracies and improving robustness.
4 Experiments
4.1 N-Queens Problem
Setup. The N-queens problem involves placing N queens on an N N chessboard such that two queens threaten each other, meaning no two queens can be placed in the same row, column or diagonal. We evaluate how well different methods can generate valid solutions to the problem. During training we use only one single instance of the N-queens problem for a given value N. In our approach, we use the rows of a single instance to train, and then compose this model row-wise, folumn-wisc, wise, and diagonal-wid wise to form 2D chessboard. That is, we train model to generate a valid row and then reuse it simultaneously for rows, columns and diagonals (see Figure 3).
Er
Ed
Exota!
E
Figure 3: N-Queens Problem Composition. To compose
model to solve the N-queens prob- lem, add energy of each row (E, ), each
column (E ). each diagonal & (E of the chessboard. We then sample from the resulting ergy function Ejotal to generate valid solutions.
Baselines. We compare against existing baselines for neural combinatorial optimization solvers, for which the N-queens problem is represented as graph, and the solution corresponds to Maximum Independent Set (MIS). As baselines we include: a reinforcement learning approach, where the model learns to defer harder nodes when solving the problem (LWD, [3]), an unsupervised method, where the model learns a Markov decision process over graphs (GFlowNets, [72]), and a supervised categorical diffusion solver (DIFUSCO, [60]). Furthermore, we also compare against previous state
5

---

Correct
Model
Type
Size
Instances
LWD
RL +
22
7.1000 = 0.5744
GFlowNets
+
6.9293 = 0.5904
DIFUSCO (7=50)
6,9400 0.6452
Fast T2T (T==1TT==1)
+
6.8200 0.8761
Fast 12T (1==1,
SL + GS
6.7000 = 0.7141
5. Fast T2T (7y=5, 5
SL +
20
7.0600
1 0.5800
Fast T2T
SL | GS
41
7,3800
1 0,5436
FBM (P=1024) (Ours)
ST. + PEM
97
7.9699 + 0.1714
Table 8-Queens Problem Evaluation. We compare the perfor-
mance against state-of ofate-of-the-arr 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 ergy (top, right) and incorrect solutions (bot- pled 100 8-queens solutions. IR: Iterative Refinement, BP: Belief Propagation, TS: Tree Search, S: Sampling, GS: Guided Sam- pling.
Figure 4: Energy Map Visualization. Cor- rect solutions (top, left) are assigned low en-
tom, left) assigned higher energy (bottom, right). Energy each position the sum the row, column, and diagonal energy.
T=100
T=75
T=50
T=25
T=10
T=4
T=3
T=2
T=1
Solation
Reverse dirfusion
PEM
(P-8)
Figure 5: Optimized Samples Across Timesteps. Generated samples where yellow squares represent queens placed in the chesshoard. Reverse diffusion fails to find valid solutions (top). Increasing the number of particles with PEM leads from invalid solutions (middle) valid solutions (bottom).
of the combinatorial optimization models (Fast T2T [42]), with different inference steps T, and gradient search steps Tg. For the latter, we also compared against the guided sampling version, where a penalty function added for the MIS problem guide denoising.
For all the methods evaluated, we report the number of correct instances of the problem found, and the average number of queens placed in the chessboard. We follow previous works approach to decode solutions, where, given model heatmap of the chessboard, we perform greedy decoding by sequentially placing queens in the board until conflict found.
Quantitative Results. We report the comparison of our approach with the previous baselines in Table 1. For all the methods reported, we sampled 100 different solutions. In this table, that our approach able to generate nearly all perfect solutions to the problem. Furthermore, our method significantly outperforms the previous ttate-of-the-an solvers. Out of 100 generated samples, 97 are valid 8-queens solutions, while state-of-the-are art methods are able to find correct instances at most.
Qualitative Results. In Figure we visualize the sampling process our approach using reverse diffusion and PEM. We can observe that PEM is able generate much better quality samples than reverse diffusion. Additionally, more particles leads to the generation of valid solutions. We include example of parallel sampling with particles in Appendix B. In Figure 4 we energy is assigned to rows and columns where the constraints are violated.
that higher
Performance with Increased Computation. In Table 2 we report the performance of our approach on the 8-queens problem with increasing number of particles during sampling. We show that increasing the number of particles significantly improves the quality of the generated samples and, as a consequence, a larger number of correct instances are found. In Figure 6, we visualize performance of our model with different number of particles on different complexity levels of the problem. We can see that by adjusting the number particles, we can adapt the ability of our model to solve more difficult problems.

---

80
Num. Particles
Correct Instances
Size
Carrect Instunces
8
6.6599 + 0.7550
64
34
7.2500 0.6092
128
7.8800 0.3265
20
256
89
7.9000 + 0..015
(lggace
512
92
7.9400 = 0.2336
Number of Particees
1024
7.9699 0.1714
Figure 6: Number of Particles vs Correct Instances Across Problem Difficulties. We sampled 100 solu- tions each from N-queens problems of increasing diffi- culty (7 10 queens). In all cases, a higher number particles results in more correct instances.
Table 2: Number of Particles vs Correct Instances. We sampled 100 solutions from the 8-queens problem. Increasing the number of particles with PEM signifi- cantly improves the number of correct instances.
Correct
Sampler
Size
Diffusion
Contrastive Loss
Correct
Satisfied Clauses
Instances
Loss
Instances
Reverse Diffusion
12
2.6400 + 0.7722
ULA
5
2.5200 + 0.6432
No
Yes
1.3200 + 0,5482
MALA
8
2.6700 + 0.6824
Yes
No
6,6799 + 0.7089
UHMC
9
2.6700 1 0.6971
Yes
Yes
97
7.9699 + 0.1714
HMC
11
2.6900 1 0.7204
PEM (P=8)
99
3,9900 + O. 1000
Table Loss Ablation. Ablations proposed the loss function on the performance on the 8-queens prohlem. We sampled 100 solutions from the 8-queens problem. combination of both diffusion and contrastive loss to shape the landscape produces the best results on the task. In all cases we sampled using PEM (P=1024).
Table 3: Sampler Ablation. Ablations proposed for different samplers. We sampled 100 solutions from the q-quees problem. Compared to other samplers. PEM able to consistently produce accurate solutions.
Ablation Study. In Table 3, we compare the performance with different methods proposed for EBM sampling, including Unadjusted Langevin Dynamics (ULA). Metropolis Adjusted Langevin Dynamies (MALA), Unadjusted Hamiltonian Monte Carlo (UHMC) and Hamiltonian Monte Carlo (HMC). We show that our approach substantially outperforms existing samplers in the 4-queens task. In Table 4, compare three models trained with different loss functions. We show that the combination of the diffusion and contrastive losses produces improved results on the task.
4.2
SAT Problem
Setup. In this section we evaluate the performance of our approach on the Boolean satisfiability problem (SAT), well known to be an NP-complete problem. The 3-SAT problem is binary decision problem where a Boolean formula given in Conjunctive Normal Form (CNF), with each clause having exactly 3 literals. The task is to find truth assignment to the variables (true or false) such that the formula evaluates to true. For training, we generated random 3-SAT instances with number of variables within [10, 20). The number of clauses set to be in phase transition, that is, it was set to be 4.258 x n, where n is the number variables [57]. For our approach, we train a model to generate satisfiable assignment to only individual clause of the 3-SAT problem. We then compose the model to generate a solution to the entire problem. This enables the generalization to an arbitrary number of clauses. We evaluate using the SATLIB benchmark [34]. For distribution similar to the training one, we used 100 instances with 20 variables and 91 clauses. For a larger distribution, we used 100 instances with 50 variables and 218 clauses.
Baselines. We compare against cxisting baselines for neural SAT solvers, including: the seminal neural SAT solver, where the solution is iteratively refined with increasing number of steps ANe [58]), and the state-of-the-art neural solver based on proicf-propagation (NSNet [43]). both cases, we use different number of steps T for solution refinement. When feasible, we also compare with combinatorial optimization models by encoding the 3-SAT problem as a graph.
Quantitative Results. In Table we find that our method significantly outperforms the previous stateoof-the-an neural SAT solvers, and able to find a larger number of correct instances of the problem. In the similar distribution 91 instances are solved compared to 58 instances solved by NSNet. the larger distribution our method still outperforms the previous other methods with 43 correct instances, with NSNet solving 37 correct instances.
7

---

Similar Distribution
Larger Distribution
Correct Instances f
Satisfied
Correct
Satisfied
Model
Type
Clauses
Instances f
Clauses f
GCN
SL
5
0,9617 0.0264
0
0.9569 +: 0.0203
DGL
SL TS
10
0,9520 0.0330
0
0.8705 0.0405
DIFUSCO (T- 50)
SL +
0.9734 0.0156
0
0.9738 0.0156
Fast T2T (7s= Tc=1)
SL + S
23
0,9749 0.0210
4
0.9751 0.0141
Fast T2T (7==5,7c=1)
SL +
22
0,9760 = 0.0273
20
0.9734 0.0159
NeurSAT (7==50)
SI. + IR
0,9661 0.0185
0
0.9651 0.0110
NeuroSAT (7=500)
SL +
0,9742 + 0.0154
0
0.9697 0.0111
NSNet (7=50)
SL + BP
0,9845 + 0.0272
34
0,9817 0.0237
NSNet (7=500)
+ BP
0,9856 = 0.0266
37
0,9846 0.0205
EBM (P-10244 (Ours)
SL + PEM
91
0.9985 + 0.0048
43
0.9963 0.0046
Table 5: 3-SAT Problem Evaluation. We compare the performance against the state-of of the- combinatorial optimization models and neural SAT solvers on the 3-SAT task. Models evaluated on distribution similar to the training distribution and larger distribution. Similar distribution has 100 instances with 20 variables and 91 clauses, while larger distribution has 100 instances with 50 variahles and 218 clauses. Our approach outperforms existing methods. IR: Iterative Refinement, BP: Belief Propagation, TS: Tree Search, S: Sampling.
Similar Distribution
Larger Distribution
Finetuning Correct
Satisfied Clauses
Num.
Correct Instances t Clauses
Satisfied
Correct
Satisfied Clauses
Instances
Particles
Instances
No Yes
57
0.9951 - 0.0068
30
0.9874
2
0,9910
91
0.9985 0.0048
70
0.9962
0,9910
Table 7: Fine-tuning Ablation. Abla- tions proposed for the finetuning of the model on the performance the 3-SAT problem. We show that finetuning the the evaluation performance of our approach on the 3-SAT problem composed model with complete instances with increasing number of particles. Increasing the number of particles leads to better performance. In all cases substantially improves the generalization performance of the model. we sampled using PEM (P=1024).
128
78
0.9975
18
0.9927
1024
91
0.9985
43
0.9963
Table Number of Particles 3-SAT Performance. We compare
Qualitative Results. In Appendix B, we present additional qualitative results for 3-SAT, where we show that unsatisfied clauses are assigned higher energy, while satisfied clauses are assigned lower.
Performance with Increased Computation. We assess in Table 6, the impact of the number of particles on the performance on the 3-SAT problem. We show that increasing the number of particles improves the results on both the similar and larger distributions. By formulating the problem as an en- ergy minimization problem, we can adjust the number of particles to adapt to the difficulty of the task.
Ablation Study. In Appendix C we include ablation for the sampling procedure, showing that PEM outperforms existing methods, as well as ablations on the training loss. Moreover, in Table 7 we report that finetuning the composed model improves the overall performance.
4.3 Graph Coloring
Setup. In this section, evaluate approach the graphical problem of graph coloring. Given a graph instance, the task is to assign a color to each node in the graph such that no two adjacent nodes share the same color using at most & colors. This problem is known be NP-complete. The chromatic number of graph is the minimum number of colors needed to color the graph. To train baselines, we followed the same approach in [41], and generated random graphs with number of nodes within [20, 40). density within [0.01, 0.5), and chromatic number x within [3,8], 8]. For our approach, we train model to generate a valid coloring of individual edge given set of colors. We then compose the model for all the edges of the graph to generate a valid coloring solution. To train our model we generate random pairs of different colors. For evaluation, we use graphs from the well-known COLOR benchmark P. Additionally, we also evaluate on random graph instances generated following different graph distributions, namely: Renos-Renyi [25]. Holme- Kim [33], and random regular expander graphs [4]. For each distribution, we generate smaller graphs with nodes within [20, 40) and larger graphs with nodes within [80, 100). Moreover, we also evaluate on densely connected regular graphs such as Paley graphs [51] and complete graphs. For all the methods, we
'nttps://mat. tepper. cmu, wwt.:c............eee edu/COLORO2/
8

---

(a) Original Graph (b) Correct Solution (c) Correct Solution (d) Incorrect Solution (c) Incorrect Solution Energy
Energy
Figure 7: Qualitative Illustration Energy Maps. We show a graph instance (a) along with valid solution
(b). In (c), we visualize the energy map of a correct solution, where we show the energy of each edge individually. We present an incorrect solution with two conflicting edges in (d) and the corresponding energy for each edge in (e). As expected, higher energy is assigned to conflicting edges.
Distribation
d
GCN
GAT
XIVIN
GIN--CPP
EBM (Ours)
(P=128)
Renyi
39]
[81. 99]
[193, 225]
0.05
[3.4]
151.60
12.08
130:20
11.47
03.80
-1.112
53.80
8.34
29.20
8:05
Holme Kim
[22. 31]
[56. 92]
0.26
[1.4]
T1.00
14.74
51.20 1: 10.03
29.00
7.75
15.20 1
7.66
10.60
10.60111
2.70
[86. 100]
[398, 469)
0.10
408
26.40
253.20 +: 50.71
182.60
55.20
12:63
58.00
3.74
Regular Expander
38
[86, 100)
[3.
144.0
6.90
112.60
Paley
[19. 37)
[171, 465
0.80
16.
10)
285.00
117.70
230.20
+ 159.43
151.80
91.20
34.80
20.27
Complete
[8. 121
[36, 66]
L.00
[8, 121
36,00
15.04
46.00
+: 15.04
34.80
16.42
30,00
4:
2.54
3.40
1.14
Table 8: Graph Coloring Evaluation. We compare the performance against canonical GNNs and GNN-based methods for graph coloring different random graph distributions and the COLOR benchmark. Performance measured 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 generalizes better to larger and denser graphs. Here V= Nodes, E Edges, d= Average density, x Chromatic number.
generate a coloring of the graph with & colors, where k the chromatic number of the given graph, and report the number of conflicting edges in the generated solution.
Baselines. We compare against existing baselines for neural solvers for graph coloring that trained to generalize to novel graph instances (GNN-GCP [41]). We also include a comparison with canonical Graph Neural Networks (GCN [39] and GAT [62]), and RL guided by Neural Algorithmic Reasoners (XLVIN [16)).
Quantitative Results. We compare our approach with GNN-based methods for graph coloring. In Table 8, we report the performance on random graphs generated following different distributions and the average performance on the COLOR benchmark. We show that our approach is able to generalize to larger graphs and different distributions better than existing methods. The detailed performance on the COLOR benchmark can be found in Appendix B, where approach significantly outperforms existing methods. Notably, while GNN-based methods show increasingly worse performance on larger graphs, our method maintains good performance across scales.
Qualitative Results. We present in Figure a graph instance with a valid coloring solution. We visualize the energy map and show that low energy is assigned to all the edges that compose the whole graph. We also show an incorrect solution with two confiictigg edges. The energy of the two conflicting edges is higher than the non-confficting ones.
Performance with Increased Computation. In Appendix B we report that increasing the number of particles from 8 to 1024 decreases the average number of conflicting edges from 15.0 to 8.0.
Ablation Study. We include in Appendix C ablations the sampling procedure, which yields 8.0 conflicting edges on average with PEM compared to 12.3 edges with UHMC. Additionally, we ablate the training losses, where obtain average of 15.0 conflicting edges with diffusion loss only and 9.0 when using contrastive loss only.
4.4 Crosswords
Setup. In this section, we report the results on crosswords puzzle solving. Crosswords are word puzzles where letters are arranged in a grid, with words intersecting both horizontally and vertically. Each word is associated with clue that provides definition, context or hint for the answer. The goal is to fill the grid so that all words satisfy both the clues and the grid constraints. For approach,

---

P1
P2
P3
Pa
Ps
P6
P7
Ps
R
FIFE
y (50)
3
222
LADE
AAD
(20))
((1)
Figure 8: Optimized Samples Across Timesteps and Particles. We show samples y generated on the crossword puzzle with PEM (I*=8) at timestep for different particles Pt, where indicates the particle number. The given crossword has five horizontal clues and five vertical clues. At the end of the optimization process, the PEM is able generate valid solutions the puzzle. Cells in red indicate incorrect letters.
train model generate valid word given precomputed embeddings of the corresponding hint. To solve complete crossword, we compose horizontally and vertically the model to form the given grid. We evaluate on the Crosswords Mini Benchmark introduced in [71].
Baselines. We compare against different inference algorithms for Large Language Models, including: Standard Input-Output (IO), Chain Thought (CoT)
Model
Letter Success Rate
Word Success Rate
10
38.7
14.0
CoT
40.6
15.6
[65] and Tree of Thought (ToT) [71].
ToT
EBM (Ours)
80.4
50.5
Quantitative Results. In Table we compare our approach with various LLM inference methods. Our approach significantly outperforms both the and CoT baselines. Furthermore, it achieves performance competitive with ToT. While ToT attains a slightly higher average word success rate (60.0% vs. 50.5%), our method achieves higher overall grid completion rate (80.4% vs. 78.0%).
Table 9: Crosswords Evaluation. We compare against différent strategies for LLM inference. We report average over crosswords. We used P=1024 for sampling with PEM.
Qualitative Results. Figure shows the particles generated for a crossword across timesteps. It can be seen that, during the optimization process, different particles explore different solutions the puzzle. In the end, the optimization algorithm successfully finds valid solution to the crossword.
Limitations and Conclusion
Limitations. limitation of our method is that it assumes starting Gaussian distribution and models optimization sequence of Gaussian increments. Future work could explore non-Gaussian objectives and initializations that enable recurrent improvement of initial solutions. Another limitation is that, while our method excels on N-Queens and 3-SAT, there is still room for improvement in achieving optimal solutions for graph coloring. Further research could investigate alternative training strategies for EBMs to produce more accurate energy landscapes
Conclusion. In this work, we propose novel approach solving reasoning problems using EBMs. By formulating the problem as an energy minimization, we decompose them into smaller subproblems, each with its own energy landscape. We then combine the objective function of each landscape to construct energy landscapes for more complex instances. To train EBMs, we have proposed a combination of diffusion and contrastive loss, which yielded superior performance. We also introduced PEM, parallel optimization method with an adaptable computation budget for sampling from the resulting energy functions. Our approach outperforms existing methods on several reasoning tasks, including N-queens, 3-SAT, graph coloring, and crossword puzzles, pointing to the promise of this approach.
10

---

Acknowledgments
We thank Siba Smarak Panigrahi for his feedback during the preparation of this manuscript.
References
[1] Rakefet Ackerman and Valerie A Thompson. Meta-reasoning: Monitoring and control of thinking and reasoning. Trends in cognitive sciences, 21(8):607-617, 2017.
[2] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd. Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019. 2
[3] Sungsoo Ahn, Younggyo Seo, and Jinwoo Shin. Learning what to defer for maximum indepen- dent sets. In International conference on machine learning, pages 134-144. PMLR, 2020. 3,
[4] Noga Alon and Yuval Roichman. Random cayley graphs and expanders. Random Structures Algorithms, 5(2):271-2844 1994. 8
[5] Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as a layer in neural networks. International conference on machine learning, pages 136-145. PMLR, 2017. 2
[6] Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. arXiv preprint urXiv:2003.05033 2020. 3
[7] Arpit Bansal, Avi Schwarzschild, Fitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35:20232-20242, 2022.
[8] Irwan Bello, Hieu Pham, Quoe V Le, Mohammad Norouzi, and Samy Bengio. Neural com- binatorial optimization with reinforcement learning. arXiv preprint 1oXiv:/611.9900 2016.
[9] Léon Bottou. From machine learning machine reasoning: An essay. Machine learning, 94:133-149, 2014.
[10] Quentin Cappart, Didier Chételat, Elias B Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković. Combinatorial optimization and reasoning with graph neural networks. Journal of Machine Learning Research, 24(130): 2023. 3
[11] Davide Carbone, Mengjian Hua, Simon Coste, and Eric Vanden- Bijuden- Efficient training of energy-based models using jarzynski equality. Advances in Neural Information Processing Systems, 36:52583- 52614, 2023. 3
[12] Xinyun Chen, Chen Liang, Adams Wei Yu, Dawn Song, and Denny Zhou. Compositional generalization via neural-symbolit stack machines. Advances in Neural Information Processing Systems, 33:1690-1701. 1701, 2020.
[13] Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint 1rXiv:1609.01704, 2016.
[14] Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, et al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261. 2025. 23
[15] Armand Comas, Yilun Du Christian Fernandez Lopez, Sandesh Ghimire, Mario Sznaier, Joshua B Tenenbaum, and Octavia Camps. Inferring relational potentials in interacting systems. In International Conference on Machine Learning, pages 6364-6383, PMLR, 2023.
[16] Andreea Deac, Petar Veličković, Ognjen Milinković, Pierre-Luc Bacon, Jian Tang, and Mladen Nikolić. Xlvin: executed latent value iteration nets. arXiv preprint 20Xiv:2010.13144 2020.
11

---

[17] Carles
Alberto Bietti, Eric Hunden-Eijnden and Joan Bruna. On energy based models with overparametrizeed shallow neural networks. In International Conference on Machine Learning, pages 2771-2782. PMIR, 2021. 3
[18] Honghua Dong, Jiayuan Mao, Tian Lin, Chong Wang, Lihong Li, and Denny Zhou. Neural logic machines. arXiv preprint arXiv:1904 1694, 2019.
[19] Arnaud Doucet, Nando De Freitas, and Neil Gordon. introduction to sequential monte carlo methods. Sequential Monte Carlo methods in practice, pages 3-14, 2001. 4
[20] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. In International conference on machine learning, pages 8489-8510. PMLR, 2023. 3
[21] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. arXiv preprini arXiv:2012.013 2020. 3
[22] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning, pages 5570- 5582. PMLR, 2022. 2,
[23] Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion. arXiv preprint arXiv:2406. 2024. 2,
[24] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019. 3,
[25] Paul Erdös and Alfréd Rényi. On the strength of connectedness of random graph. Acta Mathematica Hungarica, 12(1):261-267, 1961.
[26] Richard Evans and Edward Grefenstette. Learning explanatory rules from noisy data. Journal of Artificial Intelligence Research, 61:1-64, 2018. 2
[27] Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks on out-r--distribution graphs. IEEE transactions on pattern analysis and machine intelligence, 46(1):322-333 2023. 3
(28) Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobatt Benjamín Sánchez-Lengellnng Dennis Sheberla, Jorge Aguilera-Iparraguire Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molccules. ACS central science, 4(2):268-21) 276, 2018.
[29] Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.088 2016. 3
[30] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint
arXiv:1410.5401, 2014. 3
[31] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-rl incentivizes reasoning in lims through reinforcement learning. Nature, 645(808 1)x633-633 2025. 23
[32] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 1473) 2002. 3
[33] Petter Holme and Beom Jun Kim. Growing scale-free networks with tunable clustering, Physical review E, 65(2):026107, 2002.
[34] Holger H Hoos and Thomas Stützle. Satlib: An online resource for research on sat. Sat, 2000:283-292, 2000. 7
[35] André Hottung, Bhanu Bhandari, and Kevin Tiemeey Learning a latent search space for routing problems using variational autoencoders. In International Conference on Learning Representations, 2021. 2
12

---

[36] Chaitanya K Joshi, Thomas Laurent, and Xavier Bresson. An efficient graph convolutional network technique for the travelling salesman problem. arXiv preprint arXiv:1906.01227, 2019.
3
[37] Daniel Kahneman. Thinking, fast and slow. macmillan, 2011. 1
[38] Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv:1511.08228, 2015. 3
arXiv preprint
[39] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint me0p::0009007 2016.
[40] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. A tutorial on energy-based learning. Predicting structured data, 1(0). 2006.
[41] Henrique Lemos, Marcelo Prates, Pedro Avelar, and Luis Lamb. Graph colouring meets deep learning: Effective graph neural network models for combinatorial problems. In 2019 IEEE 31st International Conference on Tools with Artificial Intelligence (ICTAI), pages 879-885.
IEFE, 2019. 8, 9, 23
[42] Yang Li, Jinpei Guo, Runzhong Wang, Hongyuan Zha, and Junchi Yan. Fast t2t: Optimization consistency speeds up diffusion- grainingtt--tsstingg solving for combinatorial optimization. Advances in Neural Information Processing Systems, 37:30179-3003 2024. 3, 6
[43] Zhaoyu Li and Xujie Si. Nsnet: A general neural probabilistic framework for satisfiability problems. Advances in Neural Information Processing Systems, 35:25573-25585, 2022. 3,
[44] Lukas Lüken and Sergio Lucia. Self-supervised learning of iterative solvers for constrained optimization. arXiv preprint XXXv:2440.088066, 2024. 3
[45] Robin Manhaeve, Sebastijan Dumancic, Angelika Kimmig, Thomas Demeester, and Luc De Raedt. Deepproblog: Neural probabilistic logic programming. Advances in neural informa- tion processing systems, 31, 2018. 2
[46] Sean Michael McLeish and Long Tran-Thanhe [re] end-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking. ML Reproducibility Challenge 2022, 2022. 25
[47] Pasquale Minervini, Sebastian Riedel, Pontus Stenetorp. Edward Grefenstette, and Tim Rock- täschel. Learning reasoning strategies in end to-tto-end differentiable proving. In International Conference Machine Learning, pages
PMLR, 2020. 2
[48] Arvind Neelakantan, Quoe V Le, and Ilya Sutskever. Neural programmer: Inducing latent programs with gradient descent. arXiv preprint arXiv:1511.04833 2015. 3
[49] Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: backpropagating through discrete exponential family distributions. Advances in Neural Information Processing Systems, 34:14567-145799, 2021. 2
[50] Erik Nijkamp, Mitch Hill, Tian Han, Song-Chun Zhu, and Ying Nian Wu. On the anatomy of meme-based maximum likelihood learning of energy-based models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5272-5280, 2020. 3, 4
[51] Raymond EAC Paley. On orthogonal matrices. Journal of mathematics and Physies, 12(1-
4):311-320, 1933. 8
[52] Scott Reed and Nando De Freitas. Neural programmer-interpretess. arXiv preprint 1rXiv:1511.06279, 2015. 3
[53] Tim Rocktäschel and Sebastian Riedel, End-to-end differentiable proving. Advances in neural information processing systems, 30, 2017.
[54] Yulia Rubanova, Alvaro Sannhez- Tobias Pfaff, and Peter Battaglia. Constraint-basee graph network simulator. arXiv preprint arXiv:2112.091611 2021. 2
13

---

[55] Martin JA Schuetz, J Kyle Brubaker, and Helmut G Katzgraber. Combinatorial optimization with physics-inspired graph neural networks. Nature Machine Intelligence, 4(4):367-377, 2022.
3
[56] Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy hard problems with recurrent networks. Advances in Neural Information Processing Systems, 34:6695-677 6706, 2021.
3
[57] Bart Selman, David G Mitchell, and Hector Levesque. Generating hard satisfiability problems. Artificial intelligence, 81 (11-2):17729, 1996. 7, 22
[58] Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David L. Dill. Learning sat solver from single-bit supervision. arXiv preprint arXiv:1802.03685, 2018.
3,
[59] Chen Shengyuan, Yunfeng Cai, Huang Fang, Xiao Huang, and Mingming Sun. Differentiable neuro-symbolic reasoning on large-scale knowledge graphs. Advances in Neural Information Processing Systems, 36:28139-2222888 299-2815444 2023. 2
[60] Zhiqing Sun and Yiming Yang. Difusco: Graph- diffusion solvers for combinatorial optimization. Advances in neural information processing systems, 06:3706-3731, 2023. 3,5
[61] Mokanarangan Thayaparan, Marco Valentino, and André Freitas. A differentiable integer linear programming solver for explanation-based natural language inference. arXiv preprint 2rX1v:2404.02625, 2024. 2
[62] Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio, et al. Graph attention networks. stat, 1050(20):10 2017. 9
[63] Eric Wallace, Nicholas Tomlin, Albert Xu, Kevin Yang. Eshaan Pathak, Matthew Ginsberg, and Dan Klein. Automated crossword solving. arXiv preprint arXiv:2205. 09660, 2022. 24
[64] Po-Wei Wang, Priya Donti, Bryan Wilder, and Zico Kolter. Satnet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver. International Conference on Machine Learning, pages 6545-6554. PMLR, 2019. 2,
[65] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. of-innoof-thouhh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24884774-44877, 2022. 10
[66] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational autoencoders and energy- based models. arXiv preprint wxXiv:2010.00654 2020.
[67] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009.1 11848, 2020. 3
[68] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu L.,, et al. Qwen3 technical report. arXiv preprint arXiv.2505.09388, 2025. 23
[69] Fan Yang, Zhilin Yang, and William W Cohen. Differentiable learning of logical rules for knowledge base reasoning. Advances in newral information processing systems, 30, 2017. 3
[70] Zhun Yang, Adam Ishay, and Joohyung Lee. Learning to solve constraint satisfaction problems with recurrent transformer. arXiv preprint arXiv:2307.0445 2023. 3
[71] Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths. Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in neural information processing systems, 36: 1111-18111888 2023. 10, 24
[72] Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron C Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial problems with gflownets. Advances in neural information processing systems, 36:1 1199,, 2023. 5,
14

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
15

---

Justification:
Not Applicable
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Sections 3, and Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification:
Data and code for reproduction are released as supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..
Sewiidde/CodSSubiissee asssnnollcccyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /CodeSubmic / //nips. wttpccccpppccccceeececccccccccct /tt.:///ggiiieeeeeeeuuubiiiiii ssisnnolllll onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: See Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See Appendix A.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: h/w/ueuurpp...e https///c//.......cccciiieeeeee.
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: No direct societal impact can be derived as an outcome from our work. The work proposes core machine learning contributions for improving generalization in reasoning tasks and does not involve sensitive data or high-risk applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: No high risk exist in the misuse of the models. Released data involves only abstract reasoning tasks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: See Section 4.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: See Sections 4 and Appendix A.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing experiments were performed in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No crowdsourcing experiments were performed in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Core method development does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https /20p:/pp...//ooeeeeeeee///////// for what should or should not be described.
21

---

Appendix Overview
The Appendix is organized follows: Section A describes the experimental setting, including hyperparameters and training details, Section B presents additional quantitative and qualitative results on the problems considered in the paper, and Section C includes ablation studies on the sampling procedure and training losses.
A Experimental Settings
A.1 N-Queens Problem
Setup. We used a single instance solution of 8-Queens for both training and validation. For all approaches, the best model was selected based on the validation performance.
Compositional Approach. In our approach, train model using the 8 rows from the instance. We use each of the rows as targets to be generated by the model. As negative samples, we used a row without queens and row with two queens (see Figure 9). In other words, we train model to generate a vector of length 8 with exclusively one I and seven Os. To compose a 8-queens solver, we simultaneously add the energy of all rows, columns and diagonals using the same model. For the diagonals, we pad the rows with zeros to the right. Notice that this composition, as it is, assumes as a constraint that queen has to be placed in each diagonal, however, this constraint does not exist in the original problem. Alternatively, we could have trained two separate models: (I) one for the rows and columns, where the constraint is to always have one queen placed, and (2) one for the diagonals, where the constraint is to have either zero or one queen placed. Empirically, we observed that the single model approach led to better heatmaps than training two separate models for rows/columns and diagonals.
Training. As model, used 3-layer MLP, with each layer having: layer normalization and linear layers of dimensions 128, 256, 128, followed by ReLU activation. We added skip connections for each layer. The model was trained with learning rate of le-4 with Adam W optimizer for 20000 epochs with a batch size of 2048. For the contrastive loss, we used a weight of 0.5. For scheduled noise we used a linear schedule with T 100 timesteps. With single Nvidia A10 GPU with 24GB of memory, the model was trained in approximately 5 hours.
Baselines. For all baselines used the default hyperparameters proposed in each work for MIS solving on SATLIB unless stated otherwise. To encode the N-queens problem as a MIS problem for each position of the chessboard nodes in the same row, column and diagonal. For LWD we included self-loops for each node. For GFlowNets, trained model with 8 layers for 1500 epochs with batch size of 128. For DIFUSCO we trained a model with 8 layers for 5000 epochs. For FastT2T we trained a model with 8 layers for 20000 epochs. In Table 10 report the wall clock time required to sample 25 solutions of the N-queens problem for all baselines considered. Results are obtained using single Nvidia A10 GPU with 24GB of memory
created node, and then added edge with each of the other
A.2 3-SAT Problem
Setup. We generated 4000 random satisfiable 3-SAT instances for training and 1000 for validation, using the cnf enfgen Python package. For each instance, we randomly sample the number of variables within [10, 20) and assign a number of clauses equal to 4.258 x n, where n is the number of variables. This is, we generate samples to be in phase transition. which technically makes it more difficult instances to solve [57]. For evaluation, we used instances from SATLIB, which are also generated to be in phase transition. We used 100 satisfiable instances with 20 variables and 91 clauses, and 100 satisfiable instances with 50 variables and 218 clauses. For all approaches considered, the best model was selected based on validation performance.
Compositional Approach. As base model, we train a model to generate an assignment that satisfies a single clause. For instance, given a clause (aVbV-e), two example valid assignments are (1, 0, 1) or (0, 0, 0), where 1-True and 0=False. Notice that for each clause with three literals, we have seven valid assignments and one invalid assignment. For each clause, negative sample we use the invalid assignment, which is given by the negation of the clause. For instance, for the previous clause (0, 0, 1) is not a satisfiable assignment. To obtain the clauses for training split each of the
22

---

Model
LwD
1.16 = 0.04
GFlowNets
1.12 = 0.08
DIPUSCO
33.46 1.36
Correct
Unique
FaxT2T (1s=1. 1)=1)
1.82 = 0.18
Model
Instances
Solutions
HarlZT (fx=1. 1a GS)
1.91 = 0.19
FaxT2T (7:=5. (5
3.37
DeepT FBM (P-128) (Ours)
FasT22 (7:=5. T5-5, GS)
7.75 = 0.00
100
37
EBM (P=1024) (Ours)
84.99 = 1.61
Table 11: 8-Queens Problem Evaluation. We com-
Table 10: 8-Queens Problem Inference Time. We report the wall clock time required generate 25 so- lutions of the problem. Values are averaged over five different runs.
pare the number of unique solutions generated against the DeepTT approach. While Deep'T is deterministic approach, our method generates 37 different solutions to the problem out of 100 correct instances.
Num. Particles
Energy
Solution
Energy
2
136.62 2.26
Random
151.43 + 3.42 137.33 + 1.86
4
136.29 = 1.98
One invalid queen
8
133.93 1.67
Correct
130.78 + 0.37
16
133.82 2.05
32
132. 1.16
Table 13: Solution Type vs Energy. We show the comparison of energy values in the 8-Queens prob- lem for three solution types: incorrect solutions with 8 randomly placed queens, nearly-correct solutions with one misplaced queen, and correct solutions. Correct solutions have, on average, the lowest energy values compared to the other solution types. Values are aver- aged over five runs.
64
132.36 0.99
128
132.00 1.03
256
131.84 0.97
Table 12: Sampled Energy vs Number of Particles. A larger number of particles enables sampling lower energy values in the N- I-Queens problem. Values are averaged over five runs.
4000 instances into separate clauses. To compose the models, we add the energy of all clauses that make an instance.
Training. We trained 3-layer MLP with skip connections and each layer having layer normalization,
3 linear layers of dimensions 128, 256, 128, followed by ReLU activation. As input, we use the concatenation of the generated sample (dimension 3) and the clause sign (dimension 3). We trained using the AdamW optimizer with a learning rate of le- for 20000 epochs and batch size 1024. We used 0.5 for the contrastive loss weight. For finetuning, we trained for an additional 10000 epochs with learning rate of 1e 4 and batch size 1024. For scheduled noise, we used linear schedule with T 100 timesteps. With a single Nvidia A10 GPU with 24GB of memory, the model was trained in approximately 12 hours.
Baselines. For all baselines, we used the default hyperparameters proposed in each work for MIS solving on SATLIB unless stated otherwise. For all the combinatorial optimization models, we modified the architecture to include the clause sign input. For GCN, we trained a model with 12 layers for 100 epochs with batch size 256. For DIFUSCO, we trained a model for 250 epochs with batch size 256. For FastT2T, we trained a model for 500 epochs with batch size 256. With neural SAT solvers, we used the hyperparameters specified in the corresponding works except for the following: for both Ae and NSNet, we trained for 300 epochs with batch size 128.
A.3 Graph Coloring Problem
Baselines. We extend the baselines presented in the main paper to include additional comparisons with state-of-the-an reasoning Large Language Models (LLMs). In particular, we consider Gemini 2.5 Pro [14], DeepSeek RI [31], and Qwen 235B-A22B-2555 [68].
Setup. We generated 1000 random graphs following the approach from [41]. The graphs have number of nodes within [20, 40), density within [0.01, 0.5), and chromatic number X within [3. 8). We then make a 90-10 split for training and validation. For evaluation, we generated ten random instances from each of the distributions: Erdos-Renyi, Holme-Kim, and random regular expander graphs, with five instances with nodes within [20, 40) and five instances with nodes within [80, 100). Additionally, we generated five Paley graphs with prime numbers between 19 and 37, and complete graphs from 8 to 12 nodes.
23

---

2 1
aa a avbv-
Positive Sample
18 12
(IVVD
Negative Samples
1 207
OVOVD
Figure 9: Negative Samples. (Left) In the N-Queens problem, the model is trained to generate rows containing exactly one queen. Negative samples consist of invalid rows with no queens multiple queens. (Middle) For the 3-SAT problem, the model leams to produce valid assignments obtained directly by negating the clause sign. (Right) In the graph coloring problem, the model is trained to assign different colors to the nodes of an edge. Edges with the same color for both nodes are used as negative samples.
single clause. The negative sample used
Compositional Approach. We trained base model to generate valid coloring for a single edge. We define a fixed set of colors, in this case k=14 colors. The model generates a one-hot vector encoding the left and right colors of the edge (in total, an output of dimension 28). As negative samples, we used edges with the same color for both nodes. Additionally, to enforce the generation of valid one-hot vectors, we also used as negatives vectors with random perturbations in a random position. To compose the model for the whole graph, we add the energy of all edges that compose the graph.
Training. We trained a 4-layer MLP with skip connections and each layer having layer normalization, 3 linear layers of dimensions 128, 256, 128, followed by ReLU activation. We trained using the A optimizer with a learning rate of 1e- for 50000 epochs and batch size 1024. We used 0.5 for the contrastive loss weight. For scheduled noise, used linear schedule with T 100 timesteps. With a single Nvidia A10 GPU with 24GB of memory, the model was trained in approximately hours.
Baselines. For both GCN and GAT, we trained a model with 8 layers with hidden dimension 128, dropout 0.1 with AdamW for 1000 epochs and batch size 512. We train the models using cross-entropy loss to predict the color of each node out of 14 colors. To train GNN-GCP we follow the same setting as in the original work and stop the training when the model achieves 82% accuracy and 0.35 binary cross-entropy loss on the training set averaged over 128 batches containing 16 instances. For XLVIN, we train the XLVIN-CP variant using CartPole-style synthetic graphs and use the same hyperparameteras reported in the original work.
A.4 Crosswords
Setup. We evaluate approaches based on their ability solve 5x5 crosswords with 10 words each. For evaluation, we use the 20 crosswords from the Mini Crosswords dataset from [71]. To train approach, we sample 32.7k and 6.8k five-letter words from the Crosswords QA dataset [63] for training and validation, respectively. Each entry in the dataset consists of a hint and a five-letter word. To generate embeddings of the hints,
the text sont---medddig-33 3--mmll model from OpenAI.
Compositional Approach. As a base model, we train a model that, given the embedding of a hint. is able to generate a five-letter word. We later compose the model to solve a 5x5 crossword by combining the energy functions of the five horizontal rows and five vertical columns, where each of the ten words has separate hint input.
Training. As a model, we use 3-layer MLP, with each layer having layer normalization and three linear layers of dimensions 1024. 1024, and 1024, followed by a ReLU activation. We add skip connections for each layer. The model receives embeddings of dimension 1536 as input. The model is trained with a learning rate of 10-1 using the AdamW optimizer for 20,000 epochs with batch size of 2048. We did not use contrastive loss for training the model. For scheduled noise, we linear schedule with = 100 timesteps.
24

---

P
P3
PA
PA
Po
Pa
g0000)
500
200
3/0
(8)
Figure 10: Optimized Samples Across Timesteps and Particles. We show samples y generated 8-queens problem with PEM (P-8) at timestep t for different particles P., where indicates the particle number. In the figure. yellow squares represent queens placed in the chessboard. PEM is able to generate valid instance 8-queen problem () with Pl, P3 and Ps). A first valid solution appears at y with PI.
B Additional Results
B.1 N-queens Problem
Quantitative Results. In Table 11 we provide quantitative comparison where we compare with Deep Thinking (DeepT [46]). While were able to successfully solve the 8-queens problem using both approaches, the DeepT approach is purely deterministic, producing the same solution each time. In the case of N-queens, where there multiple solutions, and input provided, this is a limitation. Using our approach, we are able to generate 37 unique solutions out 100 sampled solutions.
Qualitative Results. Figure 10 we show the samples generated with PEM different timesteps and particles. It can be seen that at the end of the sampling process, the procedure reaches valid solution for the 8-queens problem. During the optimization process, can seen how different particles partially approximate different solutions until they converge the same solution.
B.2 3-SAT Problem
Qualitative Results. In Figure 11 we consider an instance of the 3-SAT problem with four clauses and three variables. We compare two variable assignments: correct assignment satisfies all clauses, and an incorrect one that satisfies only three. The figure shows the energy computed by the model for each clause individually.
the correct solution, all clauses are assigned low energy. In contrast, in the incorrect solution, the unsatisfied clause has comparatively higher energy than the others. This consequence, when
highlights how the energy function effectively reflects clause satisfaction. the energy of all clauses is composed, higher energy is assigned to the incorrect solution.
25

---

(a Va Ve) A (-a -V) A (a (a (bv
Correct Solution: a=0
b=0 c=0
(OVOV 1)
(1V 1
(0V TV
(0V I 1)
Incorrect Solution: a
(OVOV 1) (1V OV 1) (0V OVO) A (1V 1
Figure 11: Qualitative Visualization of Energy Maps. Example of a 3-SAT instance with four clauses and three variables (top), along with correct solution (middle) and an incorrect solution (bottom), having 1 ==True 0= We show energy of each clause individually. A higher energy is assigned clause that evaluates to false (unsatisfied clause), while clauses that evaluate to true (satisfied clauses) are assigned lower energy.
Confficting Edges
Model
Erdos-Renyl Small
Erdos-Renyi Large
Gemini 2.5 Pro
20.00 14.22
142.66 + 6.11
Deepseek R1
4.00 1: 3.46
78.00 1. 36.38
Qwen 135B-A22B-22 PEM (P=128) (Ours)
16.00 +: 13.06
102.66 1: 26.10
3.15 + 2,00
45.81 + 11.88
Table 15: Comparison of Conflicting Edges across Models. Lower values indicate fewer conflicts. PEM (ours) achieves the lowest number of conflicting edges across both settings. Values averaged five graphs.
B.3
Graph Coloring Problem
Baslines
Quantitative Results. Table 16 we provide an evalua- tion of our approach on the COLOR benchmark. We compare against existing methods for graph coloring methods and canon- ical GNNs. We can see that methods based GNNs generalize worse with increasingly larger graphs. On the larger graph con- sidered, our approach generates solution with 69 conflicting edges, while GNN-GCP generates a solution with 667 conflict- ing edges, and GCN and GAT generate solutions with 1625 and 1454 conflicting edges, respectively. In Table 15 we compare our approach with state of the reasoning Large Language Models on graphs following the Brddo-Renyí distribution. While some models able solve complete instances in-context (e.g. DeepSeek achieves an average of 4.0 conflicting edges on small Erdos-Renyí graphs), their performance deteriorates as graph size increases. For larger graphs, these models fail to find effective in-context solutions, with the best achieving an average of 78.0 conflicting edges compared 45.81 with our approach.
Num.
Conflicting Edges J
Particles
8
15.0 + 2.64
64
14.3 + 3.51
128
10.3 1 2.51
1024
1. 2.64
Table 14: Number of Particles vs Conflicting Edges. We sampled three solutions for given graph instance. Increasing the number of particles with PEM leads on average to more optimal solutions.
Qualitative Results. In Figure 12 we show the samples generated with PEM at different timesteps 30, particle Pa generates the first valid solution and
and particles. We can that, at timestep that, at timestep / = 10, all particles have already converged to a valid solution. In Figure 13 we show an example energy landscape resulting from the composition of two edges, where the optimal solution corresponds to the minimum of the function. Additionally, in Figure 14 we show the evolution of the landscape over different timesteps.
Performance with Increased Computation. We assess the effect of increased computation in Tables 14 and 17. We find that a larger number of particles in sampling slightly improves performance on the graph coloring task. On average, the number of conflicting edges in the generated solution is lower with larger number of particles, meaning that the solution is closer the optimal solution.
Table 18 we evaluate the effect of the timesteps hyperparameter in the graph coloring performance using the Holme Kim distribution. We observe that increasing the number timesteps from 20 to 1000 leads to an average reduction in the number of conflicting edges, from 6.89 to 6.51 in small instances, and from 44.27 to 42.71 in larger instances.
Performance with Increased Timesteps.
26

---

P1
P2
P3
P4
Ps
P6
P1
Ps
(100)
y150)
((3))
(10)
g())
Figure 12: Optimized Samples Across Timesteps and Particles. We show samples y/ generated on the
graph coloring problem with PEM (P 8) at timestep f for different particles P., where indicates the particle number. The graph instance has eight edges, nine nodes and chromatic number x=2. PEM able to generate valid coloring for the graph (red and gray in the figure).
Enee
color
Figure 13: Composition of Energy Landscapes. (Left) Energy landscape for different values of from simple graph coloring problem with edge and fixed color red for node. The plot shows that the energy assigned gray color the lowest. By composing two energy landscapes, we can create a new function corresponding to a larger problem with two edges. (Right) The energy landscape resulting from composing the energy landscapes of two edges with one node fixed to color red. The plot shows the energy for combinations of colors for nodes nu and The assignment gray and ny=red yields the lowest energy, indicating that this the optimal solution.
Eo(g, 405
(v. (=35)
Ea(y 30)
Ea (y. (=25)
Ea(y, (=5)
Um
1008
Figure 14: Energy Landscapes Across Timesteps. Evolution of energy landscapes over time for graph coloring problem with two edges (shown in Figure 13). The landscapes transform a Gaussian distribution into gradually the target distribution. Eventually, the optimal solution becomes the global minimum of the function.
27

---

EBM (Ours)
Graph
V
d
GCN
GAT
GNN-GCP
(P=128)
mycic13
0.36
17
10
myciel4
23
0.28
71
31
10
queen5_5
25
160
0.53
5
160
160
96
23
queen6_ 6
36
290
0.46
290
290
116
37
myciel5
236
0.22
171
97
42
26
queen7_
0.40
476
476
216
queen8
64
728
0.36
728
728
272
73
1-Insertions 4
67
232
0.10
5
220
100
42
21
huck
74
301
0.11
11
216
234
172
20
jean
254
0.09
208
146
206
david
406
0.11
11
333
352
156
mug88_1
88
146
0.04
4
117
127
98
17
myciel6
95
755
0.17
755
755
100
98
queen8_ 12
96
1368
0.30
12
1368
1368
408
96
games120
120
0.09
574
438
418
anna
138
493
0.05
11
260
392
110
2-Insertions_4 _
149
541
0.05
5
270
304
198
64
homer
556
1629
0.01
13
1625
1454
667
69
Table 16: Graph Coloring Evaluation. We compare the performance against canonical GNNs and GNN-based methods for graph coloring on the COLOR benchmark. For each graph, we report the number of conflicting edges in the coloring solution. where lower is better. Our methods outperform existing methods almost all the instances and shows better generalization to larger graphs
Conflicting Edges
Num, Particles
Holme Kim Small
Holme Kim Large
128
10.60 + 2.70
59.00 5.24
256
8.38 + 3.28
56.89:112.75
512
9.68 + 3.13
54.68 8.10
1024
7.04 + 3.64
57.04 + 3.27
2048
7.79 + 4.20
55.57 + 4.58
Table Number Particles vs Conflicting Edges. report solutions averaged over five instances each distribution. Increasing the number of particles with PEM on average decreases the number of conflicting edges in Holme Kim graph distributions.
Conflicting Edges
Num. Timesteps
Holme Kim Small
Holme Kim Large
2.19
44.27 7.29
7.01 2.68
54.40
100
6.71 2.60
52.97 8.07
200
6.00 3.83
50.31 5.76
500
6.50 2.58
49.63 3.19
1000
6.51 + 2.35
42.71 3.93
Table 18: Number Timesteps vs Conflicting Edges. Training with higher number of timesteps decreases on average the number conflicting edges in Holme Kim graph distributions. We report solutions averaged over five instances of cach distribution. In all cases we sample using PEM (P=1024).
28

---

Similar Distribution
Larger Distribution
Diftusion Contrastive
Correct Instances
Satisfied
Correct
Satisfied
Correct Instances
Satisfied Clauses
Loss
Loss
Clauses
Sampler
Instances 1 Clauses
No
Yes
0
0,9331 = 0.0258
Yes
No
0,7742 : 0.0157
Reverse Diffusion
1
0,9521
0,9519
Yes
Yes
57
0,9951 0.0068
ULA
0
0,9524
0,9516
MALA
0
0,9519
0.9535
Table 20: Loss Ablation. Ablations proposed
UHMC
0,9502
0,9537
HMC
1
0,9553
0,9533
for the loss function the performance the 3-SAT problem. A combination of both diffusion loss to train the EBM and a con- trastive loss to shape the landscape leads to the best results. In all cases sampled using PEM (P=1024).
EBM (P = 1024)
91
0,9985
43
0,9963
Table 19: Sampler Ablation. Ablations proposed for samplers
the 3-SAT task. PEM significantly outperforms other samplers on the 3-SAT problem for both similar and larger distributions.
Conflicting Edges
Diffusion
Contrastive Loss
Conflicting Edges
Sampler
Loss
Reverse Diffusion
19.6 +: 3,51
No
Yes
9.0 2.00
ULA
28.3 + 4,93
No
= 4.00
Yes
Yes
8.0 2.64
MALA
4: 2.64
UHMC
12.3 + 2.51
Table 22: Loss Ablation. Ablations proposed for the loss function on the performance on the graph coloring task. We sample three so- lutions for given graph instance. combi- samplers on the graph coloring task. We sample three solutions nation of both diffusion and contrastive loss for given graph instance. On average PEM finds solution with leads to the best results. In all cases lower number conflicting edges.
HMC
14.6 + 0.57
PEM (P = 1024)
8.0 2.64
Table 21: Sampler Ablation. Ablations proposed for different
pled using PEM (P=1024).
C Ablation Study
3-SAT Problem. We ablate the sampling procedure in Table 19. We compare the performance of our method with different samplers. In the similar distribution setting, our method successfully finds satisfiable solutions 91 out of 100 instances, whereas baseline samples find at most one. In the larger distribution setting, our method finds 43 satisfiable solutions out of 100, while baseline samples do not succeed. We also ablate the training losses Table 20. A model trained with both diffusion and contrastive loss solves 57 out of 100 instances, compared to 11 with diffusion loss only and 0 with contrastive loss only.
We propose ablations for both training and sampling of EBMs on the graph coloring task. In Table 21, we show that on average our sampling procedure produces on average more optimal solutions. In Table 22 we ablate the diffusion loss and contrastive loss used to train the model. The combination of both losses leads the best performance.
Graph Coloring Problem.
29


Reviews:
[{"review_id": "yWdo5nwZI5", "reviewer": "Reviewer_nBoZ", "summary": "This paper introduces a novel approach to improving generalization in machine learning models for reasoning tasks by leveraging compositional energy minimization. The authors propose learning energy landscapes over subproblems and combining them during inference to tackle more complex problems. They introduce Parallel Energy Minimization (PEM) to optimize these composed energy functions effectively. The method is evaluated on reasoning tasks like N-Queens, 3-SAT, and Graph Coloring, demonstrating superior performance compared to existing state-of-the-art approaches.", "questions": "1. If the number of particles in PEM continues to increase (e.g. beyond 1024), will this persistently lead to better results? Additionally, for the EBM results in Table 8 of the main paper, would using more particles than 128 lead to better performance, particularly on the Holme-Kim graph distribution?\n2. It seems the contrastive loss is important. How to tune the hyper-parameters of contrastive weight and MSE weight?\n3. Based on the experimental settings in appendix, the total number of steps used for gradient descent is set to 100. Does the total timesteps parameter affect the optimization results?\n4. Typo: the generation formulas for the positive and negative samples y+ and y- defined in line 128 are the same.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths：\n1. The paper is well‑written, with clear figures.\n2. By decomposing complex reasoning problems into simpler subproblems, the approach aligns well with human-like reasoning processes.\n3. The paper provides experimental validation across multiple reasoning tasks, showcasing significant improvements over baseline methods. \n\nWeaknesses:\n1. While PEM improves optimization, the requirement for multiple particles and iterations may increase computational overhead. A discussion on trade-offs between performance and computational resources would be beneficial.\n2. The evaluation is limited to combinatorial problems which can be cleanly decomposed into subproblems. How to extend it to more complex reasoning tasks？", "was_revised": true, "final_justification": "The authors provided a detailed response and additional analysis, including new experimental results on the crossword dataset, which addressed the initial concerns about generalization and applicability. Therefore, I raise my rating to 5.", "rebuttal": "We thank the reviewer for appreciating the clarity of our paper and the experimental validation. We believe that the feedback provided will help us improve the paper. We address the concerns below:\n\n---\n\n### **Q1) Increasing number of particles**\n\nWe agree that we should have provided additional evaluation on the effect of increasing the number of particles. We have scaled the experiments from Table 8, particularly those from the Holme-Kim graph distribution, up to our hardware and time limits. We report the results in the table below. In this table, we show that increasing the number of particles does lead to an improvement in the number of conflicting edges, from 59.00 on average to 55.57 in the case of the larger instances. Notice, however, that we double the number of particles for each entry of the table to obtain improved results. That is, while there is some potential for further improvement, one needs to take into account the computational cost of increasing the number of particles. We will add this additional information to the appendix of the paper.\n\n| Num. Particles | Conflicting Edges (Holme Kim [22-34 nodes]) | Conflicting Edges (Holme Kim [86-100 nodes]) |\n|---------------| -------------------------------------------- | --------------------------------------------- |\n| 128        \t| $10.60 \\pm 2.70$          \t| $59.00 \\pm 5.24$\n| 256        \t| $8.38 \\pm 3.28$           \t| $56.89 \\pm 12.75$ |\n| 512        \t| $9.68 \\pm 3.13$           \t| $54.68 \\pm 8.10$  |\n| 1024       \t| $7.04 \\pm 3.64$           \t| $57.04 \\pm 3.27$  |\n| 2048       \t| $7.79 \\pm 4.20$           \t| $55.57 \\pm 4.58$  |\n\n---\n\n### **Q2) Loss weight tuning**\n\nLoss weight hyperparameters can indeed affect the end performance. In all our experiments, we used a weight of 1.0 for MSE and 0.5 for contrastive loss, since we observed that this combination led to the best performance in all cases. In early experiments, we also tried a weight of 1.0 for both losses without observing a significant difference in performance. Generally, we recommend values for the contrastive loss in the range 0.05-0.5. A higher weight for the contrastive loss might lead to a more unstable optimization, while a lower weight might render this loss ineffective.\n\n---\n\n### **Q3) Effect of number of timesteps**\n\nThe optimization results are potentially affected by the number of timesteps, since more timesteps allow a smoother transition between states and enable more update steps. To evaluate the impact of this hyperparameter, we have trained a graph coloring model for timesteps in the range $[20, 50, 100, 200, 500, 1000]$. Below, we report the number of conflicting edges found by the model using the Holme-Kim random graph distribution. Each value is averaged over five random instances:\n\n| Num. Timesteps | Conflicting Edges (Holme-Kim \\[22–34 nodes]) | Conflicting Edges (Holme-Kim \\[86–100 nodes]) |\n| -------------- | -------------------------------------------- | --------------------------------------------- |\n| 20         \t| $6.98 \\pm 2.19$                          \t| $44.27 \\pm 7.29$                          \t|\n| 50         \t| $7.01 \\pm 2.68$                          \t| $54.40 \\pm 8.61$                          \t|\n| 100        \t| $6.71 \\pm 2.60$                          \t| $52.97 \\pm 8.07$                          \t|\n| 200        \t| $6.00 \\pm 3.83$                          \t| $50.31 \\pm 5.76$                          \t|\n| 500        \t| $6.50 \\pm 2.58$                         \t| $49.63 \\pm 3.19$                          \t|\n| 1000       \t| $6.51 \\pm 2.35$                          \t| $42.71 \\pm 3.93$                             |\n\nGenerally, we observe an improvement on average in the optimality of the solutions, going from 6.98 conflicting edges for 20 timesteps to 6.51 conflicting edges for 1000 timesteps in the smaller instances. In the case of the larger instances, we observe that with a number of timesteps equal to 20, the model does especially well, finding on average 44.27 conflicting edges. However, for the rest of the timesteps, the same trend is observed as in the smaller instances, with the number of conflicting edges decreasing from 54.40 to 42.71 as the number of timesteps increases from 50 to 1000. We can observe from this that the number of timesteps is a hyperparameter that, depending on the computational budget, can be tuned to improve the performance of the model. We will add these additional results to the appendix of the paper.\n\n---\n\n### **Q4) Typo in contrastive formulas**\n\nThank you for pointing this out. We will rewrite the formulas, differentiating the positive and negative samples in line 128. The correct formulas are:\n\n$\\tilde{y}^{+} = \\sqrt{1 - \\sigma_t} y^+ + \\sigma_t \\epsilon$ and $\\tilde{y}^{-} = \\sqrt{1 - \\sigma_t} y^- + \\sigma_t \\epsilon$, where $\\tilde{y}^{+}$ and $\\tilde{y}^{-}$ are noise-corrupted versions of the positive and negative samples $y^+$ and $y^-$, respectively."}, {"review_id": "vcW91LRgaz", "reviewer": "Reviewer_xJaU", "summary": "This paper discusses a framing for reasoning based on energy minimization. The authors propose parallel energy minimization, a method that composes an energy landscape given a new problem instance.", "questions": "See the weaknesses section. I'm excited to engage with the authors and I hope that with some iteration, this paper becomes a strong submission.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Interesting approach to reasoning.\n\n\nWeaknesses:\n- Limited benchmark datasets. The modern setting for reasoning is missing from this paper. Recent work on AI reasoning models use long form open ended math, programming, and science questions. I acknowledge the importance of non-LLM work, but the questions remains, can state-of-the-art reasoning LLMs solve the N-Queens problem or the graph coloring problem? Do we still need specialist models for this type of problem?\n- Formatting/latex issue. I printed this paper to read and mark up and none of the instances of \"fi\" print correctly. In fact, even in my PDF viewer application, the find function fails to find any instances of \"fi\". I'm extremely curious if the authors can respond to this. Is this because of a tex issue you know about? Is it the result of some formatting issue that might be a problem for the conference template? \n- Limited baselines. How do existing specialist iterative reasoning models do on these problems. I'm thinking of models like those in [7] for N-Queens and in https://proceedings.mlr.press/v202/bevilacqua23a/bevilacqua23a.pdf and https://proceedings.neurips.cc/paper/2021/file/82e9e7a12665240d13d0b928be28f230-Paper.pdf for graph coloring.", "was_revised": true, "final_justification": "The authors addressed my concerns and provided clarity and new results in response to my questions.", "rebuttal": "We thank the reviewer for finding our approach interesting and for the constructive feedback. We believe that the feedback is useful and will help us improve the paper. We discuss the concerns raised below:\n\n---\n\n### **Q1) Comparison with reasoning LLMs**\n\nWe agree that LLM reasoning shows impressive results in certain domains. While we also agree on the importance of non-LLM approaches, we want to note that the approach presented in this paper is not necessarily limited to the algorithmic domain and can be adapted to other modalities too (e.g. text, images, etc.).\n\nWe also agree that, given the results of LLMs in reasoning, a comparison would improve the paper coverage of current reasoning approaches. To this end, here we discuss the comparison of our method against reasoning LLMs on the N-Queens problem and the graph coloring problem.\n\n- **N-Queens:** We observed that all the LLMs we evaluated (see graph coloring below) have memorized at least one solution instance from 8-queens up to 16-queens. This makes the comparison difficult, as it is hard to know whether the LLMs are actually reasoning or just recalling a memorized solution.\n\n- **Graph coloring:** We evaluated using random Erdos-Renyi graphs, which guarantees that the LLMs do not have memorized solutions. We provide quantitative results below on a smaller graph with 38 nodes, 146 edges and k=3, and a larger graph with 96 nodes, 386 edges and k=3. We observe that some LLMs, such as Qwen-235B, are, in some cases, able to solve the smaller graph by backtracking fully in-context. However, once the graph size increases, the context window becomes a limitation, and the models usually output random guesses or heuristic guesses based on the degree of the nodes. These results indicate that while LLMs can reason over smaller, more tractable instances, solving larger or more complex reasoning problems remains a limitation for them as of today. We will include these additional results in the appendix of the paper.\n\n    | Model                | Conflicting Edges (38 nodes)     | Conflicting Edges (96 nodes)     |\n    |----------------------|----------------------------------|----------------------------------|\n    | Gemini 2.5 Pro       | $20.00 \\pm 14.22$                | $142.66 \\pm 6.11$                |\n    | Deepseek R1          | $4.00 \\pm 3.46$                  | $78.00 \\pm 36.38$                |\n    | Qwen 235B-A22B-2507  | $16.00 \\pm 13.06$                | $102.66 \\pm 26.10$               |\n    | PEM (P=128) (ours)   | $3.15 \\pm 2.00$                  | $45.81 \\pm 11.88$                |\n\n---\n\n### **Q2) Issue with “fi” formatting**\n\nUnfortunately, we were unable to reproduce the issue with the \"fi\" characters. We printed the paper and used different PDF viewers, and in all cases, the \"fi\" characters printed correctly. We have used the default formatting packages included in the Neurips template: `\\usepackage[utf8]{inputenc}`, `\\usepackage[T1]{fontenc}` and `\\usepackage{microtype}`. The document was compiled using PDFLaTeX. If the reviewer has any additional information on how to reproduce the issue, we are happy to update the paper to solve this issue.\n\n---\n\n### **Q3) Comparison with iterative/neural algorithmic reasoning**\n\nThank you for suggesting relevant baselines for our work. We discuss the comparison with two of the proposed works below:\n\n1. **Bansal et al.**\nWe agree that the work in [1] is relevant to our work, as both works present methods for end-to-end reasoning for a given set of problems, with the potential to be applied to larger instances. For the sake of comparison, we have trained a model for the N-Queens problem using the method in [1]. We provide a discussion of the limitations of the method compared to ours below:\n\t- **Uniqueness:** While we were able to successfully solve the 8-queens problem using the approach in [1], notice that the method, as it is, is purely deterministic, meaning that given the same input, it will always output the same solution. In the case of N-queens, where there are multiple solutions, and no input is provided, this is a limitation. To provide a comparison, while we can generate a correct solution using the method in [1] (DeepT), our approach can generate 37 distinct solutions out of 100 sampled solutions:\n\t| Model              \t| Correct Instances       \t| Unique Solutions          \t|\n\t| -----------------------| ----------------------------| ------------------------------|\n\t| DeepT              \t| 1                       \t| 1                         \t|\n\t| EBM (P=128) (ours) \t| 100                     \t| 37                        \t|\n\n\t- **Increasing computation:** while the method of [1] has the potential to solve larger instances by performing multiple iterative steps, we observe that increasing the number of steps from 30 (used during training) to 1000 led to a degeneracy in the solutions, which suggests that the model has learned suboptimal heuristics:\n\t| Num. Iterations | Size (Queens placed) |\n\t| ----------------| ----------------|\n\t| 30          \t| 8           \t|\n\t| 60          \t| 7           \t|\n\t| 100         \t| 6           \t|\n\t| 500         \t| 6           \t|\n\t| 1000        \t| 6           \t|\n\n      We will add these additional results to the appendix of the paper.\n\n2. **Bevilacqua et al.**\nThe work in [2] is closely related to ours, since both approaches try to learn iterative steps to solve a given problem and generalize to potentially larger problems. The approach to algorithmic reasoning presented in [2], however, proposes to supervise the learning of reasoning by directly supervising each algorithmic step, such that the model learns to execute a given algorithm. In other words, training requires access to a complete trace of the algorithm which the model learns to simulate. In contrast, our method relies only on input-output pairs (that is, no intermediate steps) and the reasoning is done through energy minimization, which makes both methods not directly comparable.\n\n---\n\n### References \n\n[1] Bansal et al., End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking\n\n[2] Bevilacqua et al., Neural Algorithmic Reasoning with Causal Regularisation"}, {"review_id": "vaj3AihHqW", "reviewer": "Reviewer_6UFB", "summary": "This work proposes a novel energy-based method for generating solutions to reasoning problems by learning energy landscapes of subproblems and combining them. The method combines diffusion and contrastive loss, and is evaluated against state-of-the-art methods on problems like N-queens, SAT, and graph coloring. All methods are trained and then compared in out-of-distribution setting. Proposed method significantly outperforms existing SoTA methods.", "questions": "- What are the limits on performance of increasing the number of particles?\n- What is the computational time of your method compared to other compared methods?\n- How are the N-queens decomposed into subproblems? Is your method applicable to tasks that are not decomposable?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The paper is well written, with clear formulas and descriptions of the method. The supplement contains sufficient details about implementation and experiments.\n- The training method is novel and a major improvement over existing SoTA methods, while allowing for some interpretability through plotting energy values of solutions.\n\nWeaknesses\n- Increasing the number of particles seems to always improve performance. Limits of increasing this parameter to improve performance are not explored.\n- The computational time of the method is not compared against other methods. \n- There is a minor mistake in equation 3 notation, where the expectation should mark that it's over ε ∼ N (0, I), not N(ε,0,I)\n- When solving specific problems like N-Queens, it is not entirely clear how the algorithm decomposes them into subproblems.", "was_revised": true, "final_justification": "Authors answered questions, added requested experiments, and argued well for certain design choices of this work. As it was pointed out by other reviewers, measuring the method's performance on a wider problem variety would be beneficial to this work. Additionally, the method is quite computationally intensive, although this is somewhat warranted due to the strong results. Overall, i suggest to accept this paper, but lower my score from 6 to 5 due to the previously mentioned limitations.", "rebuttal": "We thank the reviewer for appreciating the novelty of our method and the clarity of our paper. The reviewer has raised interesting points that will help us improve the paper.\nWe address the questions below:\n\n---\n\n### **Q1) Scaling the number of particles**\n\nWe have scaled the results from the graph coloring problem, particularly those from the Holme-Kim graph distribution, up to our hardware and time limits. We report the results in the table below. In this table, we show that increasing the number of particles does lead to an improvement in the number of conflicting edges, from 59.00 on average to 55.57 in the case of the larger instances. Notice, however, that we double the number of particles for each entry of the table to obtain improved results. That is, while there is some potential for further improvement of the results, one needs to take into account the computational cost of increasing the number of particles. We will add this additional information to the appendix of the paper.\n\n| Num. Particles | Conflicting Edges (Holme-Kim [22-34 nodes]) | Conflicting Edges (Holme-Kim [86-100 nodes]) |\n|---------------| -------------------------------------------- | --------------------------------------------- |\n| 128        \t| $10.60 \\pm 2.70$          | $59.00 \\pm 5.24$ |\n| 256        \t| $8.38 \\pm 3.28$           \t| $56.89 \\pm 12.75$ |\n| 512        \t| $9.68 \\pm 3.13$           \t| $54.68 \\pm 8.10$  |\n| 1024       \t| $7.04 \\pm 3.64$           \t| $57.04 \\pm 3.27$  |\n| 2048       \t| $7.79 \\pm 4.20$           \t| $55.57 \\pm 4.58$  |\n\n---\n\n### **Q2) Computational time**\nTo provide a comparison of the computational time, in the table below we report the wall clock time required to solve 25 instances of the 8-queens problem:\n\n| Model      \t| Time (s)     \t|\n| ------------------| -----------------------|\n| LwD        \t| $1.16 \\pm 0.04$  |\n| GFlowNets  \t| $1.12 \\pm 0.08$  |\n| DIFUSCO    \t| $33.46 \\pm 1.36$  |\n| FastT2T ($T_S = 1, T_G = 1$) | $1.82 \\pm 0.18$ |\n| FastT2T ($T_S = 1, T_G = 1$, GS) | $1.91 \\pm 0.19$ |\n| FastT2T ($T_S = 5, T_G = 5$) | $3.37 \\pm 1.12$ |\n| FastT2T ($T_S = 5, T_G = 5$, GS) | $7.78 \\pm 0.60$ |\n| EBM (P=1024) (ours) | $84.99 \\pm 1.61$ |\n\nIt can be observed that, compared to other methods, our approach shows a comparable overhead in terms of computational time. This is expected, taking into account that during inference, we sample from multiple composed models instead of a single monolithic model. While we show that this approach leads to better generalization, there is a trade-off in terms of computational time. Note that, since we use a diffusion objective for training, multiple methods that exist in the literature for speeding up diffusion models can be potentially applied to our method too. We will include this additional information in the appendix of the paper.\n\n---\n\n### **Q3) N-queens decomposition**\n\nThe N-queens problem can be decomposed into the constraints that a solution must satisfy, that is, no queen can be placed in the same row, column or diagonal as another queen. This can be formulated as the following constraints: each row should have exactly one queen, each column should have exactly one queen, and each diagonal should have at most one queen. If we train a model that can satisfy these two constraints:\n1. Place one queen in each row.\n2. Place at most one queen on a diagonal.\n\nWe can then arrange these two models forming a chessboard and sample from the composition to obtain a solution to the N-queens problem. In the end, model 1 is repeated 8 times (for the rows), 8 times (for the columns) and model 2 is repeated 15 times (for the diagonals). Empirically, we found that using model 1 for the 15 diagonals provided decent results too. We think including an additional visualisation of this decomposition in the paper would significantly improve the clarity for the reader.\n\nRegarding the applicability to other tasks, composition is the means by which we achieve generalization in this work. The method presented in this work is potentially also applicable to e.g. composing multiple EBMs, even if they are not components of a decomposed problem. In this work in particular, we focus on decomposable tasks to show that by learning simpler instances, we are able to combine them to solve larger or out-of-distribution instances."}, {"review_id": "QgqItV9cs2", "reviewer": "Reviewer_7Yz8", "summary": "This paper proposes Parallel Energy Minimization (PEM) a particle-based optimization strategy for sampling at inference time. The key hypothesis is that valid (correct) solutions would be assigned lower energy and invalid/incorrect solutions would have high energy. The method essentially conducts search by decomposing the problem, i.e., the energy function into more tractable sub-problems that are optimized in parallel. The authors test the effectiveness of their method on three algorithmic/compositional tasks: N-Queens, 3-SAT and graph coloring with the common theme of given a question generating an answer (final state configuration) that satisfies the problem constraints and find that PEM outperforms other sampling/search baselines.", "questions": "- Is it a reasonable/realistic assumption to know the test time complexity ahead of time? If not how should the number of particles be set for new tasks?\n- Lots of reasoning tasks have unique final answer but several correct ways  or chains to achieve them. For the tasks studied in the paper the correct configuration/state may also may not be unique -- do these have the same energy -- how does energy map to correctness?\n- Seems like the method would only work for algorithmic formalization. Would the sampling also be useful if the constraints and actions (large k in graph coloring, closest which makes the problem easier and not harder) are in text or with a large action space? Results on tasks such as BlocksWorld, crossword, and Game-24 would improve the strength of this work.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n- The paper is well-written, easy to follow, and addresses the important topic of sampling strategies dedicated to \"reasoning\" tasks that require planning and search to some extent. The method itself seems intuitive. \n-  The method achieves impressive results especially on N-queens and 3-SAT tasks along with ablations that back up the effectiveness of their design. \n\n**Weaknesses:**\n- The reasoning emphasis in the introduction, abstract, and title are misleading and not consistent with the nature of the method or tasks studied. The paper predominately focuses on algorithmic tasks with search and constraint satisfaction components. However, the aforementioned section give the impression that method can be used for textual reasoning tasks (like the ones frontier LLMs excel at) and per my understanding their method would not easily generalize to such tasks.\n- Complexity of the tasks studied is not be high enough: most of the tasks have a binary action space/constraints. See question below. Establishing that their method works on more planning or search-based benchmarks with a larger action space would improve the strength of the submission.", "was_revised": true, "final_justification": "Increased the score since I find the results from the crosswords experiment convincing", "rebuttal": "We thank the reviewer for appreciating the clarity of the paper and results reported in the paper. We believe the reviewer has made significant points that will help us make a stronger submission. We address these concerns in detail below:\n\n### **Q1) Number of particles for new tasks**\n\nWe agree that the number of particles is a hyperparameter that is not straightforward to set for new tasks. One potential option for setting the number of particles is to increase the number adaptively until convergence is achieved. For instance, one can use the energy of the sampled solution as a measure of convergence, and increase the number of particles until the energy does not change significantly. To provide an example, the following table shows the average energy of the sampled solutions for the 8-Queens problem, increasing from 2 particles to 256. It can be seen that after 32 particles, the difference in energy compared to the previous value becomes, on average, less than 0.1, which can be used as a stopping criterion.\n| Num. Particles | Energy       \t|\n| -------------- | ----------------- |\n| 2          \t| $136.62 \\pm 2.26$ |\n| 4          \t| $136.29 \\pm 1.98$ |\n| 8          \t| $133.93 \\pm 1.67$ |\n| 16         \t| $133.82 \\pm 2.05$ |\n| 32         \t| $132.47 \\pm 1.16$ |\n| 64         \t| $132.36 \\pm 0.99$ |\n| 128        \t| $132.00 \\pm 1.03$ |\n| 256        \t| $131.84 \\pm 0.97$ |\n\n### **Q2) Energy and correctness relationship**\n\nThis is a significant point in the work. The presented method is based on the assumption that, if one of the models that is being composed is not satisfied or solved, the energy of the solution will be higher, while if the model is satisfied, the energy will be lower. As a consequence of this, a correct solution, i.e., a solution where all the subproblems are satisfied, will have the lowest energy. If at least one of the subproblems is not satisfied, the corresponding energy will be higher and, as a consequence, the composed energy will also be higher. To provide quantitative results on the energy of the solutions, we use the 8-Queens problem as an example, with the table below showing the energy of the solutions for the 8-Queens problem at timestep $T = 0$. We compare the following solutions:\n1. The 92 correct solutions to the 8-Queens problem.\n2. 92 solutions with 8 queens where at least one queen is not placed in a valid position.\n3. 92 random solutions where 8 queens are placed randomly.\n\nEach of the entries reports the energy averaged over the 92 instances:\n\n| Solution Type | Energy        \t|\n| -----------------|-----------------|\n| Random   | $151.43 \\pm 3.42$ |\n| One invalid queen | $137.33 \\pm 1.86$ |\n| Correct | $130.78 \\pm 0.37$ |\n\nFrom this, we can observe that all the correct solutions of the 8-Queens problem share a similar energy, and that, moreover, this is the minimum energy if we compare it to other suboptimal solutions.\n\n\n### **Q3) Additional benchmark**\n\nRegarding the search space of the problems studied, notice that for some of the graph coloring problems, the action space is considerably large. For instance, for graph \"homer\" from the COLOR benchmark, with 1629 edges and k=13, the resulting search space is of $1629 \\times 13$ binary variables.\n\nWe agree that the method can also be applied to additional tasks to improve the strength of the work. To this end, here we provide additional evaluation on the Crosswords tasks.\n\n- **Crosswords:** \nWe evaluate using the Mini Crosswords dataset from [3]. This dataset is a collection of 5x5 crossword puzzles. Each crossword puzzle has five clues for the rows and five clues for the columns. To be able to generate the given words, we used the CrosswordsQA datasets from [4]. From this dataset, we extracted only the words that appear in the Mini Crosswords dataset, which resulted in 38113 pairs of clue words. We then trained our EBM to predict a given word given a clue. To encode the clues, we used embeddings from the text-embedding-ada-002 model from OpenAI. \nTo solve a 5x5 crossword puzzle, we compose the trained model spatially, forming the five rows and five columns of the crossword puzzle. We then sample from the composed model using PEM.\n\n  We compare our approach against the results from GPT-4 [3], including:\n  - Standard Input-Output (IO)\n  - Chain of Thought (CoT)\n  - Tree of Thoughts (ToT) [3]. \n\n  In the table below, we report the success rate at the level of letters and words. We show that, despite not exceeding current results from ToT, our approach shows competitive performance in a complex search space task, with 72% of words correctly generated, compared to 78% from ToT. Additionally, we show that our method significantly outperforms both IO and CoT with 43% of correct words versus 14% and 15% respectively. This result is especially relevant, taking into account that GPT-4 has been trained on a large corpus of text data and contains significantly more semantic knowledge than our model.\n\n  | Model      \t| Letter Success Rate | Word Success Rate |\n  | ---------------| -------------------- | ------------------ |\n  | IO                 | 38.7               | 14.0               |\n  | CoT               |  40.6              | 15.6               |\n  | ToT        \t| 78.0             \t| 60.0           \t|\n  | EBM (P=256) (ours) |  72.0            \t|  43.0              \t|"}]

Author rebuttal:

Meta-review: This paper introduces a novel and intuitive method for improving generalization on reasoning tasks by composing learned energy landscapes of smaller subproblems. The reviewers were in strong agreement about the quality and contributions of the work.

Summary of Strengths:

* The proposed approach is novel, well-motivated, and aligns well with compositional reasoning (nBoZ, 6UFB, 7Yz8).

* The paper is well-written, clear, and easy to follow (nBoZ, 6UFB, 7Yz8).

* The method demonstrates strong empirical performance, showing significant improvements over state-of-the-art methods on a variety of combinatorial tasks (nBoZ, 6UFB, 7Yz8).

* The authors provided a thorough rebuttal with new experiments that addressed nearly all reviewer concerns, including new baselines (xJaU), application to a new domain (crosswords, per 7Yz8's suggestion), and analysis of computational trade-offs (nBoZ, 6UFB).

The reviewers unanimously recommend acceptance. The core idea is technically solid, and the experimental results are impressive. Initial concerns regarding the scope of evaluation, comparison to modern baselines like LLMs, and computational cost were effectively addressed by the authors in their detailed rebuttal, leading all reviewers to either confirm or raise their scores. This paper presents a valuable contribution to the field of reasoning and generalization.
Decision: Accept (spotlight)