Learning with Embedded Linear Equality Constraints via Variational Bayesian Inference

1. Problem Statement

In scientific and engineering domains, predictive models must often adhere to known physical laws, typically expressed as linear equality constraints (e.g., mass and energy balances). However, standard Machine Learning approaches face two critical limitations:

1. Lack of Uncertainty Quantification: Deterministic neural networks provide point predictions without principled confidence intervals.
2. Physical Inconsistency: Standard Bayesian Neural Networks (BNNs) quantify uncertainty but do not explicitly enforce physical constraints, leading to predictions that may violate known physics (e.g., violating Kirchhoff's laws or energy conservation).

Existing methods to enforce constraints (e.g., Physics-Informed Neural Networks or hard projection layers) often treat constraints deterministically, failing to propagate uncertainty or account for measurement noise and partial observability. There is a need for a framework that embeds constraints probabilistically while maintaining calibrated uncertainty estimates.

2. Methodology

The authors propose a Bayesian Constrained Probabilistic Neural Network (BCPNN) framework that integrates linear equality constraints directly into the variational inference process.

Core Concept: Gaussian Conditioning

The method assumes the predictive distribution of the neural network is Gaussian. When constraints are linear, enforcing them reduces to closed-form Gaussian conditioning.

• Constraint Formulation: The authors define a constraint residual $z = Ax + By - b + \epsilon$, where $\epsilon$ represents a tolerance term modeled as Gaussian noise with variance $r$.
• Conditioning: By conditioning the predictive distribution on $z=0$, the mean ($\mu_P$) and variance ($\sigma^2_P$) of the network output are updated to a constrained posterior ($\mu_C, \sigma^2_C$) using analytical formulas derived from the properties of multivariate Gaussians.
• Differentiability: This conditioning step is fully differentiable, allowing it to be inserted as a layer within the network architecture.

Joint Variational Inference

Unlike standard BNNs that infer only model parameters ($\theta$), this framework performs joint inference over:

1. Network Parameters ($\theta$): The weights and biases of the neural network.
2. Constraint Tolerance ($r$): A random variable representing the strictness of the constraint.

The authors treat the tolerance $r$ as a learnable quantity (reparameterized as $r = \exp(\rho)$ to ensure positivity). They optimize a modified Evidence Lower Bound (ELBO) that balances:

• Data fit (likelihood of observations).
• Regularization of network parameters (KL divergence from priors).
• Regularization of the constraint tolerance (KL divergence from priors).

This allows the model to learn how strictly to enforce constraints based on the data, rather than assuming constraints hold perfectly or are violated arbitrarily.

Uncertainty Decomposition

The paper provides a theoretical decomposition of predictive variance under constraint conditioning. The total variance is shown to consist of:

• Aleatoric Uncertainty: Reduced by the conditioning step (constraint reduces marginal variance).
• Epistemic Uncertainty: Reduced due to tighter posterior concentration.
• Tolerance Uncertainty: Variance arising from the learned uncertainty in the constraint tolerance itself.
• Interaction Terms: Covariance between parameters and the constraint update.

3. Key Contributions

1. Probabilistic Constraint Embedding: A novel framework that embeds linear equality constraints into BNNs via Gaussian conditioning, ensuring predictions are physically consistent in expectation while preserving uncertainty quantification.
2. Learnable Constraint Tolerance: The ability to infer the appropriate strictness of constraints from data, distinguishing between "hard" physical laws and "soft" relationships affected by measurement noise.
3. Theoretical Uncertainty Analysis: A derivation showing that constraint conditioning systematically reduces predictive variance (both aleatoric and epistemic) while maintaining calibrated coverage.
4. Tractable Optimization: The use of closed-form updates allows for efficient training via standard variational inference techniques (reparameterization trick) without requiring complex optimization layers (like QP solvers) during inference.

4. Experimental Results

The framework was evaluated on a Single Particle Model (SPM) of a lithium-ion battery, a complex system governed by spherical-diffusion equations and thermal dynamics.

• Setup: The model predicted 8 outputs (electrochemical and thermal) given 3 inputs (current, temperature, etc.), subject to two linear constraints: Kirchhoff's Voltage Law and Energy Balance.
• Performance:
  • Accuracy: The BCPNN achieved point prediction accuracy and coverage ratios (0.99) identical to a standard unconstrained BNN.
  • Uncertainty Reduction: The BCPNN produced tighter credible intervals (reduced variance) compared to the standard BNN, confirming the theoretical variance reduction.
  • Constraint Adherence: The most significant improvement was in constraint violation. The BCPNN reduced median constraint violations by two orders of magnitude for voltage and four orders of magnitude for energy balance compared to the unconstrained BNN.
  • Adaptive Tolerance: The learned posterior for the constraint tolerance successfully distinguished between the two constraints. The voltage law was learned to be extremely strict (low tolerance), while the energy balance retained higher uncertainty (higher tolerance), reflecting the differing noise levels in the simulated data.

5. Significance and Future Work

This work bridges the gap between structured probabilistic modeling and modern deep learning. It provides a principled way to incorporate domain knowledge (physical laws) into neural networks without sacrificing the ability to quantify uncertainty.

• Impact: It enables the deployment of ML models in safety-critical engineering applications where physical consistency is non-negotiable, yet data is noisy or incomplete.
• Limitations & Future Directions: The current formulation is limited to linear equality constraints. Future work aims to extend this to inequality constraints and nonlinear relationships. Additionally, the learned tolerance posteriors could be integrated into active learning or experimental design pipelines to guide data collection strategies based on both predictive and constraint uncertainty.