ENFORCING PHYSICAL CONSTRAINTS IN NEURAL NETWORKS THROUGH DIFFERENTIABLE PDE LAYER
Year: 2019

Abstract:
Recent studies at the intersection of physics and deep learning have illustrated successes in the application of deep neural networks to partially or fully replace costly physics simulations. Enforcing physical constraints to solutions generated by neural networks remains a challenge, yet it is essential to the accuracy and trustworthiness of such model predictions. Many systems in the physical sciences are governed by Partial Differential Equations (PDEs). Enforcing these as hard constraints, we show, are inefficient in conventional frameworks due to the high dimensionality of the generated fields. To this end, we propose the use of a novel differentiable spectral projection layer for neural networks that efficiently enforces spatial PDE constraints using spectral methods, yet is fully differentiable, allowing for its use as a layer in neural networks that supports end-to-end training. We show that its computational cost is cheaper than a regular convolution layer. We apply it to an important class of physical systems – incompressible turbulent flows, where the divergence-free PDE constraint is required. We train a 3D Conditional Generative Adversarial Network (CGAN) for turbulent flow superresolution efficiently, whilst guaranteeing the spatial PDE constraint of zero divergence. Furthermore, our empirical results show that the model produces realistic flow fields with more accurate flow statistics when trained with hard constraints imposed via the proposed novel differentiable spectral projection layer, as compared to soft constrained and unconstrained counterparts.

heading: 1 INTRODUCTION
text: Convolutional Neural Network (CNN) based deep learning architectures have achieved huge success in many tasks across computer vision, but their use in the physical sciences have only recently been explored. Many parallels exist between physical science problems and those in computer vision. For instance, grid-based simulations generate a physical scalar or vector field which can been compared to multidimensional arrays in computer vision. However, unlike computer vision problems, physical fields are often constrained by PDEs that arise from the governing equations of the physical system. For example, the Poisson equation of the form∇2φ = f is often encountered in heat diffusion problems, whereas the divergence-free (also known as solenoidal) conditions in the form of ∇ · φ = 0 is fundamental to magnetic fields, as well as incompressible fluid velocity fields to ensure conservation of mass. For meaningful application of deep learning to a range of important physical problems it is essential to enforce such spatial PDE constraints to guarantee physical consistency and reliability of the model output for scientific applications. Yet, general means of enforcing these constraints do not exist and the existing methods do not scale well with high dimensional, high resolution outputs.
In this paper, we address this issue by proposing a novel differentiable PDE layer (PDEL) that efficiently enforces spatial PDE constraints for neural networks, at costs on par with a single CNN layer. We use spectral methods, which leverages the highly efficient Fast Fourier Transform (FFT) algorithm for enforcing such constraints. Using this formulation, we are able to exploit the structures of the spectral matrices corresponding to these differential operators that renders the entire layer O(n log n) for processing a 3 dimensional field of size n. The method is general for enforcing arbitrary linear combinations of differential operators on these fields, which encompasses physical constraints from a broad range of important scientific and engineering systems. We apply this hard constraining layer to the problem of turbulence superresolution, where we show that training with the
hard constraining layer in-the-loop not only guarantees that the imposed constraint is strictly satisfied, but also generates solutions that are more accurate measured via a variety of fluid flow metrics.
In summary, the main contributions of this paper are:
• We propose the highly efficient differentiable spatial linear PDE layer (PDEL), which strictly enforces linear spatial PDEs constraints.
• We apply the PDE layer towards the superresolution task for turbulent flows, showing that training with hard constraints in-the-loop results in solutions that not only strictly satisfy the imposed constraint but also produce flow fields with more accurate fluid flow statistics.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 CONSTRAINTS IN NEURAL NETWORKS
text: Many studies in machine learning have considered imposing some form of constraints for their respective applications. ? proposed an approach for constraining the prediction of a discriminative predictor and showed that a Gaussian Process can be forced to satisfy linear and quadratic constraints. ? proposed training a kernalized latent variable model that imposes equality and inequality constraints. ? proposed a constrained CNN, which phrases the training objective as a biconvex optimization for linear models, which is then relaxed to nonlinear deep networks for any set of linear constraints on the output space. ? proposed an alternative approach by randomly subsampling a set of constraints at each optimization step and projecting the gradients onto the feasible solution space. OptNet (?) solves a generic quadratic programming problem differentiably within the neural network, but its cubed complexity does not handle high dimensional output. ? proposed parameterizing the feasible solution space for imposing inequality constraints. However, methods to impose physical constraints into machine learning and deep learning models remains largely unexplored.

