RAISE Lab — Responsible AI for Science and Engineering
RAISE LabResponsible AI for Science & Engineering
Research area guide · University of Virginia

Constraint-aware
generative AI

How to make diffusion and flow-matching models obey physics, geometry, logic, and safety requirements—without giving up the diversity that makes generative models useful.

9 papers synthesized 4 application areas ~22 min read
The recurring idea
Denoise
Correct
Learned priorConstraint set

Let the model propose. Let optimization enforce. Repeat throughout generation.

From data to noise—and back again

A diffusion model learns to generate by first studying how data is destroyed. Given a clean sample \(x_0 \sim p_{\mathrm{data}}\), the forward process adds a small amount of Gaussian noise at each step:

\[q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(\sqrt{1-\beta_t}\,x_{t-1},\,\beta_t I\right).\]

Writing \(\alpha_t = 1-\beta_t\) and \(\bar\alpha_t=\prod_{s=1}^{t}\alpha_s\) lets us jump directly to any noise level:

\[x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I).\]

Early states retain most of the structure in \(x_0\); late states are almost indistinguishable from Gaussian noise. The generative problem is to learn the reverse direction: starting from \(x_T \sim \mathcal{N}(0,I)\), recover a sequence of progressively cleaner states until reaching a sample \(x_0\).

Forward processData → noiseA fixed variance schedule gradually erases structure.
Learned reverse processNoise → dataA neural score estimates which direction increases probability.

Score-based denoising

Rather than learning the normalized density \(p_t(x)\) itself, a score-based model learns its spatial gradient:

\[s_\theta(x_t,t) \approx \nabla_{x_t}\log p_t(x_t).\]

The score is a vector field. At any noisy point \(x_t\), it points toward nearby configurations that have higher probability under the perturbed data distribution. During training, denoising score matching provides the target analytically because the conditional corruption kernel is Gaussian:

\[\nabla_{x_t}\log q(x_t\mid x_0) = -\frac{x_t-\sqrt{\bar\alpha_t}x_0}{1-\bar\alpha_t} = -\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}.\]

The network is therefore trained across noise levels to minimize a weighted objective of the form

\[\mathcal{L}_{\mathrm{score}}(\theta)= \mathbb{E}_{t,x_0,\epsilon}\!\left[ \lambda(t)\left\|s_\theta(x_t,t)+\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}\right\|_2^2 \right].\]
What the score learns

The network does not memorize a single route from noise to data. It learns a direction field at every noise level: “from this corrupted point, move this way to become more typical of the data.”

Sampling with Langevin dynamics

Once the score is learned, we can sample by following that field while retaining controlled stochasticity. A discretized Langevin update at noise level \(t\) is

\[x_t^{k+1}=x_t^k+\gamma_t s_\theta(x_t^k,t)+\sqrt{2\gamma_t}\,\xi^k, \qquad \xi^k\sim\mathcal{N}(0,I).\]

The score term moves the sample toward regions of higher learned density. The Gaussian term preserves exploration, preventing every trajectory from collapsing onto the same local mode. Sampling is annealed from large to small noise levels: the process explores broadly when \(t\) is large, then becomes increasingly precise as \(t\to0\) and the variance schedule shrinks.

In the small-noise limit, the stochastic term becomes negligible and Langevin dynamics approaches gradient ascent on \(\log p_t(x)\). Equivalently, denoising approximately minimizes the negative log-likelihood \(-\log p_t(x)\). This optimization view is the bridge to constraint-aware generation.

Why a realistic sample is not necessarily a valid one

Diffusion models are excellent at learning what their training data looks like. Scientific and engineering systems demand something stricter: every output may also need to satisfy equations, geometric relationships, safety margins, or logical rules.

That distinction is easy to miss. A protein backbone can look natural while violating a bond-angle requirement. A robot path can be smooth while crossing an obstacle. A generated material can resemble the training distribution while having impossible interatomic distances.

The thesis Generation should be treated as a constrained decision process—not as “sample first, repair later.”

Our work inserts constraint enforcement inside the generative trajectory so that plausibility and feasibility are negotiated at every step.

A

Generate, then filter

Simple, but wastes samples and can collapse diversity when feasible outputs are rare.

B

