Embedding Linear Equality Constraints in Probabilistic Neural Networks for Dynamic Modelling

Problem Statement
Machine learning models are increasingly deployed for surrogate modelling in chemical process systems, offering the advantage of building models directly from data without requiring extensive specialist knowledge. However, these data-driven approaches often lack integrated physical knowledge, leading to predictions that may violate known physical laws, particularly when training data is limited or noisy. Furthermore, standard deterministic models fail to provide principled uncertainty quantification, which is critical for applications in optimization and control where data is often partially observable or stochastic. While Physics-Informed Neural Networks (PINNs) attempt to enforce physical consistency via soft constraints in the loss function, they do not guarantee constraint satisfaction and can introduce hyperparameter sensitivity and convergence issues. Conversely, hard-constrained methods often rely on output projections or reduced subspaces that may compromise model expressivity or increase computational burdens, especially for non-linear constraints. Additionally, existing constrained learning methods are predominantly deterministic, lacking the ability to model aleatoric (inherent) uncertainty.

Methodology
The authors propose a Probabilistic Neural Network (PNN) framework that enforces linear equality constraints within a specified tolerance while explicitly modeling output aleatoric uncertainty. The core of the methodology involves:

1. Probabilistic Formulation: Instead of predicting point estimates, the neural network $f_\theta$ predicts the parameters of a multivariate Gaussian distribution for the output $y$. Specifically, it outputs a mean vector $\mu_P$ and a Cholesky factor $L_P$ of the covariance matrix $\Sigma_P = L_P L_P^\top$. The model is trained to minimize the negative log-likelihood (NLL) of this distribution.
2. Constraint Definition: Linear equality constraints are defined as $Ax + By = b$. To allow for tractable conditioning, the authors introduce a small numerical diffusion term $\epsilon \sim \mathcal{N}(0, \epsilon I)$, treating the constraint residual $z = Ax + By - b$ as a Gaussian variable rather than a hard delta function.
3. Closed-Form Conditioning: Leveraging the conjugacy of linear Gaussian systems, the framework derives a closed-form update to condition the unconstrained output distribution $p(y|x)$ on the constraint $z=0$. This results in a new conditional distribution $p(y|z=0, x) = \mathcal{N}(\mu_Q, \Sigma_Q)$, where the mean and covariance are updated analytically. This projection ensures the output satisfies the constraints within the tolerance $\epsilon$ and reduces uncertainty in the direction of the constraint.
4. Implementation: The conditioning is implemented via non-trainable, differentiable layers embedded within the computational graph. The model is trained by minimizing the NLL of the conditioned distribution. This approach avoids iterative optimization during inference and maintains differentiability for backpropagation.

Key Contributions

• Guaranteed Constraint Satisfaction: Unlike soft-constrained PINNs, this framework guarantees that linear equality constraints are satisfied within a user-defined tolerance at both training and inference points.
• Uncertainty Quantification: The method captures aleatoric uncertainty, providing a predictive distribution rather than a deterministic point estimate, which is essential for risk-aware decision-making.
• Computational Efficiency: By utilizing closed-form conditioning, the method avoids the computational overhead associated with iterative projection methods or convex optimization layers often used for constrained learning.
• Improved Performance on Limited Data: The framework demonstrates that embedding physical constraints directly into the probabilistic output significantly improves predictive accuracy and uncertainty calibration when data is scarce.

Results
The framework was evaluated on two batch reactor case studies (one with irreversible reactions and one with reversible reactions) across four data regimes (minimal, small, medium, and large). The CPNN was benchmarked against a standard unconstrained ANN and a soft-constrained PINN.

• Data Regime Dependency: In minimal and small data regimes, the CPNN significantly outperformed both baselines in terms of Mean Squared Error (MSE), Continuous Ranked Probability Score (CRPS), and constraint adherence. The embedding of constraints reduced overfitting and guided the model toward physically feasible solutions. As data volume increased to the medium and large regimes, the performance gap narrowed, with unconstrained models eventually achieving similar accuracy as they learned the invariants from the data.
• Training Efficiency: The CPNN consistently demonstrated significantly faster training times (often reducing time by half or more) compared to the ANN and PINN across all data regimes. The authors attribute this to the conditioning step altering the loss landscape, focusing the search on feasible solutions and enabling faster convergence.
• Uncertainty Calibration: While the 95% coverage ratios were consistently above 0.95, the authors noted that the models tended to overestimate variance, leading to poorly calibrated confidence intervals. This suggests a tendency to overfit to noise rather than the underlying dynamics, a known challenge in probabilistic deep learning.
• Trajectory Prediction: Visual analysis of trajectory predictions showed that the CPNN provided tighter confidence intervals and predictions closer to the true system values in low-data regimes compared to baselines, which often aligned with noisy data points.

Significance and Claims
The paper claims that the proposed framework offers a principled approach to integrating physical knowledge into probabilistic machine learning without sacrificing computational tractability. The primary significance lies in its ability to deliver superior predictive accuracy and uncertainty quantification when data is limited, a common scenario in chemical process engineering. The authors note that while the performance advantage diminishes with abundant data, the conditioning step continues to offer benefits in terms of faster training and inference speeds. They suggest this efficiency makes the approach promising for real-time control and optimization on larger systems. The authors also modestly acknowledge that the current work is limited to linear equality constraints due to analytical tractability and that future work could explore extensions to non-linear or inequality constraints, as well as integration with sequential models like RNNs or Transformers.