heading: 2.2 APPLICATIONS OF MACHINE LEARNING FOR PDE AND TURBULENCE
text: More recently there have been some work in the literature to apply machine learning to PDE and physical problems. ??? developed physics informed deep learning frameworks for assimilating observational data. In the context of applying machine learning methodologies to turbulence problems, earlier works using data-driven modeling approaches have used optimization and Bayesian inference approaches to calibrate existing turbulence models (?????). With the advancement of efficient and accurate modeling tools in machine learning, recent studies have looked at data driven approaches for turbulence modeling (?), including the direct use of random forests for modeling Reynolds-Averaged Navier Stokes (RANS) model errors (?), the use of multilayer perceptrons for modeling Reynolds stress anisotropy tensor from simulation data (?), and using random forests to predict mean velocities of turbulent flow fields. ? use a random forest to compute particle location and ? use a CNN to approximate part of a numerical PDE scheme. A recent application of deep learning for generating realistic fluid flow fields is TempoGAN (?) that uses a specialized discriminator network for temporal coherence. None of these methods, however, enforce constraints that are necessary for physically consistent fluid flow fields. ? addressed this with a customized loss function for the divergence-free constraint of fluid flow, but since it is a loss-based soft constraint, the conditions are ultimately not exactly satisfied.

heading: 3 PROBLEM STATEMENT
text: The main focus of this paper is to introduce a novel and efficient method for imposing spatial linear PDE constraints to the outputs of Convolutional Neural Networks (CNNs). This is discussed in the context of underdetermined systems, since solutions do not exist for overdetermined systems, while solutions for determined systems do not fit in the context of constraining the outputs. More specifically, given the output of the network to be a discretization of a 3D vector field f : R3 7→ R3 and a linear spatial PDE operator A(( ∂∂xj ) 0, ( ∂∂xj ) 1, · · · ) that maps vector fields to scalar fields Af : R3 7→ R, we seek a means of efficiently imposing the spatial linear PDE constraints within
CNNs, i.e.,
Af = b (1)
Note that this form encompasses a wide range of physically relevant constraints. In particular, all spatial PDE constraints composed of divergence, curl, Laplacian and other higher order partial differential terms in linear combination may be expressed in this form. Depending on the domain of application, this includes mass conservation for incompressible fluid flows, the heat equation, the wave equation, Laplace’s equation, Helmholtz equation, Klein-Gordon equation, and Poisson’s equation. For the important constraint of mass conservation in incompressible flows, we investigate the divergence-free (solenoidal) constraint of:
∇ · f = ∑ j ∂ ∂xj fj = 0 (2)

heading: 4 METHODS
text: Before presenting our proposed method for enforcing the solenoidal condition on CNN outputs, we present an overview of two commonly utilized strategies for enforcing general linear constraints, which we will compare and benchmark against in the experiments section (Sec. 5).
We first discuss enforcing linear constraints on the outputs of the neural network, where we have a neural network that learns the function mapping f : Rt 7→ Rm, where the function f(x;θ) is parameterized by learnable parameters θ, and is subject to the linear constraint Af(x;θ) = b, where A ∈ Rn×m, b ∈ Rn. For this to be an underconstrained system, we have n < m.

