The Theory and Practice of MAP Inference over Non-Convex Constraints

In many safety-critical settings, probabilistic ML systems have to make predictions subject to algebraic constraints, e.g., predicting the most likely trajectory that does not cross obstacles. These real-world constraints are rarely convex, nor the densities considered are (log-)concave. This makes computing this constrained maximum a posteriori (MAP) prediction efficiently and reliably extremely challenging. In this paper, we first investigate under which conditions we can perform constrained MAP inference over continuous variables exactly and efficiently and devise a scalable message-passing algorithm for this tractable fragment. Then, we devise a general constrained MAP strategy that interleaves partitioning the domain into convex feasible regions with numerical constrained optimization. We evaluate both methods on synthetic and real-world benchmarks, showing our approaches outperform constraint-agnostic baselines, and scale to complex densities intractable for SoTA exact solvers.

💡 Research Summary

The paper tackles the problem of performing constrained maximum‑a‑posteriori (MAP) inference when both the algebraic constraints and the probability density are non‑convex and non‑log‑concave, a setting that arises in safety‑critical applications such as robot trajectory planning, physical simulation, and process control. Traditional MAP methods assume either discrete variables, convex constraints, or log‑concave densities, which guarantee that any local optimum is global. In the general case, however, MAP under non‑convex constraints is NP‑hard, and existing exact solvers (e.g., CAD‑based methods, OMT extensions, OptiMathSAT, CDCL‑OCAC) quickly become intractable.

The authors contribute two complementary approaches.
1. Exact message‑passing for a tractable fragment (C1).
They identify a fragment where the constraint formula (expressed in SMT‑LRA) and the density both factorize over a graph whose treewidth is one (i.e., a tree or forest). Moreover, the density must belong to a function family Ω that satisfies three “tractable MAP conditions” (TMC): closure under multiplication, tractable symbolic supremum over a bounded interval for any bivariate component, and tractable pointwise maximum for any pair of univariate components. Two concrete families are introduced: Ω_PP (piecewise polynomials where each piece factorizes into a product of two univariate polynomials) and Ω_PEP (piecewise exponentiated polynomials, e.g., independent Gaussian components). Under these conditions, the global MAP problem can be solved exactly by a fixed‑parameter tractable message‑passing algorithm (MPMAP). The algorithm proceeds from leaves to root, computing for each variable a message that encodes the optimal value of the sub‑tree conditioned on that variable. Because each message requires only the operations guaranteed by TMC, the overall runtime is linear in the number of variables and exponential only in the bounded tree depth, which is negligible for typical problems. Theorem 4.5 formally proves that if the global graph has treewidth one, bounded diameter, and the density satisfies TMC, then MAP can be solved exactly.

2. General approximation via convex partitioning (C2).
For problems that do not satisfy the tree‑structure assumption, the authors propose PAMAP, a framework that partitions the non‑convex feasible region defined by the SMT constraints into a collection of convex polytopes using an SMT solver. Each convex region is then tackled independently with off‑the‑shelf continuous optimizers (Adam, Basin‑Hopping, SHGO, etc.). The solutions from all regions provide lower bounds, while a hierarchy of Lasserre relaxations or sum‑of‑squares (SoS) programs yields provable upper bounds for polynomial densities. By iteratively refining the partition and updating bounds, PAMAP converges to a high‑quality approximation with explicit guarantees. The method also leverages the Boolean structure of the constraints to prune infeasible regions early, dramatically reducing the search space.

Experimental evaluation.
The authors evaluate both MPMAP and PAMAP on synthetic 2‑D/3‑D benchmarks and three real‑world domains: (i) robot trajectory planning with obstacle avoidance, (ii) a physics‑based simulation of particle dynamics, and (iii) a fair‑prediction task with complex logical constraints. In the tree‑structured cases, MPMAP recovers the exact global optimum and outperforms exact solvers (OptiMathSAT, CDCL‑OCAC) by a factor of 2–5 in runtime. In the general cases, PAMAP achieves higher objective values than constraint‑agnostic baselines (penalty‑augmented Adam) and reduces runtime by 30–50 %. Notably, on a high‑dimensional trajectory problem where CDCL‑OCAC timed out after one hour, PAMAP found a feasible solution within ten minutes, with a 7 % improvement over the best baseline. The experiments also demonstrate that the bound‑tightening procedure of PAMAP yields provable gaps that shrink rapidly as the partition is refined.

Discussion and implications.
The paper clarifies the theoretical boundary between tractable and intractable MAP under non‑convex constraints, showing that tree‑width one combined with TMC‑satisfying densities yields a polynomial‑time exact algorithm. It also argues that weighted model integration (WMI) and MAP are incomparable: functions tractable for WMI may violate TMC, and vice‑versa. For problems lacking a tree structure, the convex‑partition strategy provides a practical compromise, offering both empirical performance and theoretical guarantees. The authors suggest future work on hybrid schemes that apply MPMAP on tree‑like subgraphs while using PAMAP on the remaining parts, and on extending the function families that satisfy TMC.

Overall, the work makes a significant contribution by (i) identifying a new tractable fragment of constrained MAP, (ii) delivering a concrete message‑passing algorithm with provable optimality, and (iii) presenting a scalable, bound‑aware approximation framework for the general case. These advances open the door to reliable, real‑time MAP inference in safety‑critical AI systems where non‑convex constraints are the norm.