Convex Compositional Reasoning Models

2026-05-22

AI Summary

Combinatorial inference is notoriously difficult to optimize due to the non-convexity of its energy landscape. This work proposes a convex combinatorial energy minimization framework that, for the first time, integrates input-convex neural networks (ICNNs) into combinatorial energy models to parameterize local factor energies. By optimizing the global objective over a tight convex relaxation of the feasible set, the approach guarantees overall convexity, enabling deterministic first-order optimization. The method employs a two-stage training strategy combining contrastive learning with an end-to-end unrolled solver. Remarkably, the model trained solely on small-scale or single-scale problems exhibits strong zero-shot generalization to substantially larger instances, significantly enhancing scalability and transferability.

Abstract

Compositional energy-based models can generalize to larger combinatorial reasoning problems by reusing a learned factor energy across many local constraints. In our paper, we show that a key bottleneck in compositional reasoning is not composition itself, but the non-convex geometry of the learned energy landscape. To solve this problem, we introduce Convex Compositional Energy Minimization (CCEM), a framework that parameterizes each factor with an input-convex neural network and optimizes the composed energy over a tight convex relaxation of the feasible set. Because convexity is preserved under summation, the global relaxed objective remains convex, enabling deterministic projected first-order optimization. CCEM is trained in two stages: factor-level contrastive learning to shape local energy basins, followed by end-to-end refinement through an unrolled projected solver. Our experiments show that our models trained on small subproblems or a single problem size transfer to larger instances without retraining.

Problem

• compositional reasoning
• non-convex energy landscape
• convexity
• energy-based models
• combinatorial optimization

Innovation

• Convex Compositional Energy Minimization
• input-convex neural network
• compositional reasoning
• convex relaxation
• energy-based model