Convex Compositional Reasoning Models

arXiv:2605.23395 · cs.LG · Submitted 2026-05-22 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Hi, I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: We're looking at the paper: "Convex Compositional Reasoning Models".

Jane: Here's the gist of the paper. The gist: Convex Compositional Energy Minimization (CCEM) introduces a framework that uses input-convex neural networks to parameterize local factors, ensuring that non-negative factor summation preserves convexity of the global relaxed objective, thereby enabling deterministic first-order optimization for compositional reasoning.

Tom: How it works Compositional energy-based models are used to solve reasoning problems by defining an energy function where valid solutions should have low energy and invalid ones higher energy.

Jane: The core idea is to learn a local rule for a small constraint and reuse it across many constraints in a larger instance.

Tom: The composed energy is defined as Eθ(x, y) = Σk=1 wkfθ(ySk; ck), where each factor fθ scores a local scope Sk under context ck.

Paper summary: Tom: So we're looking at this paper called "Convex Compositional Reasoning Models," and it’s about taking these compositional energy models, which are great for reasoning problems like N-Queens or graph coloring, and making them work better on bigger problems.

Jane: That’s right. The main idea here is that the problem isn't actually putting all the pieces together—the composition itself isn't the hard part—but how the learned energy landscape looks, which can be really non-convex.

Lu: Exactly. The authors show that a key bottleneck in compositional reasoning is this geometry of the learned energy landscape, not just the process of composing constraints. They introduce this framework called Convex Compositional Energy Minimization, or CCEM, to fix that.

Meng: So what does CCEM actually do? I need to know if it’s just another way to tune a neural network or if it fixes the fundamental mathematical issue with these models being non-convex.

Tom: It fixes the math by parameterizing each factor with an input-convex neural network, or ICNN. This ensures that when you sum up those factors, which is what you do in compositional energy minimization, the resulting global relaxed objective remains convex.

Jane: That’s a big deal because convexity is preserved under summation, meaning the global relaxed objective stays convex even with these local rules being added together.

Lalam: The paper claims that this convexity enables deterministic projected first-order optimization. This means instead of using some kind of sampling method to find the answer, you can use a direct minimization approach.

Meng: Deterministic optimization sounds promising for deployment because it avoids the randomness that often comes with sampling methods when you need a reliable result at scale. What’s the training process like then?

Tom: They train it in two stages first: factor-level contrastive learning, which is used to shape those local energy basins. Then they do an end-to-end refinement through an unrolled projected solver to get it ready for actual inference.

Lu: The theory here is pretty strong too. Proposition three point one shows that unconstrained compositional energies can have stable but invalid local minima, but by enforcing convexity at the factor level, the global geometry is preserved <ref:2605.23395#pg1>.

Jane: That’s interesting because it means you don't have to worry about those spurious local minima messing up your search for the best solution on the larger problem.

Tom: They give a condition for validity in Lemma A.two: if the learned factors uniformly approximate certificate factors with error at most epsilon, and two times epsilon times K of K is less than one, then every global minimizer of that composed energy over Yx is valid <ref:2605.23395#pg1>.

Meng: So it’s not just theoretical; they’ve given us a condition for when we can actually trust the model to find a correct answer. How does this affect deployment costs or complexity?

Jane: The inference pipeline replaces particle-based exploration with deterministic first-order minimization directly on that tight convex relaxation of the feasible set. They use projected gradient descent with a step size of one over L for that optimization.

Paper summary: Lalam: They also mention using S random simplex warm-starts for inference, and for harder cases, they can use Parallel Energy Minimization by replacing the solver with a P-particle ensemble resampled using w(t)i proportional to exp minus theta of y at t i.

Tom: The empirical results show that models trained on small subproblems or even just a single problem size transfer to much larger instances without needing any retraining.

Jane: They tested this on structured reasoning benchmarks, and the CCEM approach can match or even improve particle-based compositional energy-based models while using a simpler inference procedure.

Meng: That matching performance with less complex inference sounds like something we could actually put into a production system if the convergence properties hold up under real traffic loads.

