Skip to content

Sobolev Metrics

Sobolev Metrics¤

While \(L^p\) metrics only measure the error in the predicted function values, Sobolev metrics (\(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 (higher order) derivatives in AD requires consideration of input and output dimensions. Choosing the wrong method or differentiation mode can determine whether training is 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", k==1)¤

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 Jacobian column-by-column, scaling with the input dimension \(n\).
    • Reverse-mode (bwd): Computes Jacobian row-by-row, scaling with the output dimension \(m\).
    • Auto (auto): neojax uses forward-mode AD if \(n \le m\), and backward-mode AD if \(n > m\).

Even when \(n \approx m\), diff_mode="fwd" is typically superior due to its memory efficiency.

2. Interpolated Mode (method="interpolate", k > 1)¤

For higher-order derivatives (\(k \ge 2\)), 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 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\) or output dimension \(m\) (in case of diff_mode="bwd") is very large (e.g., high-resolution grids), computing the exact or interpolated derivatives is intractable.

  • The stochastic method uses an unbiased estimator of the norm instead of the exact norm.
  • Instead of computing derivatives along every basis direction, it samples a small number of random basis vectors/tensors and evaluates the derivatives only in those random directions.

Auto Routing (method="auto")¤

If you set method="auto", neojax automatically routes to the most suitable strategy:

  • \(k = 1\): Uses "exact" if the input or output dimensions are small (configurable) or "stochastic" for large grids (random subsampling of basis vectors).
  • \(k \ge 2\): Uses "interpolate" if the input dimensions are small (configurable), otherwise "stochastic" (using STDE; random subsampling of basis tensors).

neojax.metrics.sobolev_metrics.SobolevMetric ¤

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

Computes the metric 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}\)-metric (known as \(H^1\)-metric).

Parameters:

Name Type Description Default
k int

Order of partial derivatives. Default is 1.

 1
p float | int

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

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

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

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

Inner AD mode for first order derivates. See documentation for details. Default is auto.

 'auto'
weight float

(Learnable) scalar weight coefficient of the metric. The metric is computed as 'metric * weight'. Default is 1.0. If learnable_weight is True, weight has to be strictly positive and will be optimized via gradient descent.

 1.0
learnable_weight bool

Optional flag indicating whether weight is learnable. If True, this flag is used in the is_learnable_metric_weight filter function to indicate to eqx.filter_... or eqx.partition that weight should be adapted by the optimizer. Default is False.

 False
n_random_samples int

Number of random vectors to sample (without replacement) in "stochastic" method. Ignored in other methods. Default is 10.

 10
direction_threshold int

Maximum number of basis tensors for derivative seeding until "stochastic" method is used. Default is 32. Only change if needed.

 1024

Raises:

Type Description
ValueError

If the order of derivatives is negative or zero.
ValueError

If the method or diff_mode is invalid.
ValueError

If the direction_threshold is negative or zero.
ValueError

If k > 1 and the method is "exact".
Internal Attributes

These fields store the internal state of the metric.

  • 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 metric 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.
  • direction_threshold (int): Maximum number of basis tensors for derivative seeding until "stochastic" method is used. Default is 1024.

Warning

Computing higher-order derivatives is expensive. Given a function \(f: \mathbb{R}^n \rightarrow \mathbb{R}^m\) the k-th derivative has \(m \cdot n^k\) components, so the size of the derivative representation can grow as \(\mathcal{O}(m\cdot n^k)\). The actual computational and memory cost also depends on the model and differentiation strategy. Only change method if you know what you are doing; otherwise, it's a fast way to run out of memory.

Unsupported primitive rules for jet

jax.experimental.jet doesn't yet implement rules for all jax primitives. If a model uses a primitive without corresponding jet rule (e.g. neojax.models.GeoFNO uses jax.lax.scatter), the methods for computing higher order partial derivatives with jax.experimental.jet (interpolated_sobolev and stochastic_sobolev_stde) are not available for that model. This condition raises a KeyError containing the name of the primitive if encountered. This is a limitation of jax.experimental.jet's current implementation.

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

Computes the first order Sobolev metric exactly.

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 the output size is larger or the same as the input size, 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 metric.
interpolated_sobolev(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...']) -> Float[Array, ''] ¤

Computes Sobolev metric for higher orders 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 metric.
stochastic_sobolev_first_order(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...'], key: PRNGKeyArray) -> Float[Array, ''] ¤

Computes the stochastic Sobolev metric for first order derivatives.

Uses a stochastic estimation approach inspired by sparse STDE in stochastic_sobolev_stde by seeding the JVP/VJP with a random sparse subset of basis vectors. This yields an exact subset of exact columns/rows of the Jacobian and converges to the exact Jacobian. It is an unbiased estimator of the Sobolev norm for k=1; useful for large input/output grids.

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 metric.
stochastic_sobolev_stde(model: Module, *, x: Float[Array, 'in_c ...'], target: Float[Array, 'c ...'], key: PRNGKeyArray) -> Float[Array, ''] ¤

Computes the stochastic Sobolev metric for higher order derivatives.

Uses jax.experimental.jet, which forwards a higher order Taylor polynomial which is more efficient than nested first order methods.

Uses the sparse STDE approach, making jet more efficient by sampling random basis tensors instead of using all basis tensors. This yields random subsamples of exact and pure partial derivatives. The idea is the same as in stochastic_sobolev_first_order: An unbiased estimator of the Sobolev norm for k>1 and large grids.

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 metric.
__call__(model: equinox._module._module.Module | None = None, *, target: Float[Array, 'b c ...'], x: Float[Array, 'b in_c ...'] | None = None, pred: Float[Array, 'b c ...'] | None = None, key: PRNGKeyArray | None = None, **kwargs: typing.Any) -> Float[Array, ''] ¤

Computes the general Sobolev metric.

Parameters:

Name Type Description Default
model Module | None

The model being trained.

 None
target Float[Array, 'b c ...']

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

 required
x Float[Array, 'b in_c ...'] | None

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

 None
pred Float[Array, 'b c ...'] | None

Unused in Sobolev metric.

 None
key PRNGKeyArray | None

Optional random key needed for stochastic Sobolev. Default is None.

 None
**kwargs Any

Additional arguments for specific metrics.

 required

Returns:

Type Description
Float[Array, '']

Scalar metric.

Raises:

Type Description
ValueError

If method is "stochastic" but no random key was passed. Or if model or x are not provided.
KeyError

If an outstanding jet primitive rule is encountered for higher order derivatives for some models.