PaperlayerView exact revision on arXiv

Source-linked AI summary

Meta-learning PINN loss functions

Apostolos F Psaros, Kenji Kawaguchi, George Em Karniadakis

arXiv:2107.05544v1cs.LG

TL;DR

PINN loss selection is difficult because training couples optimization choices with physics-based constraints across PDE tasks. The paper develops offline gradient-based meta-learning of loss functions using parametrized PDE task distributions and enforces desirable properties through regularization or parametrization. Shared offline-learned losses improve performance even on out-of-distribution tasks, while the paper also identifies limitations and design dependencies.

  • Problem

    PINN training involves difficult optimization and loss-selection choices, while existing online adaptive losses increase computational cost and do not use prior problem knowledge.

  • Method

    The paper learns PINN loss functions offline by differentiating PINN optimization paths across parametrized PDE task distributions and enforcing desirable properties through regularization or loss parametrization.

  • Results

    Shared offline-learned losses improve performance even for out-of-distribution meta-testing, including comparisons with online loss adaptation.

  • Takeaways & Limitations

    Offline loss learning can transfer across task distributions and outperform online adaptive loss approaches in the reported PINN examples.

  • Takeaways & Limitations

    The reported capabilities depend on algorithm design choices, loss parametrizations, and optimization behavior, including Adam-related limitations in some PINN examples.

Abstract

from arXiv · show

We propose a meta-learning technique for offline discovery of physics-informed neural network (PINN) loss functions. We extend earlier works on meta-learning, and develop a gradient-based meta-learning algorithm for addressing diverse task distributions based on parametrized partial differential equations (PDEs) that are solved with PINNs. Furthermore, based on new theory we identify two desirable properties of meta-learned losses in PINN problems, which we enforce by proposing a new regularization method or using a specific parametrization of the loss function. In the computational examples, the meta-learned losses are employed at test time for addressing regression and PDE task distributions. Our results indicate that significant performance improvement can be achieved by using a shared-among-tasks offline-learned loss function even for out-of-distribution meta-testing. In this case, we solve for test tasks that do not belong to the task distribution used in meta-training, and we also employ PINN architectures that are different from the PINN architecture used in meta-training. To better understand the capabilities and limitations of the proposed method, we consider various parametrizations of the loss function and describe different algorithm design options and how they may affect meta-learning performance.

1. Introduction

PINNs require coordinated choices about architectures, optimization, residual sampling, and loss terms, making PDE learning theoretically and practically challenging. Because loss functions shape optimization and are usually selected by trial and error, the paper motivates learning them automatically while incorporating task-specific information.

  • PINNs: PINNs construct a neural-network PDE approximator through automatic differentiation and minimize residual, boundary-condition, and initial-condition terms in a composite objective.
  • Motivation: PINN training requires selecting architectures, optimizers, learning-rate schedules, residual-point locations, and loss-term weights, creating a challenging learning problem.These choices govern both the neural approximation and enforcement of PDE, boundary, and initial-condition constraints.
  • Motivation: Loss functions interact with optimization and neural-network loss landscapes, affecting convergence rate and the performance of the resulting minimum.
  • Motivation: Choosing MSE, MAE, or another loss is an additional hyperparameter, and PINN practice commonly relies on trial and error with MSE as the popular option.
  • Related work: Existing parametrized adaptive losses include standard losses and can be optimized online, but they increase training cost and do not use prior knowledge of the problem.
  • Related work and motivation: The paper frames the open question as whether physics from parametrized PDE task distributions can be encoded into PINN loss functions.

PDE in a loss function optimized offline?

The paper develops an offline, gradient-based meta-learning procedure that discovers shared PINN loss functions for diverse parametrized PDE task distributions. Meta-training differentiates through short PINN optimization paths, after which the learned loss is used on unseen tasks; the method is supported by theory, regularization, alternative parametrizations, and computational benchmarks examining performance and limitations.

  • Overview of the proposed method: Meta-training alternates short inner PINN optimizations with outer updates of loss parameters based on the MSE of optimized PINN parameters.The learned loss is evaluated on task-specific PDE, boundary, and initial-condition data in the inner loop, while outer updates use validation residuals and solution data when available.
  • Meta-testing: After meta-training, the learned loss is applied to unseen tasks until convergence, including tasks or architectures that may differ from those used during meta-training.The computational examples evaluate the learned loss at meta-testing rather than continuing to adapt it online for each test task.
  • Overview of the proposed method: The proposed method performs offline discovery of PINN loss functions for diverse PDE task distributions using gradient-based meta-learning.PDE tasks are sampled during meta-training, and the method is explicitly designed to learn losses shared across tasks.
  • Computational evaluation: The study compares alternative loss parametrizations and algorithm design options, and defines benchmarks to assess performance, applicability, and limitations.The experiments vary loss parametrizations and other design choices while considering diverse PDE task distributions.
  • Theory and loss design: The paper identifies two desirable learned-loss properties—optimal stationarity and an MSE relation condition—and proposes regularization or parametrization choices to enforce them.The parametrization proposed in prior work is stated to guarantee both properties, while the paper also introduces a new regularization method.