Lu: The paper also notes that the main gain comes from the compositional energy structure itself, like row, column, and diagonal factors which already impose a strong inductive bias.

Tom: The authors flag one limitation: convexity can be too restrictive and produce flat regions or fractional relaxations when dealing with multimodal discrete solution spaces.

Jane: So while it’s great for structured problems where the geometry is well-behaved, it might struggle if the answer space has many distinct, widely separated modes that aren't smoothly connected.

Meng: That sounds like something we’d need to watch closely when we try to apply this to more open-ended reasoning tasks outside of pure combinatorial constraints.

Lalam: Overall, what this paper does is shift compositional reasoning from a search problem involving diffusion sampling into one that uses deterministic first-order minimization on a convex relaxation.

Tom: It’s a major shift because it removes the source of spurious local minima at the parameterization level, and they provide a way to certify global optimality through that first-order condition at the target solution.

Jane: The title "Convex Compositional Reasoning Models" points to how they solved this specific bottleneck by making sure the underlying math supports a reliable optimization path.

Lu: This framework allows for global optimality certification through a first-order condition, which is a major step away from prior work that relied on auxiliary landscape-shaping losses to rule out spurious minima.

Tom: So, to wrap up on "Convex Compositional Reasoning Models," this framework lets us use input-convex neural networks for factors so that the summation preserves convexity of the relaxed global objective.

Jane: That’s the core mechanism that enables deterministic projected first-order optimization instead of relying on slower sampling techniques.

Lu: The theoretical guarantees show that this convex composition rules out spurious relaxed local minima, and first-order supervision is sufficient to certify global optimality in this setting.

Meng: For deployment, it means we can potentially achieve better performance than particle-based models using a simpler inference procedure for structured problems.

Lalam: The paper demonstrates strong generalization from small training instances to larger combinatorial problems without needing retraining, which is a very useful property.

Conclusion: Tom: So, we’re wrapping up on "Convex Compositional Reasoning Models," and what I gotta say is this paper tackles the whole issue of how these compositional models get messy when they try to solve big problems.

Jane: Yeah, it’s about taking those energy-based models—the ones that use rules to score solutions—and making sure the math stays clean, especially when you stack a bunch of those rules on top of each other.

Tom: Right. The authors are showing how they fix that by using something called Convex Compositional Energy Minimization, or CCEM. It’s about parameterizing each local rule with an input-convex neural network so that the total energy function stays convex under summation.

Lu: Exactly what Tom said; the core idea is to use input-convex neural networks for those factors, which makes sure that when you sum them up across all those constraints, you don’t get a wild non-convex landscape with hidden traps.

Jane: So what does that mean for us? It means instead of just hoping the model finds a good answer by sampling randomly, we can use direct optimization on this relaxed, convex version of the problem.

Meng: From my side, it means if you can solve it deterministically without needing thousands of random trials just to get one good result, that’s a huge win for deployment at scale. It cuts down on the compute time per answer significantly.

Lalam: I see what Meng is getting; the method allows us to replace those slow diffusion sampling techniques with projected first-order minimization, which gives us a much more certain path to finding the solution.

Tom: And theoretically, they’re showing that this convexity at the factor level actually means we can certify that any solution we find is globally optimal within that relaxation.

Jane: That's a big deal because it shifts the problem from "maybe this is good enough" to "this is mathematically guaranteed to be the best thing possible in this simplified space."

Lu: The authors show they’ve trained these factors using contrastive learning and then refined them with an unrolled solver, which really shows how you can build a reliable system piece by piece.

Meng: But I gotta ask, what about when the real problem is really multimodal? If the true answer space has a bunch of totally separate islands that don't connect smoothly, does this convex relaxation still work?

Jane: That’s actually a limitation they admit—that convexity can be too restrictive and might create flat regions or fractional answers if the discrete solution space is really complex.

Tom: It suggests that for many structured reasoning problems, like N-Queens, the geometry is good enough to let this work. But it tells us that if we hit a truly multimodal space, we still might need more sophisticated ways to handle those gaps in the energy landscape.

