Deep Generative Models with Hard Linear Equality Constraints

TL;DR

This work tackles the challenge of enforcing hard linear equality constraints within deep generative models by learning and sampling from a constrained Gaussian distribution $p_{\boldsymbol{\theta}}(\boldsymbol{z} \mid \boldsymbol{A}\boldsymbol{z}=\boldsymbol{k})$ instead of post-hoc sample correction. It introduces gradient estimators that exploit constraint information, notably a Marginal Expectation proxy, and derives closed-form forms for Gaussian (and Poisson) cases to enable end-to-end training with exact constraint satisfaction. The approach is validated across VAEs, diffusion models, and graph neural networks on five image datasets and three scientific applications, showing consistent constraint satisfaction and superior generative performance compared to existing baselines, including methods that enforce constraints only at inference time. The results highlight the practical impact of principled constraint integration for constraint-sensitive domains such as chemistry, chemical engineering, and finance, enabling more accurate and physically plausible generative modeling. Overall, the paper provides a versatile, theoretically grounded framework for incorporating hard linear equality constraints into DGMs with broad methodological and applied implications.

Abstract

While deep generative models~(DGMs) have demonstrated remarkable success in capturing complex data distributions, they consistently fail to learn constraints that encode domain knowledge and thus require constraint integration. Existing solutions to this challenge have primarily relied on heuristic methods and often ignore the underlying data distribution, harming the generative performance. In this work, we propose a probabilistically sound approach for enforcing the hard constraints into DGMs to generate constraint-compliant and realistic data. This is achieved by our proposed gradient estimators that allow the constrained distribution, the data distribution conditioned on constraints, to be differentiably learned. We carry out extensive experiments with various DGM model architectures over five image datasets and three scientific applications in which domain knowledge is governed by linear equality constraints. We validate that the standard DGMs almost surely generate data violating the constraints. Among all the constraint integration strategies, ours not only guarantees the satisfaction of constraints in generation but also archives superior generative performance than the other methods across every benchmark.

Figures (9)

• Figure 1: Comparison of different methods for generating samples that satisfy linear equality constraints. The left panel shows the original unconstrained distribution in a $2$-dimensional plane, with the purple line representing the constraint $x_1 + x_2 = 0$. Our proposed method generates the most realistic sample as indicated by the right figure, outperforming existing methods that optimize for L1 distance (CL) and L2 distance (Euc).
• Figure 2: The constrained model considered in this work. It involves an encoder $h_{\boldsymbol{v}\xspace}\xspace$ that outputs $\bm{\theta}$ to parameterize a latent distribution constrained by the linear equality constraint $\bm{A} \bm{z} = \bm{k}$. We first study when the objective admits a closed-form expression such that standard training is amenable. We further propose and study various gradient estimators for the general case by combining exact sampling in the forward pass and gradient approximations in the backward pass.
• Figure 3: Comparisons of gradient estimators for point-wise loss $\ell$ being L1 loss (upper plot) and L2 loss (lower plot) applied to Gaussian variable are conducted. To compare the directions of the estimated and ground-truth gradients, we utilize the cosine distance. The bias, variance, and error of the gradient estimators are measured using a sample size of $10,000$.
• Figure 4: Comparison of gradient estimators for VAE with constrained latent space. Negative log-likelihood (NLL), negative ELBO (NELBO), and reconstruction loss (RL) are averaged over $5$ trials.
• Figure 5: The first block displays the original MNIST images and the ones modified by the brightness constraint as inputs. For the following blocks, each displays the reconstructed images by different VAE architectures. Within each block, the first row is generated by the unconstrained VAE, the second by VAE constrained by the baseline Constrained Layer and the last one by VAE constrained by our method.

Theorems & Definitions (18)

• Proposition 4.1: Gaussian Conditional Marginal and Expectations
• Proposition 5.1: Gaussian Closed-form Expected Loss
• Proposition 7.1: Poisson Closed-form Expected Loss
• Proposition 7.2: Poisson Constrained Distribution
• Proposition 7.3: Poisson Conditional Marginal and Expectations
• Proposition 8.1: Gaussian Constrained Distribution
• proof
• Proposition 8.2: Gaussian Closed-form Expected Loss
• proof
• Proposition 8.3: Gaussian Conditional Marginal and Expectations