2. Preliminaries

PINNs solve parametrized forward or inverse PDE problems by minimizing composite residual and data losses, with gradients determined by the chosen discrepancy function.

  • 2.1. PINN solution technique overview: PINNs represent PDE solutions with neural networks and minimize residuals for the PDE, boundary conditions, and initial conditions over discrete points.Inverse problems additionally include solution-data misfit and jointly learn the PDE parameters.
  • 2.1. PINN solution technique overview: The composite PINN objective sums losses for physics, boundary conditions, initial conditions, and optional solution data, with weights that are often selected by trial and error.The optimal weights are not known a priori and are frequently set equally or obtained through experimentation.
  • 2.2. Loss functions in neural network training: The loss function maps predictions and targets to nonnegative discrepancies, and its gradient determines how prediction errors affect neural-network parameter updates.For multidimensional outputs, the loss gradient also controls how discrepancies across dimensions contribute to the final gradient.
  • 2.2. Loss functions in neural network training: Squared ℓ2 discrepancy treats output dimensions uniformly and reduces the objective terms to mean squared errors when averaged over datasets.The paper uses MSE and squared ℓ2 discrepancy interchangeably depending on context.
  • 2.2. Loss functions in neural network training: SGD updates network parameters using first-order objective information, while algorithms such as AdaGrad and Adam also depend on gradient history.The parameter-to-output Jacobian connects prediction gradients to parameter gradients.

3. Meta-learning loss functions for PINNs

This section formulates PINN loss-function discovery as gradient-based meta-learning across related parametrized PDE tasks. It uses a bi-level objective to optimize shared hyperparameters according to average task performance.

  • 3. Meta-learning loss functions for PINNs: Meta-learning addresses repeated design choices in PINN training by learning settings expected to perform well across related parametrized PDE tasks.These choices affect optimizer convergence and training and test error, while task distributions are formed from related problems indexed by λ.
  • 3.1. Defining PDE task distributions: A task distribution samples related learning problems through λ drawn from p(λ), allowing information from representative tasks to support efficient solution of others.The paper distinguishes this shared-task perspective from repeatedly tuning hyperparameters independently for each novel PDE.
  • 3.2. Meta-learning as a bi-level minimization problem: The meta-learning problem is bi-level: inner optimization trains PINNs with a task-specific base objective, while outer optimization adjusts shared hyperparameters using meta-performance.The outer objective is commonly the average performance across tasks, such as validation PDE-residual error.
  • 3.2. Meta-learning as a bi-level minimization problem: The proposed algorithm alternates outer and inner optimization, requiring differentiable objectives and inner updates so automatic differentiation can compute the meta-gradient.The framework is presented for arbitrary admissible η and supports multiple SGD steps and other differentiable optimizers.
  • 3.2. Meta-learning as a bi-level minimization problem: Gradient-based meta-learning differentiates the outer objective through the inner optimization path to update hyperparameters that influence the trained PINN parameters.The total derivative contains direct dependence on η and dependence through the optimal inner parameters θ*(η).

13 end

