Constraint-Guided Prediction Refinement via Deterministic Diffusion Trajectories

TL;DR

This paper introduces a general framework using denoising diffusion models to refine predictions in machine learning tasks with complex, nonlinear constraints, improving accuracy and constraint satisfaction.

Contribution

It proposes a novel, flexible diffusion-based refinement method that handles non-convex constraints and can be applied to any base model without domain-specific modifications.

Findings

Improves constraint satisfaction in tabular data adversarial attacks.

Enhances power flow prediction accuracy under Kirchhoff's laws.

Remains lightweight and model-agnostic across tasks.

Abstract

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.