Condition or retrain

Powerful when data and compute are abundant, but expensive and tied to a particular constraint family.

C

Constrain while sampling

Wrap a pretrained model with projections, proximal steps, solvers, or simulators that steer each generation.

Figure 1 Standard sampling can leave the feasible set even when the output remains visually plausible. Constraint-aware sampling repeatedly returns the trajectory toward the set of valid solutions.

Turn every denoising step into a small constrained optimization problem

A diffusion model gradually transforms noise \(x_T\) into a structured sample \(x_0\). The key observation behind our constraint-aware foundation is that each reverse step already behaves like a stochastic optimization step.

Recall that the learned score

\[s_\theta(x_t,t) \approx \nabla_{x_t}\log p_t(x_t)\]

points toward configurations with higher probability under the model’s distribution at noise level \(t\). A Langevin reverse update follows that direction while adding Gaussian noise:

\[x_t \leftarrow x_t + \gamma_t \nabla_{x_t} \log p(x_t) + \sqrt{2\gamma_t}\,\epsilon.\]

Because ascending \(\log p_t(x)\) is equivalent to descending the negative log-likelihood \(f_t(x)=-\log p_t(x)\), this is a noisy descent step on the model’s own generative objective. At each noise level, ordinary reverse diffusion can therefore be viewed as approximately solving

\[\min_x\; -\log p_t(x).\]

The noise is important early in sampling because it lets trajectories explore multiple high-probability regions. As \(t\to0\) and the variance schedule shrinks, the stochastic term becomes smaller and the update approaches ordinary gradient descent on \(-\log p_t(x)\). This is what makes the connection to constrained optimization natural: we do not replace the diffusion objective—we restrict where its solution is allowed to live.

If scientifically valid outputs form a feasible set \(C\), the reverse step becomes

\[\min_{x\in C}\; -\log p_t(x).\]

Projected optimization gives a direct way to implement this constrained step. First, let the score model take its usual noisy denoising update:

\[\widetilde{x}_t = x_t + \gamma_t s_\theta(x_t,t) + \sqrt{2\gamma_t}\,\epsilon.\]

Then map that proposal to its nearest feasible point:

\[x_t^{+}=\mathcal{P}_C(\widetilde{x}_t), \qquad \mathcal{P}_C(z)=\arg\min_{y\in C}\lVert y-z\rVert_2^2.\]

This produces an alternating process: the score step preserves alignment with the learned data distribution, while the projection step enforces the explicit scientific or engineering requirements. Crucially, projection is applied throughout the reverse chain—not once after generation. Each feasible iterate becomes the starting point for the next denoising step, so the model and the constraint operator progressively negotiate a sample that is both likely and valid.

In plain language

The model says, “move this way to look more like the data.” The projection replies, “this is the nearest point that also obeys the rules.” Repeating that exchange lets the pretrained model supply realism while constrained optimization supplies feasibility.

Figure 2 The reusable pattern: a learned generative update alternates with an explicit feasibility update throughout the reverse process.
Technical depth: what changes when constraints are difficult?

Convex constraints. Exact projections are often efficient, and the distance to the feasible set can be shown to contract. Under the assumptions in our analysis, an \(\varepsilon\)-feasible sample is reached in \(\mathcal{O}(\gamma_{\min}^{-1}\log(1/\varepsilon))\) steps.

Non-convex constraints. Augmented Lagrangian and ADMM-style updates replace a single exact projection with local primal–dual corrections.

Latent-space models. A constraint such as porosity or stress is evaluated after decoding \(z_t\), then differentiated back through the decoder. This makes the constraint visible to a model operating in a compressed representation.

Black-box simulators. When a finite-element or physics simulator is not differentiable, calibrated perturbations estimate a useful correction direction.

One principle, several kinds of constraints

The framework is best understood as a research program rather than a single algorithm. Different domains change the representation, the constraint operator, and the appropriate solver—but preserve the same division of labor between a learned prior and explicit domain knowledge.

100%protein constraint satisfaction
4.6×better stress–strain error
100robots planned jointly
0%toxicity-rule violations
Application 01

Protein design with hard structural constraints

The challenge is not merely to generate a plausible protein backbone. It is to preserve the exact three-dimensional relationships that support molecular recognition while allowing the surrounding structure to change.