The paper formulates PINN loss-function meta-learning as bilevel optimization and derives conditions intended to make learned losses train effectively and align with MSE objectives. It compares flexible parametrizations, theory-driven regularization, and an adaptive-loss parametrization that satisfies the conditions automatically.

  • Loss parametrizations: The framework supports FFN and adaptive-loss parametrizations, with adaptive losses using robustness and scale parameters and FFNs representing loss weights and biases.The adaptive parametrization imposes α and c constraints through sigmoid and softplus transformations, while the LAL parametrization is distinct from online adaptive use.
  • Meta-learning formulation: The method learns a shared parametrized loss by optimizing task-specific PINN parameters in an inner loop and validation MSE averaged across tasks in an outer loop.Additional solution data can be included in the outer objective even when inner training uses only PDE residual and boundary or initial-condition data.
  • Meta-learning formulation: The outer objective evaluates whether a learned loss improves average task MSE, while optionally measuring solution errors using additional validation or solution data.Validation point counts may differ from training counts, and the outer objective uses MSE whereas the inner objective uses the learned loss.
  • Loss parametrizations: The adaptive loss cannot be optimized online directly in the cited manner because its loss is monotonic in α, although the outer objective avoids the trivial α solutions encountered in online optimization.The paper distinguishes the meta-learning use of the parametrization from its online adaptive-loss use.
  • Theoretical properties: The theory identifies optimal stationarity and MSE relation as desirable conditions that support efficient training and connect stationary points of the learned loss to global MSE minima.Under stated assumptions, the theorems establish global-minimum guarantees for stationary or limit points, and the corollary extends this conclusion to the MSE loss.
  • Theory-driven design: A novel outer-objective penalty can enforce the desirable conditions, whereas the LAL parametrization satisfies both conditions without regularization.The paper also proves that the LAL form automatically satisfies the required properties.

4. Computational examples

The computational examples evaluate learned PINN losses across regression and PDE task distributions, varying parametrizations, optimizers, and testing conditions. LAL often generalizes strongly in regression and some PDE settings, while performance depends on parametrization, objective-weight learning, optimizer, and task distribution.

  • 4.1. Discontinuous function approximation with varying frequencies and heteroscedastic noise: LAL losses achieved at least 15% lower average minimum rl2 error than all other considered losses on 10 OOD function-approximation tasks.LAL 4 reached approximately 17% error near iteration 20,000, versus approximately 20% for MSE near iteration 50,000.
  • 4.1. Discontinuous function approximation with varying frequencies and heteroscedastic noise: FFN losses did not generalize well in function approximation, whereas LAL losses performed better than the standard and adaptive losses considered.
  • 4.2. Task distributions defined based on advection equation with varying initial conditions and discontinuous solutions: Meta-learning objective-function weights initially improved advection performance, but final learned losses combined with final learned weights eventually deteriorated it.The learned weights increased for both parametrizations, while FFN and LAL differed in how they balanced initial-condition residuals.
  • 4.2. Task distributions defined based on advection equation with varying initial conditions and discontinuous solutions: Higher Adam β values, indicating greater dependence on the far past, deteriorated learned-loss performance, and some learned losses failed to improve over MSE.
  • 4.3–4.4. Reaction-diffusion and Burgers equations: For another PDE task distribution, FFN losses generalized across a different architecture and task distribution, whereas LAL losses did not generalize well.The FFN losses were significantly better than most standard losses after 20,000 iterations despite only 20 meta-training iterations.

5. Summary

The paper presents offline meta-learning for PINN loss discovery, supported by theoretical conditions and a regularization strategy. Experiments show benefits for shared losses under some out-of-distribution settings, while generalization remains sensitive to optimizers and parametrizations.

  • 5. Summary: The method discovers shared PINN loss functions offline for diverse parametrized PDE task distributions using alternating inner PINN optimization and outer loss-parameter updates.The outer update uses the MSE of final semi-optimized PINN parameters.
  • 5. Summary: The theory identifies optimal stationarity and MSE-relation conditions, and a proposed regularizer enforces these conditions for learned losses.The LAL parametrization satisfies both conditions automatically without regularization.
  • 5. Summary: Shared offline-learned losses produced significant improvements even on OOD tasks and with PINN architectures different from those used during meta-training.Improvement was also reported relative to online loss adaptation.
  • 5. Summary: The experiments cover discontinuous regression, advection, reaction-diffusion, and Burgers task distributions while varying loss parametrizations and meta-learning design choices.
  • 5. Summary: Neither FFN nor LAL losses showed satisfactory generalization in the PINN examples with Adam, although LAL outperformed MSE in the function-approximation example.The authors attribute one limitation to Adam's dependence on gradient history.

A. Loss functions with multi-dimensional inputs

The appendix constructs increasingly expressive loss functions for vector-valued PINN outputs. These range from sums of coordinatewise losses to functions that jointly process prediction and target vectors.

  • A. Loss functions with multi-dimensional inputs: For PINN residual terms, the vector-valued loss compares predicted PDE, boundary-condition, or initial-condition outputs with their corresponding targets and returns a scalar.
  • A. Loss functions with multi-dimensional inputs: A multidimensional loss replaces the scalar squared error with a parametrized function applied to directional prediction discrepancies.
  • A. Loss functions with multi-dimensional inputs: Directional losses can be weighted by coefficients a_j, allowing meta-learning to normalize each direction separately.
  • A. Loss functions with multi-dimensional inputs: Greater expressiveness comes from using a different parametrized loss for each output dimension instead of sharing one across dimensions.
  • A. Loss functions with multi-dimensional inputs: The most expressive construction feeds the full predicted and target vectors into a general parametrized loss rather than processing coordinatewise discrepancies.

