Skip to content

Loss Reference¤

Training Neural Operators effectively requires evaluating the error in function space rather than just on discrete grid points. Standard ML losses (like basic MSE) are discretization-dependent. To ensure our models are resolution-invariant and respect the underlying continuous functions, neojax provides discretization-invariant, norm-based losses.

Lp-Losses¤

The standard approach for training Neural Operators is using continuous \(L^p\) norms.

  • LpLoss: Computes the standard \(L^p\) norm of the error over the spatial domain.
  • RelativeLpLoss: Computes the \(L^p\) error normalized by the \(L^p\) norm of the true target function. This is standard practice in operator learning because physical data often ranges across different magnitudes, and a relative loss ensures stable gradients across samples.

neojax.losses.lp_losses.LpLoss ¤

General \(L^{p}\)-loss.

Computes $$ \Vert y - \hat{y}\Vert_{L^p} = \Bigg(\int_{\Omega} \Big| y - \hat{y} \Big|^p dx \Bigg)^{\frac{1}{p}} $$ where \(y\) is the ground truth and \(\hat{y}\) is the model prediction.

Parameters:

Name Type Description Default
p float | int

Power of the norm.

 required
weight float

(Learnable) weight. Loss is computed as weight * loss. Default is 1.0.

 1.0
learnable_weight bool

Whether weight is learnable. Used to filter trainable parameters using is_learnable_loss_weight utility function with equinox.filter_... or equinox.partition. Default is False.

 False

Internal Attributes

These fields store the internal state of the loss.

  • weight (Float[Array, ""]): Learnable loss weight. Filter during training to prevent updates.
  • p (float | int): Power of the norm.
  • learnable_weight (bool): Flag indicating whether weight is learnable.

neojax.losses.lp_losses.RelativeLpLoss ¤

General relative \(L^{p}\)-loss.

Computes $$ \frac{\Vert y - \hat{y}\Vert_{L^p}}{\Vert y \Vert_{L^p}} $$ where \(y\) is the ground truth and \(\hat{y}\) is the model prediction.

Parameters:

Name Type Description Default
p float | int

Power of the norm.

 required
weight float

(Learnable) weight. Loss is computed as weight * loss. Default is 1.0.

 1.0
learnable_weight bool

Whether weight is learnable. Used to filter trainable parameters using is_learnable_loss_weight utility function with equinox.filter_... or equinox.partition. Default is False.

 False

Internal Attributes

These fields store the internal state of the loss.

  • weight (Float[Array, ""]): Learnable loss weight. Filter during training to prevent updates.
  • p (float | int): Power of the norm.
  • learnable_weight (bool): Flag indicating whether weight is learnable.

Sobolev Losses¤

While \(L^p\) losses only measure the error in the predicted function values, Sobolev losses (\(H^1\), \(W^{k,p}\)) also penalize errors in the derivatives of the functions. By aligning the gradients of the prediction with the true gradients, the network is forced to learn the underlying physical differential operators more accurately.

Implementation Details and Efficient Usage¤

Computing derivatives in Automatic Differentiation (AD) for Neural Operator training requires careful consideration of input and output dimensions. Choosing the wrong method or differentiation mode can determine whether training is highly efficient or fails due to Out-of-Memory (OOM) issues or excessive runtimes.

Throughout this section, we assume a model mapping inputs to outputs: $$ f: \mathbb{R}^n \rightarrow \mathbb{R}^m $$ where \(n\) is the input dimension (e.g., number of input channels times spatial grid size) and \(m\) is the output dimension.

neojax supports three primary methods for computing the derivatives, chosen via the method argument:

1. Exact Mode (method="exact")¤

This method computes the exact Jacobian of the model \(f(x)\) with respect to \(x\) using standard first-order AD. - Inner AD Mode (diff_mode): - Forward-mode (fwd): Computes derivatives column-by-column, scaling with the input dimension \(n\). - Reverse-mode (bwd): Computes derivatives row-by-row, scaling with the output dimension \(m\). - Auto (auto): neojax automatically uses forward-mode AD if \(n \le m\), and backward-mode AD if \(n > m\). - Compounding Outer-Pass Effect: During model training, the outer optimizer loop takes the gradient of the scalar loss with respect to the model parameters \(\theta\) (using a backward pass, e.g. jax.grad). - If the inner mode is Forward (fwd), JAX performs Reverse-over-Forward differentiation, which is highly memory-efficient because it does not store a nested backward tape. - If the inner mode is Reverse (bwd), JAX performs Reverse-over-Reverse differentiation, which requires storing nested Wengert tapes for both passes, drastically increasing the memory footprint. Tip: Even when \(n \approx m\), diff_mode="fwd" is often superior due to its memory efficiency.

2. Interpolated Mode (method="interpolate")¤

