Constraint-Guided Prediction Refinement via Deterministic Diffusion Trajectories
This addresses the challenge of applying constraint-aware refinement to non-convex and nonlinear constraints in real-world ML tasks, offering a post hoc solution that is incremental over existing diffusion methods.
The paper tackles the problem of refining machine learning predictions to satisfy hard constraints, such as physical laws or structured dependencies, by proposing a general-purpose framework using denoising diffusion implicit models (DDIMs). The result is improved constraint satisfaction and performance in domains like tabular data and AC power flow prediction, with the method being lightweight and model-agnostic.
Many real-world machine learning tasks require outputs that satisfy hard constraints, such as physical conservation laws, structured dependencies in graphs, or column-level relationships in tabular data. Existing approaches rely either on domain-specific architectures and losses or on strong assumptions on the constraint space, restricting their applicability to linear or convex constraints. We propose a general-purpose framework for constraint-aware refinement that leverages denoising diffusion implicit models (DDIMs). Starting from a coarse prediction, our method iteratively refines it through a deterministic diffusion trajectory guided by a learned prior and augmented by constraint gradient corrections. The approach accommodates a wide class of non-convex and nonlinear equality constraints and can be applied post hoc to any base model. We demonstrate the method in two representative domains: constrained adversarial attack generation on tabular data with column-level dependencies and in AC power flow prediction under Kirchhoff's laws. Across both settings, our diffusion-guided refinement improves both constraint satisfaction and performance while remaining lightweight and model-agnostic.