B.1. Proof of Theorem 1

Theorem 1 shows that, under the stated expressivity and loss assumptions, the optimal stationarity condition rules out non-global stationary points of the learned-loss objective.

  • B.1. Proof of Theorem 1: The proof rewrites the objective gradient as a matrix-vector product and uses full row rank to infer that the loss-gradient vector is zero.
  • B.1. Proof of Theorem 1: Under the theorem's assumptions, any stationary point of the learned-loss objective is a global minimum when the loss satisfies the optimal stationarity condition.
  • B.1. Proof of Theorem 1: The argument uses the assumption that the network can realize arbitrary collections of output values at the datapoints.
  • B.1. Proof of Theorem 1: The proof constructs a stationary point that is not a global minimum when the learned loss has a nonoptimal stationary point.

B.2. Proof of Theorem 2

The proof establishes that, under either learning-rate case, gradient descent reaches vanishing gradients; with the theorem’s rank condition, any limit point is a global minimum.

  • The auxiliary lemma proves Lipschitz continuity along line segments in an open convex domain, supporting the convergence argument.
  • The loss sequence decreases monotonically and converges to a finite value because the loss is nonnegative.
  • The proof derives lim inf and lim sup gradient norms equal to zero, completing the second learning-rate case.
  • Under the rank condition on the network Jacobian, any limit point of the parameter sequence is a global minimum of L.

B.3. Proof of Corollary 1

The corollary shows that if the loss gradient vanishes exactly at prediction targets, every stationary point achieves the global minimum of the MSE loss.

  • At any stationary point, the network predictions equal the targets at all training points, so LMSE reaches its global minimum value.
  • The proof uses the condition ∇qℓ(q,u)=0 if and only if q=u to convert stationarity of the loss into exact prediction-target agreement.

B.4. Proof of Proposition 1

The proposition verifies that the proposed loss satisfies the optimal stationarity condition: zero prediction error is exactly characterized by zero loss gradient and minimizes the loss.

  • The loss satisfies ∇qℓ(q,u)=0 if and only if q=u, establishing the required stationarity property.
  • When q=u, the loss is no larger than at any alternative q′, so exact prediction agreement is globally optimal.
  • Together, these properties establish the proposition’s optimal stationarity condition for the parametrized loss.

C. Other algorithm design options

The paper discusses design choices for task and initialization sampling, stopping criteria, and snapshot use, while noting that induced randomness can make training and validation metrics noisy.

  • Sampling options: Resampling tasks and randomly initialized networks exposes the learned loss to broader task and initialization distributions, but can destabilize training.
  • Sampling options: Tasks and network initializations may be resampled every I′ and I′′ outer iterations instead of at every iteration.
  • Stopping criteria: Outer-loop stopping can use meta-training or meta-validation error, although either metric may be noisy under some design choices.
  • Stopping criteria: When the outer objective is especially noisy, a less noisy metric or moving average can be used, particularly with T = 1.
  • Experimental practice: The experiments train for 10,000 outer iterations and use all captured learned-loss snapshots during meta-testing.

D. Additional computational results related to the function approximation example

The function-approximation experiments examine initialization stability, gradient handling, and design choices for meta-training and meta-testing learned loss functions. Increasing inner iterations improves test performance and robustness, while task resampling and approximator re-initialization introduce training noise.

  • Effect of loss function initialization: A gradient explosion at outer iteration 4 causes a large jump in the loss-parameter norm under random initialization without clipping.The experiments therefore use gradient clipping to address exploding gradients rather than uniformly dividing gradients by the number of inner iterations.
  • Effect of loss function initialization: The loss-gradient norm does not necessarily increase with the number of inner iterations, despite differentiation through a longer optimization path.This motivates treating gradient growth and inner-step count as separate design considerations.
  • Design options experiment: Task resampling and approximator re-initialization introduce noise into meta-training, with resampling every 10 outer iterations producing noisier trajectories than no resampling.The design experiment varies inner steps, task-resampling frequency, and whether the approximator is re-initialized.
  • Design options experiment: Increasing inner iterations improves learned-loss test performance and robustness across outer iterations and design options.The same patterns appear for 100 and 500 test iterations, supporting consistent design-option selection.
Loading 2107.05544v1…