Constraint-Guided Prediction Refinement via Deterministic Diffusion Trajectories

2025-06-15

AI Summary

Existing methods for machine learning tasks requiring hard constraints—such as physical conservation laws, tabular column dependencies, or Kirchhoff’s laws—are often domain-specific or rely on strong assumptions of constraint convexity or linearity. Method: This paper proposes a general, model-agnostic, post-hoc prediction refinement framework that couples deterministic diffusion trajectories (based on DDIM) with gradient-based constraint correction. It supports non-convex, nonlinear equality constraints without modifying the base model or presupposing constraint structure. The framework jointly integrates constraint-gradient-guided optimization and learned prior modeling, ensuring lightweight computation and plug-and-play deployment. Contribution/Results: Evaluated on tabular adversarial attack generation and AC power flow prediction, the method significantly improves constraint satisfaction rates and prediction accuracy. It demonstrates broad applicability across diverse constraint types and domains, validating both its generality and empirical effectiveness.

Technology Category

• Constraint Satisfaction and Optimization: Constraint Learning and Acquisition
• Search and Optimization: Learning to Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

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.

Problem

• Ensures machine learning outputs satisfy hard constraints
• Handles non-convex and nonlinear equality constraints
• Improves constraint satisfaction and performance across domains

Innovation

• Deterministic diffusion trajectory for refinement
• Constraint gradient corrections for guidance
• Model-agnostic post hoc application