Why PDZ domains matter

Many proteins work not as isolated machines, but as parts of larger molecular assemblies. PDZ domains—named after the first proteins in which the shared domain was recognized: PSD-95, Discs large, and ZO-1—are small, modular protein-interaction domains that help organize such assemblies. They commonly recognize a short motif at the end of a partner protein, binding it through a groove in which the peptide forms β-sheet-like hydrogen-bond contacts with the PDZ domain.

In cells, PDZ-containing proteins act as molecular organizers. They help localize receptors, ion channels, adhesion proteins, and signaling molecules at the places where they are needed. PSD-95, for example, uses PDZ-mediated interactions to assemble and regulate protein complexes at neuronal synapses. The broader importance is not that every PDZ interaction is individually strong—many are modest-affinity and somewhat promiscuous—but that a modular domain can recruit the right partners into a larger, spatially organized signaling complex.

This modularity also makes PDZ domains useful test cases for protein engineering. One established strategy is to attach a second small domain that contacts portions of the peptide not recognized by the PDZ groove, creating an affinity clamp with a larger interaction surface and potentially improved specificity (Huang et al., 2009). Designing such a construct computationally is difficult: the added scaffold must make the intended cross-chain hydrogen bonds while preserving realistic local bond lengths, angles, and secondary structure throughout the new backbone.

What the design task asks

Our benchmark begins with a resolved PDZ domain bound to its peptide motif. The task is not to redesign the natural PDZ interaction from scratch. Instead, the model must generate an additional domain, connected to the PDZ, that reaches the opposite face of the peptide and forms complementary β-sheet-like contacts. This requires satisfying local stereochemistry and global relationships between atoms on different chains at the same time.

The benchmark was built from experimentally resolved PDZ–peptide complexes in the Protein Data Bank. After structural curation and computational filtering, it contains 31 design targets. Six include short or geometrically restrictive peptide ligands—particularly proline-containing cases—that make the requested backbone arrangement poorly posed. Keeping these difficult cases in the benchmark provides a more realistic measure of where the method succeeds and where the design specification itself leaves little room for a viable structure.

What we found

Our constrained RFdiffusion method combines proximal feasibility updates with ADMM decomposition. Across 31,000 generations from each of three baselines—93,000 baseline samples in total—none simultaneously satisfied all of the specified inter-chain bond-distance and angle constraints. The constrained method satisfied those encoded constraints in every generated sample.

Constraint satisfaction alone, however, is not the same as a successful protein design. A generated backbone was counted as usable only when it also passed structural-realism checks. Under that stricter measure, 21% of generations were usable across the complete 31-target benchmark. For the subset with well-posed peptide geometries, the usable rate reached 83%. The three baselines produced no usable samples under the same combined criteria.

Figure 3 summarizes this comparison and shows representative constrained backbones: the method preserves the requested contacts across the peptide while still producing candidates that pass the benchmark’s structural-realism checks.

Encoded constraints100%vs. 0% across 93,000 baseline samples
Usable designs21%across the complete 31-target benchmark
Well-posed ligands83%usable designs on the favorable subset
What these results show Constraint-aware sampling turns demanding geometric specifications that defeated the evaluated baselines into a meaningful pool of computational protein candidates.

This is an important step toward making functional geometry a controllable part of generative protein design. The resulting backbones are promising starting points rather than finished proteins: expression, stability, binding affinity, specificity, and biological function still require experimental validation, and the 21% overall usability rate leaves clear room for improvement.

Figure 3 In the computational benchmark, constraint-aware RFdiffusion preserves the requested inter-chain geometry and yields structurally usable candidates where the evaluated baselines fail the combined global constraints.

On a separate vacancy-constrained pocket-design benchmark, the same framework produced 97.8% usable samples, compared with 24.2% for the nearest baseline. That result supports the method’s applicability beyond PDZ scaffolding, but it measures a different geometric task and should not be conflated with the 21% overall PDZ result.

Application 02

Materials discovery with properties in the loop

In materials design, a sample can look plausible while having the wrong porosity, mechanical response, or physical geometry. The constraint operator turns those measurable properties—and, when needed, an external simulator—into feedback during generation.