heading: 4.1 GENERALIZED LINEAR CONSTRAINTS
text: Two forms of constraints are possible for explicitly enforcing a certain set of constraints for neural network outputs: soft constraints and hard constraints.
Soft constraints are easy to implement, by adding a differentiable residual loss for penalizing the network during training time for violating the explicit constraints. For simplicity, let y := f(x;θ). In the conventional unconstrained case, assume the neural network is trained under the differentiable loss function L(f(x;θ)), in the constrained case, the loss function can be augmented by an additional residual loss term defined by:
Lc(θ) = L(θ) + α((Ay − b)T (Ay − b)) (3) where α is a hyper-parameter weighing the two loss functions that can be difficult to determine and vary between applications. Although easy to implement, soft constraints provide no guarantees on the solutions satisfying the imposed constraint.
Hard linear constraints can be enforced by posing the problem as a constrained optimization problem for seeking the closest point in the solution space subject to the constraints, which can be solved by satisfying the Karush-Kuhn-Tucker (KKT) condition. The result of the projection step can be written as the stationary point of the Lagrangian:
min ŷ max λ L(ŷ,λ) (4)
where we have the Lagrangian as:
L(ŷ,λ;y) = 1 2 (y − ŷ)T (y − ŷ) + λT (Ay − b) (5)
∂L ∂ŷ = y − ŷ + λTA (6)
The KKT condition leads to the following linear system, the solution of which involves solving a linear system of dimensions (m + n) × (m + n). Given that the linear system is symmetric and positive definite, the solution can be sought by inverting the system:[
I AT
A 0 ] [ ŷ λ ] = [ y b ] ⇒ [ ŷ λ ] = [ I AT A 0 ]−1 [ y b ] (7)
While this approach is general for enforcing arbitrary linear constraints on arbitrary network outputs, it is difficult to scale it to higher dimensions, and particularly difficult for 2-dimensional and 3- dimensional outputs, by direct matrix inversion followed by matrix multiplication.

heading: 4.2 SPECTRAL METHODS
text: First, we introduce and review the spectral methods (?) for discretizing the spatial PDE operators. Spectral methods are a class of numerical methods that computes the spatial partial derivatives of a field based on the spectral decomposition of the signal. By decomposing the original signal into a linear combination with respect to trigonometric basis functions of varying wavenumbers (or frequencies), the spatial derivatives with respect to the trigonometric basis functions can be easily and efficiently computed. The Fast Fourier Transform (FFT) is a well known algorithm for efficiently computing the Direct Fourier Transform (DFT) of uniform discrete signals. The multidimensional FFT and inverse FFT respectively compute the following:
F (k) = N−1∑ n=0 f (n)e−i2πk·(n/N); f (n) = 1 N1N2N3 N−1∑ k=0 F (k)ei2πk·(n/N) (8)
where F (k) = FFT(f (n)),f (n) = IFFT(F (k)), spatial indices n = (n1, n2, n3), nj ∈ {0, 1, · · · , Nj − 1} and spectral indices k = (k1, k2, k3), kj ∈ {0, 1, · · · , Nj − 1}. The spatial derivative with respect to xj can be computed by:
∂
∂xj f (n) = IFFT(ikjF (k)) (9)
In matrix form, for the t-th component of a 3 dimensional vector field: Ft, taking its flattened vector form, and taking the flattened vector form of the wavenumber kj corresponding to the dimension xj , the spatial derivative with respect to xj can be computed using matrix multiplication:
∂
∂xj Ft = diag(ikj)Ft (10)
where diag() converts a vector into a corresponding diagonal matrix. In general, arbitrary linear combination of spatial derivatives of varying orders can be computed using a single diagonal matrix multiplication:
( ∑ j ∑ r cjr ( ∂ ∂xj )r )Ft = diag( ∑ j ∑ r cjr(ikj) r)Ft := AtFt (11)
where At is a diagonal matrix for the spatial derivatives corresponding to the t-th component of the vector field, that is a polynomial of ikj , and A = [A1, A2, A3].

