Multi-Objective Loss Balancing for Physics-Informed Deep Learning

In plain words

A physics-informed neural network learns a function that has to obey a physical rule, such as how heat or a wave spreads. Training it means adding several penalties together: one for breaking the rule inside the domain and one for each boundary or starting condition. These penalties often differ hugely in size, so the largest one dominates training and the others are ignored.

The paper compares ways of giving each penalty the right weight. The authors' method, ReLoBRaLo, looks at how much each penalty has dropped since the last training step and since the start, and turns those drops into weights with a softmax. It sometimes forgets its recent history and looks back to the start instead, which helps it escape stalls. It does not need the extra backward passes that some earlier methods use, so it adds little training time.

TL;DR

• PINN training is a multi-objective problem: the PDE residual and each boundary, initial or data term are separate losses that must be scaled against each other §III eq. 2 eq. 5.
• ReLoBRaLo sets each term's scaling from its relative progress, using only loss values, with exponential decay and a random lookback §VI eq. 11.
• The paper reports the lowest forward validation error on Burgers and Kirchhoff, and the lowest inverse error on Kirchhoff and Helmholtz; GradNorm is better on forward Helmholtz tab. II tab. IV tab. VI.
• Its overhead is lower than GradNorm's and LR annealing's on all three problems, and slightly above SoftAdapt's tab. III tab. V tab. VII.
• The code differs from the paper in several places that change results, including the activation function and the GradNorm and LR annealing formulas. Read Paper vs code before implementing.

Problem

A PINN approximates the solution of a parameterised PDE with a network whose loss sums terms for the PDE residual, boundary conditions (BC), initial conditions (IC) and, optionally, measurements §III eq. 1 eq. 2. Summing these terms with fixed weights is a linear scalarisation of a multi-objective problem §IV eq. 4 eq. 5. The scales of the terms can differ by orders of magnitude, which makes the gradients unbalanced and training stall; manual weights need a grid search that becomes impractical as the number of terms grows §I §IV.

The paper considers forward problems (PDE parameters fixed) and inverse problems (a PDE parameter is trained together with the network) §III. Its test problems are Burgers' equation, Kirchhoff plate bending and Helmholtz §VIII-A §VIII-B §VIII-C.

Method

Terms and the scalarised loss

The loss is a weighted sum of the PDE residual term L_0 and m boundary, initial or data terms §III eq. 2 eq. 5. Each term is a mean of squared residuals over a batch of sampled points §III eq. 2. Weights λ_i are chosen per step by one of the balancing rules below.

Baselines

• Learning rate annealing (LRA) §V-A eq. 6. The scaling of each boundary term is the maximum absolute PDE-gradient entry divided by the mean absolute gradient of that boundary term, then smoothed with an exponential average (α = 0.9 recommended by the original authors §V-A). It needs one backward pass per term.
• GradNorm §V-B eq. 7 eq. 8 eq. 9. The scalings are trainable. Each term's gradient norm is pulled towards the average gradient norm times its relative training rate raised to a power α; the rate is measured against the term's initial loss eq. 7. The scalings are updated by a separate optimiser. The paper gives no value for α in §V-B. It needs one backward pass per term.
• SoftAdapt §V-C eq. 10. Each scaling is a softmax over T times the change of that term's loss since the previous step, with no decay. It uses loss values only.

ReLoBRaLo

Per training step t, for each term i (term 0 is the PDE residual) §VI eq. 11:

1. Compute the progress since the previous step, λ_bal(t, t−1): a softmax over L_i(t) / (T · L_i(t−1)), multiplied by the number of terms m §VI eq. 11.
2. Compute the progress since the start of training, λ_bal(t, 0): the same softmax with L_i(0) in place of L_i(t−1) §VI eq. 11.
3. Draw ρ once per step from a Bernoulli distribution with mean E[ρ], shared by all terms §VI train.py.
4. History: λ_hist = ρ · λ_i(t−1) + (1 − ρ) · λ_bal(t, 0). If ρ = 1 the previous scalings are kept; if ρ = 0 the scalings restart from the progress since the start §VI eq. 11.
5. Update: λ_i(t) = α · λ_hist + (1 − α) · λ_bal(t, t−1) §VI eq. 11.
6. Treat the λ values as constants in the parameter gradient §VI.

# ReLoBRaLo: scalings for one training step  [§VI] [eq. 11]
# losses: (m,) current term losses, term 0 = PDE residual
# state.lambdas, state.prev_losses, state.init_losses: (m,)
# T: temperature; alpha: decay rate; rho: 0 or 1 drawn this step