Why material structure matters

A material’s behavior depends not only on what it is made of, but also on how its internal structure is organized. Pore size and connectivity influence transport through a porous medium; in cellular metamaterials, the arrangement and thickness of the struts determine how the structure bends, buckles, and carries load. Small geometric changes can therefore produce large changes in the resulting physical response.

Generative models could help explore these large design spaces and augment datasets that are costly to obtain through high-resolution imaging, fabrication, or simulation. But visual realism is not enough. A generated microstructure may resemble the training data while missing the requested porosity, and a convincing cellular pattern may fail to reproduce a target stress–strain curve—or may not define a physically valid structure at all.

What the design task asks

We study two complementary tasks. In porous-microstructure generation, the model must create an image at a prescribed porosity, including values that are sparsely represented in the data. Because porosity can be computed directly from the decoded image, the method can project each intermediate sample toward the required value. In mechanical-metamaterial inverse design, the input is instead a target nonlinear stress–strain response under compression. Here, the relevant property is evaluated by an external finite-element simulator, so the method uses simulator queries to estimate a useful correction and then applies it in the latent space of the diffusion model.

The common pattern is generate, evaluate, and correct without retraining the underlying diffusion model. Figure 4a illustrates this broader constraint-in-the-loop architecture in an atomistic setting, where a surrogate checks physical interactions during denoising. Figure 4b shows the mechanical-metamaterial experiment: as the latent sample is corrected, its simulated stress–strain curve moves toward the target and the geometry becomes physically valid.

What we found

For porous microstructures, conditional diffusion produced visually strong samples but missed the requested porosity by more than 10% in 68.4% of cases. The constrained method reduced that violation rate to 0%. Its Fréchet Inception Distance was 13.5, compared with 10.8 for conditional diffusion, indicating a modest image-quality tradeoff in exchange for much more reliable property control.

For mechanical metamaterials, the constrained method reached a stress–strain mean-squared error of 1.4 ± 0.6, compared with 6.4 ± 4.6 for a specialized inverse-design baseline and 7.1 ± 4.5 for conditional diffusion. Only 5% of its outputs were judged physically invalid, versus 20% and 55% for those two baselines, respectively. These results show that simulator feedback can improve both response matching and geometric validity on this computational benchmark.

Figure 4a The broader constraint-in-the-loop pattern for atomistic generation: a surrogate evaluates interactions and returns a correction during denoising.
Figure 4b Simulator-guided samples move toward the target stress–strain response; the final benchmark reports 5% physically invalid geometries.
Porosity error >10%0%vs. 68.4% for conditional diffusion
Stress–strain MSE1.4 ± 0.64.6× lower than a specialized baseline
Invalid geometries5%vs. 20% specialized and 55% conditional
What these results show Constraint-in-the-loop generation makes target material properties substantially more reliable while retaining competitive sample quality.

The results demonstrate a practical way to use image measurements and finite-element feedback during generation, opening a larger and more controllable computational design space. The next step is experimental validation: fabrication constraints, defects, robustness to loading conditions, measured stress–strain behavior, and other application-specific properties remain outside these benchmarks. Simulator-guided sampling also introduces a meaningful computational cost.

Application 03

Safe motion planning at generative speed

A multi-robot plan is feasible only if every robot reaches its goal without striking an obstacle or another robot—and does so along a trajectory that its motion limits can support.

Why coordinated motion matters

Multi-robot motion planning appears wherever autonomous systems share space: warehouse fleets moving inventory, coordinated drones, autonomous vehicles, and robots operating in search-and-rescue settings. Planning each path independently is not enough. A route that is safe for one robot can block another, create a collision several seconds later, or produce a deadlock in a narrow corridor. The difficulty grows quickly because each additional robot adds a trajectory whose position must remain compatible with every other trajectory over time.

Classical grid-based multi-agent path finding scales well, but its coarse spatial and temporal discretization can produce paths that are difficult to execute smoothly. Continuous trajectory optimization offers finer motion, but the joint problem becomes high-dimensional and strongly non-convex as robots and obstacles are added. A generative model provides a learned prior over plausible paths; the central question is whether that prior can be converted into a coordinated, feasible plan.

What the planning task asks