heading: 4.3 SPECTRAL PROJECTION LAYER
text: For brevity, we present our main results for computing the spectral projection operator that efficiently enforces spatial linear PDE constraints using spectral methods. We defer readers to Eqns (14 - 27) of Appendix A for detailed derivation of these results. In spectral space, the projection of the original vector field F into solution space: F̂ , can be computed by:
F̂ = PF +QB (12)
where F = FFT(f),B = FFT(b), and
P = I − 1∑3 j=0A 2 j  A21 A1A2 A1A3A1A2 A22 A2A3 A1A3 A2A3 A 2 3  ;Q = − 1∑3 j=0A 2 j [ A1 A2 A3 ] (13)

heading: 5 EXPERIMENTS AND RESULTS

heading: 5.1 COMPUTATIONAL COMPLEXITY AND COST
text: We first show that although the classic Lagrangian based hard constraining method in Eqn. 7 is general and able to enforce hard linear constraints, solving it by direct inversion leads to poor computational efficiency, especially with high-resolution 3-dimensional data outputs from CNNs.
Given Eqn. 7 for enforcing hard linear constraints using Lagrangian multipliers, we estimate the computational complexity regarding enforcing solenoidal conditions as follows. Without loss of generality, assume that the vector field on which we enforce the solenoidal constraints is 3 dimensional of resolution N in each spatial dimension, with a total of n = N3 nodes. The overall degrees of freedom in the system is 3n, and enforcing solenoidal constraint for each voxel results in n linear constraints, hence the resulting linear system in Eqn. 7 is of dimensions 4n × 4n. Though the matrix inversion is shared, hence reusable by caching, each projection involves a matrix multiplication of O((3n)2) ∼ O(9n2) operations. In comparison, the spectral projection method only involves element-wise operations, resulting in an overall complexity of O(n) operations for enforcing constraints, and O(n log(n)) for the FFT and IFFT operations. Results of an empirical analysis for computational time and memory usage is shown in Fig. 2.