For higher-order derivatives (\(k \ge 2\)), computing the exact Jacobian recursively scales exponentially. Instead, neojax uses JAX's Taylor-mode AD (jax.experimental.jet), which propagates truncated Taylor polynomials. - The interpolated method evaluates the Taylor expansion of \(f(x)\) in all \(n\) canonical basis directions. - This computes the exact higher-order derivatives, but requires a jax.vmap loop over \(n\) directions. - It is highly efficient for moderate input dimensions (\(n\)), but becomes expensive as the spatial resolution or input channel count increases.

3. Stochastic Mode (method="stochastic")¤

When the input dimension \(n\) is very large (e.g., high-resolution grids), computing the Taylor polynomial in all \(n\) directions is intractable. - The stochastic method uses the Stochastic Taylor Derivative Estimator (STDE). - Instead of computing derivatives along every canonical basis, it samples a small number of random direction vectors (Rademacher or Normal distribution) and evaluates the Taylor expansion only in those random directions. - The computational cost is independent of the input dimension \(n\) and scales only with the number of random samples (n_random_samples), which defaults to 10.

Auto Routing (method="auto")¤

If you set method="auto", neojax automatically routes to the most suitable strategy: - \(k = 1\): Uses "exact". - \(k \ge 2\): Uses "interpolate" if all spatial grid dimensions of the input \(x\) are \(\le 10\), otherwise defaults to "stochastic" to prevent OOM and runtime issues.

neojax.losses.sobolev_losses.SobolevLoss ¤

General (Sobolev) \(W^{k,p}\)-loss.

Computes the loss in the \(W^{k,p}\) Sobolev space with the associated norm: $$ \begin{aligned} \Vert (y - \hat{y}) \Vert_{W^{k,p}(\Omega)} &= \begin{cases} \left(\sum_{\vert \alpha \vert \le k} \Vert D^\alpha (y - \hat{y}) \Vert_{L^p(\Omega)}^p\right)^{1/p}, & \text{if } p < \infty, \\ \max_{\vert \alpha \vert \le k} \Vert D^\alpha (y - \hat{y}) \Vert_{L^\infty(\Omega)}, & \text{if } p = \infty. \end{cases} \end{aligned} $$ where \(D^{\alpha}(y - \hat{y})\) are all combinations of partial derivatives up to the \(k\)-th order and \(y\) is the ground truth and \(\hat{y}\) is the model prediction.

Default is the \(W^{1,2}\)-loss (known as \(H^1\)-loss).

Parameters:

Name Type Description Default
k int

Order of partial derivatives. Default is 1.

 1
p float

Power of the \(L^{p}\)-norm. Default is 2.0.

 2.0
method Literal['exact', 'interpolate', 'stochastic', 'auto']

Method for computing higher order derivatives. See documentation for details. Default is auto.

 'auto'
diff_mode Literal['fwd', 'bwd', 'auto']

Inner AD mode. See documentation for details. Default is auto.

 'auto'
weight float

(Learnable) weight. Loss is computed as weight * loss. Default is 1.0.

 1.0
learnable_weight bool

Whether weight is learnable. Used to filter trainable parameters using is_learnable_loss_weight utility function with equinox.filter_... or equinox.partition. Default is False.

 False
n_random_samples int

Number of random vectors to use in "stochastic" method. Ignored in other methods. Default is 10.

 10
random_type Literal['rademacher', 'normal']

Type of distribution to sample from in "stochastic" method. Default is standard Normal distribution.

 'normal'

Internal Attributes

These fields store the internal state of the loss.

  • k (int): Order of partial derivatives.
  • p (float): Power of the \(L^{p}\)-norm.
  • method (Literal["exact", "interpolate", "stochastic", "auto"]): Method for computing higher order derivatives. See documentation for details.
  • diff_mode (Literal["fwd", "bwd", "auto"]): Inner AD mode. See documentation for details.
  • weight (Float[Array, ""]): Learnable loss weight. Filter during training to prevent updates.
  • learnable_weight (bool): Flag indicating whether weight is learnable.
  • n_random_samples (int): Number of random vectors to sample.
  • random_type (Literal["rademacher", "normal"]): Type of distribution to sample from in "stochastic" method. Default is standard Normal distribution.

Info

Computing higher order derivates is expensive. Given a function $f: \mathbb{R}^m \rightarrow \mathbb{R}^d $, computing the derivatives up to order k, the XLA graph scales with \(\mathcal{O}(m \cdot d^k)\).

