Technical Summary: Convex Compositional Energy Minimization (CCEM)

Problem Statement

Compositional reasoning problems (e.g., N-Queens, Graph Coloring, SAT) involve finding a solution that satisfies a large number of local constraints. Compositional Energy-Based Models (EBMs) address this by learning a local energy function for a single constraint and reusing it across the entire problem instance by summing these local energies.

However, the authors identify a critical bottleneck: non-convex geometry. When local factors are parameterized by unconstrained neural networks, their sum creates a complex, non-convex energy landscape. This landscape often contains "spurious stable local minima"—invalid solutions that are locally optimal but globally incorrect. Consequently, inference requires expensive, stochastic sampling methods (e.g., long diffusion chains or Parallel Energy Minimization with many particles) to escape these local basins, rather than relying on deterministic optimization.

Methodology: Convex Compositional Energy Minimization (CCEM)

The paper proposes CCEM, a framework that ensures the global energy landscape remains convex by enforcing convexity at the factor level.

1. Convex Factor Parameterization

Instead of using unconstrained Multi-Layer Perceptrons (MLPs), CCEM parameterizes each local factor $f_\theta$ using Input-Convex Neural Networks (ICNNs) or Partially Input-Convex Neural Networks (PICNNs).

• Structure: The network is designed such that the energy is convex with respect to the decision variables (the "scope") while remaining flexible with respect to the problem context.
• Mechanism: By using non-negative weights and non-decreasing convex activations, the network guarantees convexity in the scope variables.
• Composition: Since the sum of convex functions with non-negative weights is convex, the global composed energy $E_\theta(x, y) = \sum w_k f_\theta(y_{S_k}; c_k)$ remains convex over the feasible set.

2. Inference on Tight Convex Relaxations

Inference is performed by minimizing the global energy over a tight convex relaxation of the discrete feasible set, rather than in ambient Euclidean space.

• Relaxations: Examples include the Birkhoff polytope for N-Queens and simplex-based relaxations for Graph Coloring.
• Optimization: The model uses deterministic projected first-order optimization (Projected Gradient Descent or Projected Adam). Because the objective is convex, local optimality implies global optimality over the relaxation, eliminating the need for particle-based exploration to escape local minima.
• Decoding: The continuous relaxed solution is decoded into a discrete solution (e.g., via greedy row argmax).

3. Two-Stage Training Pipeline

The training process consists of two distinct stages:

1. Factor-Level Contrastive Pretraining: The local factors are trained using a contrastive InfoNCE objective on clean samples. Positive samples are valid local patterns (vertices of the convex hull), and negative samples are invalid patterns. This shapes the local energy basins without requiring global composition.
2. End-to-End Refinement: The composed energy is refined by backpropagating through an unrolled projected solver. The loss function combines regression, an energy-margin term, and hard-negative mining over the solver trajectory. This aligns the learned energy with the behavior of the actual inference procedure while preserving the convex structure.

Key Contributions

1. Identification of the Bottleneck: The authors demonstrate that spurious local minima in compositional EBMs arise from the non-convex composition of factors, not merely from the choice of sampler. They prove (Proposition 3.1) that unconstrained composition can admit stable, invalid local minima.
2. Convex Composition Framework: They propose a novel framework using ICNN/PICNN factors where non-negative summation preserves the convexity of the global objective.
3. Deterministic Inference Pipeline: They replace diffusion sampling and particle ensembles with a deterministic projected first-order solver operating on tight convex relaxations.
4. Theoretical Guarantees:
   • Absence of Spurious Minima: Convex composition rules out spurious relaxed local minima.
   • Sufficiency of First-Order Supervision: Theorem 3.3 shows that for convex energies, satisfying the first-order optimality condition at the target solution is sufficient to certify global optimality. This removes the need for auxiliary landscape-shaping losses required in non-convex settings.
5. Empirical Validation: The method is evaluated on N-Queens, Graph Coloring, and 3-SAT, demonstrating that models trained on small subproblems generalize to larger instances without retraining.

Experimental Results

The paper evaluates CCEM against state-of-the-art baselines, including the original compositional EBM framework (using diffusion and Parallel Energy Minimization) and other neural combinatorial optimizers.

• N-Queens: Trained on a single $8 \times 8$ instance, CCEM achieves 100% correct instances on 100 sampled boards with 8 queens placed (perfect score), outperforming the EBM-Diff baseline (97% correct) and significantly outperforming non-EBM baselines. It achieves this with a single trajectory of projected Adam ($S=1$) or a small number of warm starts, whereas the baseline requires a large particle ensemble ($P=1024$) and long diffusion chains.
• Graph Coloring: CCEM significantly reduces the number of conflicts (errors) compared to GNNs, RL approaches, and the EBM baseline across various graph distributions (Erdos-Renyi, Holme-Kim, etc.). For example, on Erdos-Renyi graphs, CCEM achieves a conflict rate of ~2.2 compared to ~46.8 for GCN.
• Generalization: The models trained on small problem sizes (e.g., $N=8$ for N-Queens) successfully transfer to larger instances ($N=10$) by simply re-summing the factors, without retraining.
• Ablation: The study shows that the compositional structure itself provides a strong inductive bias; even with simplified loss functions in the refinement stage, the method remains highly effective.

Significance and Claims

The paper claims that convexity is the key to scalable and reliable compositional reasoning. By shifting the focus from "escaping" a rugged landscape via stochastic sampling to "constructing" a smooth, globally convex landscape via architectural constraints, CCEM offers a more efficient inference procedure.

The authors argue that the difficulty in compositional reasoning is not inherent to the problem structure but to the parameterization of local energy factors. By enforcing convexity, they transform the inference task into a well-behaved convex optimization problem. This allows for:

• Deterministic Inference: Eliminating the need for random seeds or particle ensembles.
• Theoretical Certifiability: Global optimality can be certified via first-order conditions.
• Scalability: The ability to solve larger problems by increasing solver precision (more iterations) rather than retraining or increasing sampling complexity.

The paper concludes that while convexity may be restrictive (potentially creating flat regions), it effectively removes the pathological geometry that plagues current compositional EBMs, offering a promising direction for neural reasoning.