heading: 5.2 TURBULENCE SUPERRESOLUTION WITH CONDITIONAL GAN
text: Conditional Generative Adversarial Networks While Generative Adversarial Networks (GANs) have been effective at generating 2D (???) and 3D (?) images, the unconditional generative modeling scenario of generating entire fields from random latent vectors is hardly useful in scientific settings. A more desirable model is one that is conditioned upon a set of inputs from either partial observations, or low resolution simulations due to computational limitations ?. Conditional GANs, which have been widely used in various image-to-image translation problems (?????), and a recent extension of conditional GANs, the GauGAN (?), utilizes a novel spatially-adaptive normalization layer that better preserves semantic information in the conditional input and produces improved texture outputs.
In this paper we use the GauGAN architecture for the task of superresolution of turbulent fluid flow fields (See Fig. 1).
Problem setup The main target application of this study is the super-resolution of turbulent flow fields. Fully resolving turbulence requires direct numerical simulation (DNS) that can resolve the smallest scales of the flow (Komogorov scale), which is prohibitively expensive. Therefore, the motivation of this study is to produce flow fields and flow statistics comparable to DNS at the cost of a low-resolution proxy, whilst strictly enforcing PDE constraints. To this end, we leverage highresolution DNS data to train a deep neural network to learn the mapping between the low-resolution flow and its high-resolution counterpart. We compare between several algorithms for the task: the conventional trilinear interpolation which is not learning based, and various deep learning methods leveraging the GauGAN architecture for conditional generative modeling. We benchmark our PDEL in-the-loop hard constraining method against unconstrained training (denoted by “none” in figures and tables) and soft constrained training (denoted by “soft”), with and without hard constraining spectral projection at test time. The goal is to satisfy the imposed constraint and to evaluate the accuracy of the predicted flow fields using key domain-specific metrics.
Dataset Description We use the Forced Isotropic Turbulence dataset from the Johns Hopkins Turbulence Database for this experiment ?. The dataset consists of DNS at 10243 resolution performed by solving the Navier-Stokes equations using the pseudo-spectral method. The dataset consists of 5028 frames (time steps) of data, each with 3 velocity components. For this experiment, we use all simulation frames starting from the 16th time step, since the initial frames consist of underdeveloped flow. Furthermore, since 10243 resolution is practically impossible to fit into modern GPUs, we use a subsampled version of the data at 1283 as high-resolution targets, and further subsampled 323 fields as low-resolution inputs. Subsampling for the high-resolution field is performed by uniformly sampling the original flow field at intervals of 8 in each dimension. For test and training splits, we take a random subset of 70/10/20% of the original data as training, validation and test splits.
Evaluation metrics Since the superresolution task constraining on the low-resolution input is mathematically underdetermined, it is not possible to recover the exact velocity field. Further, given the chaotic nature of turbulence, a single low-resolution flow field corresponds to various different realizations of high-resolution flows. Therefore, we refrain from directly comparing the norm of the difference in velocity fields. Instead, we compare the distributions of various key flow statistics, as outlined by ?, which are more informative from a turbulence modeling standpoint. In Tab. 1, we report the Kolmogorov-Smirnov (KS) statistic between the ground truth test set distributions and the distributions generated by various models conditioning on the low-resolution test inputs. We also report the mean difference between the distributions, i.e., difference between the mean of the modelled and ground truth distributions in units of ground truth distribution standard deviation.
The flow statistics in Tab. 1 is defined as below. For simplicity, we denote the different velocity components using Einstein notation, and use angle brackets <> to denote spatial averaging.
• Total kinetic energy, Etot = 12 < uiui >
• Dissipation, = 2ν < σijσij >, where σij := 12 ( ∂ui ∂xj + ∂uj ∂xi ) , and ν = 0.000185 is a
constant for fluid viscosity.
• Large eddy turnover time: TL = L/u′, where L = π2u′2 ∫ E(k) k dk and u ′ = √ (2/3)Etot
Results The main quantitative results for this experiment are presented in Tab. 1, whereas a visualization of the distributions regarding various key flow statistics are presented in Fig. 3. We observe from empirical evaluations that training with the hard constraining layer in-the-loop effectively imposes the solenoidal constraints (zero residue), and enforcing the hard constraint at training time achieves more accurate flow field distributions measured by various key flow statistics. We note that although this method is not the most accurate for the dissipation statistic, presumably because of discrepancies in the high wavenumber regime (where dissipation occurs), the overall mean statistics and individual statistics for the other metrics are superior compared to all the other methods.

heading: 6 CONCLUSIONS AND FUTURE WORK
text: Enforcing hard physical constraints to solutions generated using neural networks are essential for their application to important scientific problems. In this paper, we propose a novel differentiable spectral projection layer for neural networks that efficiently enforces spatial PDE constraints using spectral methods, yet is fully differentiable, allowing for its use as a layer in neural networks that supports end-to-end training. Further, we show that its computational cost is cheaper than a regular convolution layer. We demonstrate its use in an important class of physics problems – incompressible turbulent flows, where the divergence-free PDE constraint is required. We are able to train a 3D Conditional Generative Adversarial Network (CGAN) for turbulent flow superresolution efficiently, whilst guaranteeing the spatial PDE constraint of zero divergence. Further, our results show that the model produces more accurate flow statistics when trained with hard constraints imposed via the proposed novel differentiable spectral projection layer, as compared to soft constrained and unconstrained counterparts.
Some key limitations of this work are: (i) the method is applicable in its current form only to flows with periodic boundary conditions; (ii) we only develop a method for linear spatial constraints and (iii) we only consider statistically steady flows. In future we will address all the above limitations to extend our work to more general sets of nonlinear unsteady constraints with arbitrary boundary conditions.
APPENDIX