def relobralo_scalings(losses, state, T, alpha, rho):
    # m = number of terms in the loss
    m = losses.shape[0]
    # Progress since the previous step, softmax-normalised  [eq. 11]
    progress_prev = m * softmax(losses / (T * state.prev_losses))
    # Progress since the start of training  [eq. 11]
    progress_init = m * softmax(losses / (T * state.init_losses))
    # Random lookback: keep history if rho = 1, else use progress_init  [§VI]
    history = rho * state.lambdas + (1 - rho) * progress_init
    # Exponential decay between history and current progress  [eq. 11]
    lambdas = alpha * history + (1 - alpha) * progress_prev
    # Scalings are constants for the parameter update  [§VI]
    lambdas = stop_gradient(lambdas)
    # Keep state for the next step
    state.lambdas = lambdas
    state.prev_losses = stop_gradient(losses)
    return lambdas

The paper states the softmax can overflow when the exponent is large, and suggests subtracting a large constant (e.g. 10⁻⁹) from the softmax input §VI. The code takes a different route (see Paper vs code).

Schedule of the reference code

The following steps are in the reference training loop but not in the paper. They change the first two steps of ReLoBRaLo:

# Per-step settings, reference code  [code: src/train.py#L142-L150]
# Not stated in the paper; see Paper vs code.
def settings_for_step(step, configured_alpha):
    # alpha is 1 at step 0, 0 at step 1, then the configured value
    if step == 0:
        return 1.0
    if step == 1:
        return 0.0
    return configured_alpha

# Initial losses L_i(0) are taken at step 1, not step 0
if step == 1:
    state.init_losses = losses

Training setup used for the experiments

• Adam, learning rate 10⁻³ in the final settings, up to 10⁵ steps of gradient descent §VIII tab. VIII.
• On a plateau the learning rate is multiplied by 0.1; early stopping ends the run §VIII.
• The network is fully connected with tanh activation in the final settings tab. VIII. Layer and width counts differ between the table and the figure captions (see Pitfalls).
• For inverse problems the PDE parameter μ is a trainable scalar, initialised to 0.5 §VIII-B §VIII-C train.py.

Burgers' equation §VIII-A

Domain x ∈ [−1, 1], t ∈ [0, 1]; viscosity ν = 1/(100π) eq. 12 §VIII-A. Initial condition u(0, x) = −sin(πx); boundary u(t, ±1) = 0 eq. 12. Forward loss: PDE residual, two boundary terms and one initial term eq. 13 eq. 14. Inverse loss: PDE residual with ν as a trainable variable, plus a data term against the reference solution eq. 15. The reference solution is a finite-element solution, which the paper does not ship §VIII-A.

Kirchhoff plate bending §VIII-B

Plate width a = 10 m, length b = 10 m, load p₀ = 0.015 MN/m², Young's modulus E = 30 000 MN/m², thickness h = 0.2 m, Poisson ratio ν = 0.2 §VIII-B. Flexural stiffness D = E h³ / (12 (1 − ν²)) ≈ 20.83 eq. 17 §VIII-B. The governing equation is ∇⁴u − p/D = 0 eq. 17. The analytic solution for a sinusoidal load is given in eq. 18. There are four simply-supported boundary conditions (deflection zero) and four moment-free conditions eq. 19; bending moments are defined in eq. 20. The total loss has nine terms: one PDE term and eight boundary terms eq. 21. The inverse problem learns D from data sampled from the analytic solution, starting from μ = 0.5 §VIII-B.

Helmholtz equation §VIII-C

Δu + k²u = f on [−1, 1]², with f and u given in eq. 22 and eq. 23 for k = 1 §VIII-C. Four Dirichlet boundary terms and one PDE term give five losses eq. 24. The inverse problem learns k from data, starting from μ = 0.5 §VIII-C.

Hyperparameter search

The paper uses Bayesian optimisation (20 random points, then 80 points by expected improvement) over the ranges in tab. I, followed by a fine grid search §VII. The final values are in tab. VIII.

Hyperparameters

Results

Author-reported. Forward problems: validation error of u; inverse problems: validation error of the PDE parameter. Each cell is the median over four independent runs unless noted. The baseline uses manually chosen scalings, which the paper finds by grid search but does not list §VIII-A tab. II.

Burgers tab. II. Setup: fully connected network, see Pitfalls for the size.

Kirchhoff tab. IV.

Helmholtz tab. VI. The table does not state which E[ρ] the ReLoBRaLo column used; tab. VIII gives 0.99.

Computational overhead per 1000 optimisation steps, in seconds, relative to no balancing (baseline run time in brackets) tab. III tab. V tab. VII:

Ablation over the decay rate α, temperature T and E[ρ] fig. 11 tab. IX. The validation error of Helmholtz is lowest at E[ρ] = 0.99 (2.6·10⁻⁵), and that of Burgers and Kirchhoff at E[ρ] = 0.9999 (1.4·10⁻⁴ and 4.0·10⁻¹⁰). Setting E[ρ] = 1, which removes the lookback, gives worse values on all three tab. IX. The table caption says it is the median of three runs, not four.

Paper vs code

Repository: https://github.com/rbischof/relative_balancing, default branch, commit b3c76d2bed7c6bebb2e2628575008a04858472cf. The paper links it in its introduction §I. The owner is the first author. The repository has no LICENSE file and the README states no license, so no code is quoted and every point below is cited by path and line only. Other pieces of the repository (a data/ folder with burgers.mat, the other task files) are missing or are not used by the paper; see Pitfalls.

• Activation. The paper's final settings use tanh tab. VIII. The network defaults to tf.math.sin models.py, and train.py does not pass another activation train.py. The Bayesian search in the code has no activation parameter, though tab. I lists tanh and sigmoid. bayesian_optimization.py.
• Output layer. The output is a sigmoid models.py. The paper does not state an output activation.
• Residual connections. Hidden layers after the first add their input back (+ u) models.py. The paper describes a plain fully-connected network §III.
• Input scaling. Inputs are mapped to [−1, 1] from the domain bounds models.py. Not stated in the paper.
• LRA gradient statistic. The code uses the mean absolute PDE gradient update_rules.py, not the maximum that eq. 6 writes. It also adds 10⁻⁸ to the denominator update_rules.py.
• GradNorm initial loss. The code divides each loss by its value at the previous step, not by L_i(0) update_rules.py [code: src/update_rules.py#L112]. The initial-loss values are stored under other keys that GradNorm never reads train.py. This departs from eq. 7.
• GradNorm gradient. The gradient norm is taken only with respect to the second-last trainable tensor, which is the output layer's kernel in the layer order of models.py update_rules.py. The paper's G_i uses the network parameters θ eq. 7 §V-B.
• GradNorm restoring force. The code uses the --T flag for GradNorm's α train.py. The --alpha flag has no effect on GradNorm.
• GradNorm weights. The output weights are a softmax over the terms, so they sum to 1 models.py. The paper does not state a normalisation.
• SoftAdapt. The code multiplies the softmax by the number of terms update_rules.py, which eq. 10 does not. The decay rate stays at 0 train.py train.py, so --alpha is ignored.
• ReLoBRaLo softmax stabilisation. The code adds 10⁻¹² to the denominator update_rules.py, not the 10⁻⁹ subtraction the paper suggests §VI.
• ReLoBRaLo initial losses. L_i(0) is stored at step 1 [code: src/train.py#L142-L144], not step 0.
• ReLoBRaLo decay schedule. α is 1 at step 0, 0 at step 1 and the configured value after that train.py train.py. Neither the paper nor eq. 11 mentions this.
• ReLoBRaLo lookback. ρ is a per-step Bernoulli with mean --rho train.py. This matches §VI.
• Learning-rate patience. The paper reduces the rate after 3000 steps with no progress §VIII. The code's default patience is 4 evaluations train.py, and each evaluation is 1000 epochs train.py. The README and the search script use 3 README.md bayesian_optimization.py.
• Progress measure. The code tracks the mean training loss over the last 1000 steps, not validation loss train.py [code: src/train.py#L177]. The paper does not say which loss is tracked.
• Early stopping. The paper stops after 9000 steps without improvement §VIII. The code stops after five learning-rate reductions in a row train.py, counting reductions since the last new best training loss, and it also stops after 3.75 hours train.py. Each reduction also restores the best weights train.py.
• Kirchhoff residual scale. The forward PDE term is D∇⁴w − p₀ sin sin kirchhoff.py, which is D times the form in eq. 17 eq. 21. The inverse term uses ∇⁴w − p/D kirchhoff.py. The forward PDE loss is therefore about D² ≈ 434 times larger than the paper's form, so the same λ values mean different things.
• Helmholtz inverse. The residual uses the fixed target k in the operator helmholtz.py, and the learned k̂ only in the source term helmholtz.py. The data term is (k²·sin − k̂²·u)² helmholtz.py, which the paper does not write out.
• Burgers inverse. The random batch is discarded: the inverse loss uses the full data grid, subsampled every fourth x point burgers.py, and the random indices drawn for the batch burgers.py are not used by it. The paper says it samples from the data at every iteration §VIII-A.
• Two optimisers. The paper reports that two separate optimisers, one for θ and one for μ, were tested and helped on Kirchhoff §VIII-B. The code uses one optimiser for both train.py. The second optimiser exists only for GradNorm's weights train.py [code: src/train.py#L139-L140].
• Baseline scalings. The paper's baseline uses grid-searched scalings §VIII-A. The code's manual rule uses λ = 1 for every term train.py.
• Bayesian search. The paper uses 20 random points §VII. The script uses 15 initial points and 80 iterations bayesian_optimization.py. It fixes 3 layers and 256 nodes bayesian_optimization.py. Its T range is 10⁻³ to 10³ bayesian_optimization.py [code: src/bayesian_optimization.py#L19], not 10⁻⁶ to 10² tab. I. The fine grid search is not in the code.
• Defaults. Batch size 1024 train.py, learning rate 10⁻³ train.py, α = 0.999 train.py, T = 1 train.py, and 10⁵ epochs train.py. Fixed collocation points unless --resample is set train.py. These match the paper's stated values where it gives them.
• Initial μ. Every task starts the trainable parameter at 0.5, including Burgers train.py. The paper states 0.5 only for Kirchhoff and Helmholtz.
• Extra tasks. The repository also has Allen–Cahn, Poisson-L and diffusion–sorption tasks train.py. None of them is in the paper; they are not used here.

Pitfalls

• Output activation and negative targets. The output is a sigmoid, so predictions lie in (0, 1) models.py. The Burgers initial condition −sin(πx) is negative on (0, 1) burgers.py, and the Helmholtz solution sin(πx)sin(4πy) is negative in parts [code: src/tasks/helmholtz.py#L38]. A sigmoid output cannot fit these as written. The paper reports fitted solutions and does not state its output activation. Settle this before implementing; the code as published does not reproduce the stated setup.
• Architecture conflicts within the paper. tab. VIII gives Burgers 4 layers of 256, Kirchhoff 4 of 360 and Helmholtz 2 of 256. The figure captions give Burgers two layers of 128 fig. 2, Kirchhoff three layers of 128 fig. 5 and Helmholtz two layers of 128 fig. 8. The default run has one layer of 32 nodes train.py.
• Table II text claim. The text calls ReLoBRaLo better than the baseline on inverse Burgers §VIII-A. Its validation μ error is 2.2·10⁻¹⁰, and the baseline's is 1.9·10⁻¹⁰ tab. II. Treat the two as equal.
• Burgers reference data. The inverse Burgers and Burgers validation need data/burgers.mat burgers.py. That file is not in the repository, and src/data is excluded by .gitignore .gitignore. The finite-element reference is not reproducible from the repository alone.
• Helmholtz forward GradNorm. GradNorm has the lowest forward validation error on Helmholtz (5.6·10⁻⁶) tab. VI, so the "lowest error" claim does not hold for that setting.
• Overhead ranking. SoftAdapt is cheaper than ReLoBRaLo on all three problems tab. III tab. V tab. VII. ReLoBRaLo's advantage is accuracy over SoftAdapt, not cost.
• Stochastic runs. Four runs for most tables, three for tab. IX. Medians and standard deviations come from few runs and are sensitive to the seed; no seed is stated.
• Numerical stability. The exponentials in eq. 11 can overflow when a term's loss is much larger than its previous value divided by T. The paper's fix is to subtract a constant from the softmax input §VI; the code's fix differs (see Paper vs code). Compute the softmax with the usual max-shift.
• First two steps. Use the step-0 and step-1 α values in the code schedule if you want to match the reference runs; the paper's equations do not say them.
• Gradient cost. LRA and GradNorm need one backward pass per term, so their overhead grows with the number of terms: Kirchhoff has nine terms and the overhead is largest there tab. V. ReLoBRaLo needs none §VI.
• Temperature and scale. Small T gives near-argmax scalings, and the paper's best Helmholtz value is 10⁻⁵ §IX tab. VIII. The three problems use T values four orders of magnitude apart, so T does not transfer between problems tab. VIII.
• Table numbering. IEEEtran numbers tables in Roman numerals and sections as I to X, with subsections such as V-A. This card cites them that way (e.g. tab. II, §VIII-A). Figures and equations are Arabic.
• Version. The card reflects arXiv v2 (updated 2022-11-28). v1 was not read.

Lineage

• Builds on: learning-rate annealing for PINNs §V-A eq. 6; GradNorm from computer vision §V-B eq. 7 eq. 8 eq. 9; SoftAdapt §V-C eq. 10. ReLoBRaLo takes SoftAdapt's loss-only signal and GradNorm's progress ratio, and adds LRA-style decay and the random lookback §VI.
• Compared against: GradNorm, LRA, SoftAdapt, and a manually scaled baseline §VIII tab. II tab. IV tab. VI.
• Future work named in the paper: combining ReLoBRaLo with Sobolev training or with a mixture-of-experts approach (cited as bischof2022mixture) §X.