Technical Summary: Enforcing Constraints in Generative Sampling via Adaptive Correction Scheduling

Problem Statement

Generative models are increasingly applied in domains requiring outputs to satisfy hard constraints, such as trajectory synthesis, robot planning, and structured generation. In these settings, samples must lie on specific manifolds defined by geometric, physical, or task-specific structures. The standard approach involves running the generative process in the ambient space and enforcing feasibility via a projection operator.

Current practice typically relies on two binary strategies for timing this projection:

1. Terminal Correction: Projection is applied only once at the end of the sampling rollout.
2. Stepwise Correction: Projection is applied after every generative update.

The paper argues that this binary framing overlooks a fundamental issue: generative updates are state-dependent. Once a rollout leaves the constraint set, subsequent updates are computed from states that a constrained trajectory would never visit. While a later projection can restore feasibility at the endpoint, it cannot undo the distributional shift induced by these off-manifold intermediate states. Consequently, terminally corrected samples are often feasible but dynamically inconsistent with the intended sampling dynamics. The core problem is reframed as a correction scheduling problem: how to allocate a limited budget of projection operations across a rollout to preserve trajectory consistency, rather than merely ensuring endpoint feasibility.

Methodology: Adaptive Correction Scheduling

The authors propose Adaptive Correction Scheduling, a state-dependent policy that treats constraint enforcement as a resource allocation problem. The method operates online without requiring gradients through the projection operator or retraining the generative model.

Core Mechanism

The scheduler allocates a fixed budget $B$ of projection calls over a rollout of length $T$. At each step $t$, the scheduler:

1. Proposes a state update $\tilde{x}_{t+1} = \Phi_h(x_t)$ using the learned generative operator.
2. Computes a one-step constraint defect $s_t = d(\tilde{x}_{t+1}, \mathcal{M})$, measuring the distance of the proposed state from the constraint set $\mathcal{M}$.
3. Decides whether to project based on a budget-aware threshold $\lambda_{t,b}$, where $b$ is the remaining budget.

The update rule is:
$$x_{t+1} = \begin{cases} \Pi(\tilde{x}_{t+1}) & \text{if } b > 0 \text{ and } s_t \ge \lambda_{t,b} \\ \tilde{x}_{t+1} & \text{otherwise} \end{cases}$$

Threshold Construction

The thresholds $\lambda_{t,b}$ are precomputed using held-out calibration rollouts. They represent the marginal value of spending a correction at time $t$ with $b$ corrections remaining. Specifically, $\lambda_{t,b}$ is set to a quantile of the future defect distribution, ensuring that the scheduler reserves budget for high-defect events later in the rollout. This prevents the "failure mode" of fixed-threshold rules, which might exhaust the budget on moderate early defects, leaving the scheduler unable to correct larger, more critical deviations later.

Special Cases

This framework unifies existing strategies as limiting cases:

• Terminal Correction: Achieved by setting all thresholds to $+\infty$ (no projection until the end).
• Stepwise Correction: Achieved by setting all feasible thresholds to $0$ (project on every step).
• Periodic Correction: A uniform allocation strategy that serves as the primary baseline.

Key Contributions

1. Timing as a Design Variable: The paper demonstrates that delayed correction alters the distribution of future updates, producing endpoint deviations that terminal projection alone cannot resolve. Constraint timing is identified as a first-class design variable in generative sampling.
2. Budget Allocation Formulation: Constraint enforcement is formalized as a budgeted allocation problem over time. This allows for the direct comparison of terminal, periodic, and adaptive strategies under matched computational budgets.
3. Defect as a Value Proxy: The authors show that one-step constraint defect serves as an effective online proxy for downstream trajectory distortion. By spending corrections on high-defect steps, the scheduler approximates the optimal allocation of limited projection resources.
4. Improved Cost–Accuracy Frontier: The proposed adaptive scheduling improves the trade-off between projection cost and trajectory accuracy. It recovers near-stepwise fidelity while using significantly fewer projection calls.

Experimental Results

The method was evaluated across controlled manifold rollouts (SO(3), SE(3), and Terrain domains) and a learned Projected Diffusion Model (PDM) sampler.

• Manifold Rollouts: In domains with concentrated defect (e.g., "Impulse" or "Ridge" variants where errors occur in specific bursts), adaptive scheduling significantly outperformed periodic correction. At a fixed budget of $B/T \approx 0.25$, adaptive scheduling reduced Normalized Excess Path Error (NEPE) by 41% compared to periodic correction in the SE(3)-Lever domain.
• Projected Diffusion Models: When applied to an existing PDM sampler, adaptive scheduling recovered 71.2% of the full stepwise benefit while using only 25% of the projection calls (a 75% reduction in cost). In contrast, periodic correction at the same budget recovered only 51.4% of the benefit.
• Generalization: The results indicate that constraint violations are highly heterogeneous over time. Selectively correcting high-defect steps consistently outperforms uniform strategies, particularly when the defect profile is volatile.

Significance and Claims

The paper claims that enforcing feasibility alone is insufficient to preserve the intended constrained sampling dynamics. The primary contribution is the demonstration that constraint timing is a fundamental degree of freedom in generative sampling.

The authors assert that:

• Projection is not merely a feasibility operation but an intervention in the dynamics. Delaying it changes the states from which future updates are computed.
• Adaptive scheduling allows practitioners to spend correction effort where it matters most (high-defect regions) rather than uniformly or only at the end.
• This approach is model-agnostic, requires no gradient computation through the projection, and introduces negligible overhead beyond defect evaluation.

The work suggests that future constrained generation should treat feasibility operations as trajectory-level interventions, recognizing that the path taken to reach a feasible sample is as critical as the sample itself.