heading: A MATHEMATICAL DERIVATION FOR SPECTRAL PROJECTION
text: The solution to the Lagrangian multiplier method for enforcing the solenoidal conditions involves inverting the left-hand-side matrix in Eqn. 7. Since I, A, 0 are block matrices, the inverse can be represened by [
I AT
A 0
]−1 = [ I −AT (AAT )−1A AT (AAT )−1 (AAT )−1A −(AAT )−1 ]
(14)
Hence the projected vector in spectral space can be computed as:
F̂ = PF +QB (15)
The second term in the equation above drops out since b = 0 for solenoidal constraints. More specifically for spectral methods, the matrix A can be represented as three diagonal matrices for the wavenumbers in the three dimensions multiplied by the imaginary number i:
A =  . . . . . . . . .
A1 A2 A3 . . . . . . . . .  (16) = [A1 A2 A3] (17)
The only matrix inverse associated is AAT , whose value can be computed by block matrix multiplication:
AAT = A21 +A 2 2 +A 2 3 (18)
Given that A1, A2, A3 are diagonal matrices (eliminating the terms regarding the [0, 0, 0] mode), its inverse can be computed by directly inverting the diagonal terms:
(AAT )−1 = 1
A21 +A 2 2 +A 2 3
(19)
Hence the linear projection matrix can be written as:
I −AT (AAT )−1A = I − 1 A21 +A 2 2 +A 2 3  A21 A1A2 A1A3A1A2 A22 A2A3 A1A3 A2A3 A 2 3  (20) AT (AAT )−1 = − 1∑3
j=0A 2 j [ A1 A2 A3 ] (21)
Recovering the same solution as in Eqn. 13. More specifically for the divergence-free condition, we have:
Aj = diag(−ikj) (22) B = 0 (23)
Hence the spectral projection step can be further simplified as:
F̂ = F − k · F k · k k (24)
It is easy to show that the result is divergence-free, since:
−ikF̂ = −ikF + ikF = 0 (25) It is also easy to show that the projection is orthogonal to the solution space, since the dot product between the F̂ − F and F̂ is zero:
(F̂ − F ) · F̂ = −(k · F k · k k) · (F − k · F k · k k) (26)
= 0 (27)

heading: B MODEL AND TRAINING DETAILS
text: We use the GauGAN architecture (with schematics as shown in Fig. 1) for the conditional flow field generation task. The abbreviated names for the various modules are given in Tab. 2. Our model differs from the original GauGAN model in two distinct aspects. First, our architecture utilizes 3 dimensional convolutions instead of the 2 dimensional counterparts in the original GauGAN architecture. Second, for our hard constrained case, we append the spectral projection layer to the end of the architecture for enforcing hard constraints.
For training the model, we use multiresolution discriminator loss as in ? across 3 discriminators. We train the model with batch size of 18 (across 6 Volta V100 GPUs) with the Adam optimizer using a learning rate of 2E − 4. The soft constrained model uses a residue penalty factor of 0.01.

heading: C ADDITIONAL VISUALIZATION

Review summaries:
The paper describes a way to efficiently enforce physical constraints expressed by linear PDEs on the output of a neural network. The idea is to have, as a last layer of the network, a projection onto the constrained solution space, and to back-propagate through it. That projection layer is made efficient for high-dimensional outputs via the fast Fourier transform (FFT), exploiting a well-known numerical trick. Importantly, the proposed strategy is very general, and can indeed be used with any PDE constraint that is a linear combination of differential operators.

This work develops a differentiable spectral projection layer to enforce spatial PDE constraints using spectral methods, to achieve the introduction of the physical constraints in the end-to-end network without damaging the intrinsic property of the network. Analysis of computational cost shows the proposed layer is cheaper than the convolutional layer. The experimental comparison demonstrates the superiority of the proposed method. In my viewpoint, the novelty of this paper is somewhat novel.