Composing Metrics

Composing Metrics

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

neojax.metrics.composed_metric.ComposedMetric

Wraps metric compositions.

Composes different metrics by passing them to the class init method and evaluating them as a weighted sum.

Parameters:

Name	Type	Description	Default
metrics	BaseMetric	Arbitrary number of metric instances to compose.	required
composition_fn	Callable[[BaseMetric, ...], Float[Array, '']]	Callable combining the given (unweighted) metrics into a scalar metric. Called as composition_fn(*metrics). May be a plain function, or a stateful eqx.Module with its own learnable parameters (see the note below). It is required.	required
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

Warns:

Type	Description
UserWarning	If any of the composed metrics' learnable_weight is True. Individual learnable weights should be controlled with the given composition_fn.

Examples:

import jax.numpy as jnp
from neojax.metrics import ComposedMetric, LpMetric, RelativeLpMetric

metric_1 = LpMetric(p=2)
metric_2 = RelativeLpMetric(p=2)

composed_metric = ComposedMetric(metric_1, metric_2, composition_fn=lambda ms: jnp.sum(*ms))

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

metric_val = composed_metric(pred=x, target=y)

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable scalar metric weight. Filter during training to prevent updates.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• metrics (tuple[BaseMetric, ...]): Composed metric functions. The composition is the sum.
• composition_fn (Callable[[BaseMetric, ...], Float[Array, ""]]): Callable combining the given metrics into a scalar metric.

Info

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

A stateful, learnable composition_fn

composition_fn doesn't have to be a plain function, it can be an eqx.Module with its own learnable parameters, e.g. Kendall & Gal uncertainty weighting's per-term log-variances:

class UncertaintyWeighting(eqx.Module):
    log_vars: Float[Array, " n"]

    def __call__(self, *raw_vals):
        return sum(
            jnp.exp(-s) * v + s for v, s in zip(raw_vals, self.log_vars)
        )

composed = ComposedMetric(
    metric_1, metric_2, composition_fn=UncertaintyWeighting(jnp.zeros(2))
)

is_learnable_metric_weight (see its own docs) marks any inexact array found inside composition_fn as trainable unconditionally, so log_vars above is included in the trainable partition regardless of ComposedMetric.learnable_weight.

Why a composition_fn must be passed

A naive composition function such as sum(*metrics) gives the optimizer an easy way to cancel its components: drive the weights of its metrics to zero. To prevent such a 'free lunch', make sure to pass a composition_fn that has no such flaw, such as Kendall & Gal's uncertainty weighting (see above).

Using metric weighting with ComposedMetric

Any weight, whether learnable or fixed of the composed metrics is ignored inside ComposedMetric. A UserWarning is raised if any of the weights is not 1.0, indicating it was set by the user. To weight individual metrics, use the composition_fn instead.

__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, return_components: bool = False, **kwargs: typing.Any) -> Float[Array, ''] | tuple[Float[Array, ''], tuple[Float[Array, ''], ...]]

Computes the composed metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained.	None
target	Float[Array, 'b c ...']	Ground truth array.	required
x	Float[Array, 'b in_c ...'] | None	Model input array.	None
pred	Float[Array, 'b c ...'] | None	Model prediction array.	None
return_components	bool	If True, also returns each metric's own raw (unweighted) value alongside the composed value.	False

Returns:

Type	Description
Float[Array, ''] | tuple[Float[Array, ''], tuple[Float[Array, ''], ...]]	Scalar metric value.