Jane: So, in short, they’ve turned the challenge of non-convex reasoning into a problem of controlling input convexity at the factor level so we can use fast, deterministic optimization instead of slow sampling.

Lu: It really shows that for scaling these models up to larger instances, you don't necessarily need to invent a totally new architecture; sometimes you just need to fix the geometry of what you’re composing.

Meir Roketlishvili, Semyon Semenov, Maksim Bobrin, Viktor Kovalchuk, Albert Baichorov, Abduragim Shtanchaev, Fakhri Karray

Mohamed bin Zayed University of Artificial Intelligence

cs.LG

Submitted: 2026-05-22

Updated: 2026-10-18

License: http://creativecommons.org/licenses/by/4.0/

The gist: The gist: Convex Compositional Energy Minimization (CCEM) introduces a framework that uses input-convex neural networks to parameterize local factors, ensuring that non-negative factor summation

Key concepts

Compositional Energy Models: These models use an energy function to guide reasoning, where low energy corresponds to valid solutions. The core idea is reusing learned local rules across many constraints in a large problem instance. The challenge addressed is that the way these learned rules combine can create complex, non-convex landscapes.
Input-Convex Neural Network (ICNN): This technique parameterizes each local factor using a neural network structure specifically designed to be input-convex. This mathematical constraint guarantees that when these factors are summed up, the resulting global objective function maintains convexity, which is crucial for stable optimization.
Projected First-Order Optimization: Instead of using sampling methods like diffusion, CCEM uses deterministic first-order minimization directly on a convex relaxation of the feasible set. This involves solving an optimization problem over a simplified geometric shape (like a polytope) using projected gradient descent, leading to faster and more reliable inference.

Terminology

Contrastive Learning — 1× in this paper · explained in 8 episodes
Convex Compositional Energy Minimization — 1× in this paper

Summary

The gist: Convex Compositional Energy Minimization (CCEM) introduces a framework that uses input-convex neural networks to parameterize local factors, ensuring that non-negative factor summation preserves convexity of the global relaxed objective, thereby enabling deterministic first-order optimization for compositional reasoning.

How it works

Compositional energy-based models are used to solve reasoning problems by defining an energy function where valid solutions should have low energy and invalid ones higher energy. The core idea is to learn a local rule for a small constraint and reuse it across many constraints in a larger instance. The composed energy is defined as Eθ(x, y) = Σk=1 wkfθ(ySk; ck), where each factor fθ scores a local scope Sk under context ck.

The framework addresses the non-convexity issue by parameterizing each factor with an input-convex neural network (ICNN). This ensures that nonnegative weights preserve convexity of the relaxed global objective. The training pipeline consists of two stages:

Factor-level contrastive learning to shape local energy basins. This involves a loss like LNCE(θ) which uses factor-specific proposals as negatives.
End-to-end refinement through an unrolled projected solver. This involves backpropagating through a solver trajectory to align the energy with the actual inference procedure.

Theoretical Guarantees of Convexity

The central theoretical benefit is that convex composition cannot introduce spurious local minima over a convex relaxation. Proposition 3.1 demonstrates that unconstrained compositional energies can have stable but invalid local minima. By enforcing convexity at the factor level, the global geometry is preserved.

Lemma A.2 provides a condition for validity: if the learned factors uniformly approximate certificate factors with error at most ε, and 2εXKk=1 < γ, then every global minimizer of Eθ(x, ·) over Yx is valid. Furthermore, Theorem A.3 shows that if the first-order optimality condition ⟨∇yEθ(x, y⋆), y − y⋆⟩ ≥ 0 holds for all y ∈ Y, then y⋆ is a global minimizer of the energy.

Inference and Optimization Pipeline

The inference procedure replaces particle-based exploration with deterministic first-order minimization. The optimization is performed directly on the convex relaxation where feasibility is enforced by projection.

Key components of the inference are:

. Projected first-order solver:

The framework uses a tight convex relaxation of the feasible set, such as the Birkhoff polytope for N-Queens or simplex-based relaxations for coloring. The optimization is performed over this relaxation using projected gradient descent with step size 1/L.

. Multi-start and Sampling:

The method uses S random simplex warm-starts for inference. For harder regimes, Parallel Energy Minimization (PEM) can be used by replacing the solver with a P-particle ensemble that resamples using w(t)i ∝ exp−Eθ(y(t)i).

Empirical Results

Experiments on structured reasoning benchmarks show that CCEM can match or improve particle-based compositional EBMs while using a simpler inference procedure. For N-Queens, the CCEM approach achieved SL+Adam 100 8.00±0.00 in Table 1. The results indicate that the convex factor energy yields a decisive basin on the Birkhoff polytope rather than a smeared heatmap that requires temperature-style decoding. Compositional generalization to larger instances is demonstrated by training once on N=8 and evaluating on N ∈ [7, 8, 9, 10] without retraining.

Limitations and Broader Impact

The primary limitation noted is that convexity can be too restrictive, producing flat regions, fractional relaxations, or poor fits to multimodal discrete solution spaces. However, the paper concludes that composition turns generalization into a convex program: harder tasks mainly require more solver precision, not escape from pathological geometry. This suggests that enforcing favorable energy geometry at the factor level is a promising direction for scalable and reliable neural reasoning.

The paper also provides ablation studies showing that the main gain comes from the compositional energy structure itself: row, column, and diagonal factors already impose a strong inductive bias. The computational cost scales linearly with the number of projected solver steps and warm-starts at inference time.

The paper concludes that enforcing favorable energy geometry at the factor level is a promising direction for scalable and reliable neural reasoning. The framework allows for deterministic inference by replacing diffusion sampling with projected first-order minimization. This approach is significant because it removes the source of spurious local minima at the parameterization level. The framework enables global optimality certification through a first-order condition at the target solution. This is a major shift from prior work that relied on auxiliary landscape-shaping losses to rule out spurious minima. The framework shows strong generalization from small training instances to larger combinatorial problems. This suggests that the model learns inside a function class whose geometry makes local optimality certificates global. The framework allows for deterministic inference by replacing diffusion sampling with projected first-order minimization. This approach is significant because it removes the source of spurious local minima at the parameterization level. The framework enables global optimality certification through a first-order condition at the target solution. This is a major shift from prior work that relied on auxiliary landscape-shaping losses to rule out spurious minima. The framework shows strong generalization from small training instances to larger combinatorial problems. This suggests that the model learns inside a function class whose geometry makes local optimality certificates global. The framework allows for deterministic inference by replacing diffusion sampling with projected first-order minimization. This approach is significant because it removes the source of spurious local minima at the parameterization level.

Improvements for AI systems

Convex Compositional Energy Minimization (CCEM) enables deterministic first-order optimization instead of particle-based methods by parameterizing each factor with an input-convex neural network, which preserves convexity under non-negative summation across all constraints. This allows for inference to proceed over a tight convex relaxation, replacing long diffusion chains or particle ensembles with projected gradient descent.
Elimination of spurious local minima in the composed energy landscape ensures that the optimization converges to valid solutions even when dealing with large combinatorial problems where unconstrained models might otherwise get stuck in spurious stable minima. This removes the need for auxiliary landscape-shaping losses, simplifying the training pipeline significantly.
Compositional generalisation to larger instances is achieved by training a factor energy on small subproblems and reusing it across larger ones, as the same factor network can therefore be trained on small subproblems and applied to larger instances simply by adding more terms to the sum. This allows models trained on N=8 boards to generalize reliably to N=10 without retraining.
Inference is performed via a projected first-order solver using an unrolled trajectory, which align[s] the energy with the behavior of the actual inference procedure while preserving the convex structure of the model. This method provides a deterministic and differentiable inference procedure compatible with standard solvers, making reasoning more reliable than diffusion-based sampling.
First-order supervision is sufficient to certify global optimality in the relaxed problem, as if y⋆ ∈ Y satisfies ⟨∇yEθ(x, y⋆), y − y⋆⟩ ≥ 0 ∀y ∈ Y, then it is a global minimizer of the energy. This theoretical guarantee means auxiliary losses are not required to rule out spurious relaxed minima.

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.