Each robot is assigned a start and goal position in a shared two-dimensional workspace. The planner must produce all trajectories simultaneously while satisfying four coupled requirements: avoid static obstacles, maintain separation between robots at every time step, respect velocity and acceleration limits, and preserve the assigned start and goal states. Practical benchmark layouts include warehouse-like shelves, connected rooms, dense obstacle fields, and corridors where robots must negotiate passage rather than simply follow locally shortest routes.

Simultaneous MRMP Diffusion (SMD) addresses the joint continuous problem directly. During reverse diffusion, it repeatedly corrects the full set of trajectories using an augmented-Lagrangian projection, so obstacle, kinematic, and inter-robot constraints influence generation before the paths are finalized. Discrete-Guided Diffusion (DGD) tackles the scaling bottleneck differently: a discrete multi-agent path-finding solution supplies a coordination schedule, the continuous workspace is decomposed into tractable regions, and constrained diffusion refines the motion within those regions. A lightweight repair step handles residual violations.

The animations in Figure 5 make the constraint visible over time. They show that feasibility depends on the entire joint trajectory—particularly in corridors, shelves, and rooms—not merely on collision-free final positions.

What we found

On the SMD benchmark’s most difficult random setting—nine robots moving among 20 obstacles in a dense map—SMD solved 96% of instances, compared with 27.2% for the strongest prior multi-robot diffusion baseline. In the corridor task, SMD solved every instance while the evaluated baselines found no feasible trajectory. These results measure whether the generated plan is collision-free and satisfies the encoded motion constraints within the benchmark.

DGD extends the approach to larger problems. On the standard benchmark it achieved at least 92% success in all but the two hardest 18-robot settings, where it still led at 76%. On dense six-robot maps, DGD matched SMD’s 100% success while using 4% of SMD’s runtime. The scalability study generated feasible plans for 100 robots among 104 obstacles, compared with the 40-robot, obstacle-free scale reported for prior diffusion-based planning.

Basic
Dense
Corridor
Shelf
Room

Figure 5 Simulated multi-robot trajectories in five benchmark environments. The animation exposes coordination events and possible conflicts that a final frame would hide.

Dense-map success96%vs. 27.2% for the strongest prior baseline
Maximum scale100 robotswith 104 obstacles
DGD runtime4%of SMD's time on dense six-robot maps
What these results show Constraint-aware diffusion can coordinate many robots in cluttered environments where prior generative planners frequently fail.

The gains in feasibility and scale make diffusion-based planning a credible option for difficult coordination problems, especially when multiple valid trajectories may exist. Moving from simulation to physical deployment will require models of perception error, moving obstacles, tracking and control error, communication limits, and hardware dynamics. SMD's inference cost and consistency across DGD's decomposed regions also remain useful targets for improvement.

Application 04

Constraints in discrete spaces: molecules, text, and code

In discrete diffusion, “move to the nearest feasible point” cannot mean nudging coordinates. It means changing probability mass over tokens without destroying the model's learned preferences.

Why discrete constraints matter

Text, molecular strings, and source code are built from discrete symbols. A single misplaced token can reverse the meaning of a sentence, break a SMILES representation, or introduce a software vulnerability. Many important requirements are also global: a word may need to appear at a particular position, a molecular sequence must decode to a valid and previously unseen structure, and a program must behave correctly across an entire test suite. These constraints cannot always be enforced by choosing the most likely next token in isolation.

Masked diffusion offers a useful intervention point because it repeatedly predicts the whole sequence rather than committing permanently from left to right. Intermediate tokens can be reconsidered as the global structure becomes clearer. This makes discrete diffusion particularly well suited to constraints that depend on relationships between distant parts of a sequence.

What the generation task asks

Constrained Discrete Diffusion (CDD) operates on the probability distribution over tokens. At each reverse step, it finds a nearby distribution whose most likely decoded sequence satisfies the constraint, using KL divergence to remain close to the base model’s proposal. The same principle supports several types of evaluator: a toxicity classifier for controlled text, exact counting or positional rules for instruction following, and property or novelty checks for molecular SMILES.