exact_sobolev(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...']) -> Float[Array, ''] ¤

Computes the Sobolev loss exactly for a single sample.

Uses either eqx.filter_jacrev or eqx.filter_jacfwd to compute the exact derivative w.r.t. x, depending on diff_mode.

If diff_mode is auto, uses forward-mode AD if all output dimensions are larger or the same as all input dimensions, else backward-mode AD.

Parameters:

Name Type Description Default
model Module

The model being trained.

 required
x Float[Array, 'in_c ...']

Model input array shaped (in_c, d1, ..., dN).

 required
target Float[Array, 'c ...']

Ground truth array shaped (c, d1, ..., dN).

 required

Returns:

Type Description
Float[Array, '']

Scalar Sobolev loss.
interpolated_sobolev(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...']) -> Float[Array, ''] ¤

Compute Sobolev loss for a single sample by propagating a higher order Taylor polynomial.

Uses jax.experimental.jet, which forwards a higher order Taylor polynomial. This is more efficient than the exact method for higher order derivatives.

Parameters:

Name Type Description Default
model Module

The model being trained.

 required
x Float[Array, 'in_c ...']

Model input array shaped (in_c, d1, ..., dN).

 required
target Float[Array, 'c ...']

Ground truth array shaped (c, d1, ..., dN).

 required

Returns:

Type Description
Float[Array, '']

Scalar Sobolev loss.
stochastic_sobolev(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...'], key: PRNGKeyArray) -> Float[Array, ''] ¤

Computes the stochastic Sobolev loss.

Uses the STDE approach, making jet more efficient by sampling random vectors instead of using all basis tensors.

Cite

Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

@article{shi2024stochastic,
    title={Stochastic taylor derivative estimator:
    Efficient amortization for arbitrary differential operators},
    author={Shi, Zekun and Hu, Zheyuan and Lin, Min and Kawaguchi, Kenji},
    journal={Advances in Neural Information Processing Systems},
    volume={37},
    pages={122316--122353},
    year={2024}
}

Parameters:

Name Type Description Default
model Module

The model being trained.

 required
x Float[Array, 'in_c ...']

Model input array shaped (c, d1, ..., dN).

 required
target Float[Array, 'c ...']

Ground truth array shaped (c, d1, ..., dN).

 required
key PRNGKeyArray

Random key.

 required

Returns:

Type Description
Float[Array, '']

Scalar Sobolev loss.

Composing Losses¤

Real-world physics problems often require minimizing multiple objectives simultaneously (e.g., combining a data loss with a physics-informed regularization term). ComposedLoss allows you to pipe multiple loss functions together, computing the weighted sum of its constituent losses.

neojax.losses.composed_loss.ComposedLoss ¤

Wraps loss compositions.

Composes different losses by passing them to the classes init method.

Parameters:

Name Type Description Default
losses BaseLoss

Arbitrary number of loss instances to compose.

 required
weight float

Optional weighting for ComposedLoss. Default is 1.0, i.e., no weighting.

 1.0

Internal Attributes

These fields store the internal state of the loss.

  • weight (Float[Array, ""]): Learnable loss weight. Filter during training to prevent updates.
  • losses (tuple[BaseLoss, ...]): Composed loss functions. The composition is the sum.

Info

This loss may be called without a model by using keywords: loss_fn(pred=y_hat, target=y). The model parameter is kept for clean compatibility with JAX transformations and if composing model-dependent losses (e.g. Sobolev loss).

Examples:

import jax.numpy as jnp
from neojax.losses import ComposedLoss, LpLoss, RelativeLpLoss

loss_1 = LpLoss(p=2)
loss_2 = RelativeLpLoss(p=2)

composed_loss = ComposedLoss(loss_1, loss_2)

x = jnp.ones((1, 10))
y = jnp.zeros((1, 10))

loss_val = composed_loss(pred=x, target=y)

Learnable Loss Weights¤

When combining multiple losses, balancing their static weights can be notoriously difficult. neojax solves this by supporting dynamically learnable loss weights.

If you enable learnable weights in a ComposedLoss, custom BaseLoss, or any other loss instance, you must ensure JAX differentiates with respect to them. neojax provides the is_learnable_loss_weight utility to create the correct filter specification for equinox.

import equinox as eqx
from neojax.losses.utils import is_learnable_loss_weight

# 1. Create a filter spec that isolates the network parameters AND learnable loss weights
filter_spec = is_learnable_loss_weight(model_and_losses)

# 2. Pass the spec to filter_grad
grads = eqx.filter_grad(loss_fn, filter=filter_spec)(model_and_losses, data, targets)

neojax.losses.utils.is_learnable_loss_weight(model: PyTree) -> PyTree ¤

Creates a filter spec to correctly handle learnable weights.

Produces a boolean PyTree that can be used directly with eqx.partition, eqx.filter_grad, or eqx.filter_jit. It returns False for most fields, but True for: 1. All inexact arrays (trainable parameters) in the model. 2. Loss weight attrs only if their learnable_weight is True.

Parameters:

Name Type Description Default
model PyTree

The model or list of losses to generate a filter for.

 required

Returns:

Type Description
PyTree

A boolean PyTree of the same structure as model.
Example 
import equinox as eqx
from neojax.losses.utils import is_learnable_loss_weight

# Directly use as the filter for grad
filter_spec = is_learnable_loss_weight(model)
grads = eqx.filter_grad(loss_fn, filter=filter_spec)(model, ...)

Custom Losses¤

When implementing custom losses, inherit from BaseLoss.

Each custom loss should implement the __call__ method, taking the predicted array and the target array as inputs and returning a scalar loss.