The molecule sequence in Figure 6 provides the intuition: constraint checks accompany generation instead of being applied only after a final candidate appears. In the actual sequence benchmarks, the projection acts on token probabilities and can encode properties such as synthetic accessibility, novelty, or the absence of selected BRENK structural alerts—not only local valency.

Code requires a richer version of this loop. Constrained Diffusion for Code (CDC) decodes an intermediate program, evaluates functional or security requirements, identifies the token regions most responsible for a violation, and locally revises or remasks those regions. Differentiable surrogate feedback guides functional correctness, while static program analysis supplies concrete vulnerability information. This uses diffusion’s editable global state to repair a developing program before the final sample is committed.

What we found

Across the CDD text experiments, the constrained sampler produced 0% violations at all evaluated toxicity thresholds and on both counting and positional lexical constraints. The corresponding unconstrained MDLM violated 54.5% of counting requirements and 97.5% of lexical requirements. This exact compliance came with task-dependent quality and computational tradeoffs; for example, constrained lexical generation improved the paper’s coherence score but had somewhat higher perplexity than unconstrained MDLM.

In the molecular-generation extension, CDD generated an average of 392 novel candidates that passed the encoded BRENK screens, compared with 108 for MDLM and 5 for an autoregressive baseline, while recording no BRENK-rule violations. These are useful filters for exploring chemical space, but they are not a complete assessment of toxicity, synthesizability, stability, or biological activity.

For code, CDC increased functional correctness with Dream-Coder 7B from 34.1% to 65.2% on HumanEval-X C++ and from 27.7% to 59.2% on MBPP-C++. It also raised syntax correctness from 67.1% to 79.2% and from 41.1% to 72.0% on those benchmarks. On CWEval, joint functional-and-security success increased from 12.04% to 34.26%. Figure 7 summarizes these improvements across functionality, syntax, and security tests.

Figure 6 Conceptual view of checking molecular validity throughout generation. CDD implements this principle over discrete sequence probabilities and supports both local and sequence-level rules.
Figure 7 CDC improves benchmark pass rates for program functionality, syntax, and security across two seven-billion-parameter diffusion code models.
Rule violations0%for toxicity, counting, and lexical tests
BRENK-screened candidates392novel molecules vs. 108 for MDLM
Code functionality65.2%vs. 34.1% for the Dream-Coder base model
What these results show Intervening during denoising can materially improve rule adherence across language, molecular design, and code without retraining the base generator.

This makes discrete diffusion a promising platform for applications where outputs must be both plausible and verifiably structured. The guarantees remain tied to the supplied evaluator: toxicity classifiers, molecular screens, tests, and static analyzers each cover only part of the real requirement. Constraint optimization also adds computation, so human review and domain-specific validation remain important for consequential uses.

From correcting a pretrained model to training for constraints

Training-free guidance is flexible, but it creates a mismatch: the model learns unconstrained trajectories during training and encounters projected trajectories only at inference. Two newer directions attack that mismatch from opposite sides.

Continuous models

Constraint-Aware Flow Matching

CAFM includes projection in the training objective, teaching the velocity field to anticipate the constrained sampler. The aim is a better quality–feasibility trade-off with strict per-sample guarantees.

Christopher et al., 2026 →
Discrete models

Search-Augmented Diffusion

SearchDiff combines masked diffusion with tree search, spending extra computation to explore alternative token sequences when direct projection is not enough.

Ta et al., 2026 →
Figure 8 CAFM aligns the learned velocity field with the constrained sampling path instead of correcting an unconstrained field only at inference.

What guarantees do—and do not—buy us

Constraint satisfaction is only as meaningful as the constraint itself. A mathematically feasible output can still be undesirable if the evaluator is incomplete, misspecified, or only a surrogate for the real scientific objective.

01

Evaluator fidelity. How should uncertainty in learned surrogates or approximate simulators propagate into a feasibility claim?

02

Competing constraints. What happens when requirements are mutually inconsistent, or feasibility leaves too little room for diversity?

03

Non-convex scale. Can local projection and decomposition methods retain useful guarantees on the largest scientific design spaces?

04

Training–sampling alignment. When is a training-free wrapper enough, and when is end-to-end constraint-aware training worth the added cost?

The broader principle A generative model supplies a prior over plausible designs. Domain constraints turn that prior into a decision-making tool.

Follow the research program