Loss balancers

A balancer combines ordered, unweighted objective losses. It is independent of PhiModule, so the same module and objective can use static, gradient-norm, exact-NTK, or custom weighting.

Shared state and update contract

BalancerState.weights is aligned with objective loss_names. The traces field stores the latest adaptive diagnostic: parameter-gradient norms for GradNormBalancer, or mean pointwise diagonal NTKs for ExactNTKBalancer.

BalancerUpdatePlan contains a compiled functional update, its host-side cadence, and optional diagnostic batch sizes. A None batch-size mapping reuses the current training batches.

Adaptive balancers accept update_every_n_steps and an optional update_start_step. By default, the first update occurs after one complete interval. An explicit start anchors the cadence, so a start of 50 with an interval of 100 updates at steps 50, 150, and 250.

Logged metrics

A balancer's diagnostics() method returns names such as weight/pde/heat. PhiModule adds the train/ prefix before the values reach callbacks, progress bars, or experiment loggers. For a loss named pde/heat, the built-in balancers produce these metrics:

Balancer	Metric	Meaning
All built-ins	`train/weight/pde/heat`	Current multiplier applied to the unweighted `pde/heat` loss
`GradNormBalancer`	`train/grad_norm/pde/heat`	L2 norm of the unweighted loss gradient over all model parameters
`ExactNTKBalancer`	`train/ntk/pde/heat`	Mean pointwise diagonal empirical NTK for the configured residual or output stream

Every scalar is sent to configured experiment loggers at the Trainer's logging cadence. Loss weights also appear in the progress metrics by default. Gradient-norm and NTK diagnostics are logger-only unless a module explicitly routes them to the progress bar with self.log(..., prog_bar=True).

StaticLossBalancer reports the configured weights throughout fitting. Adaptive diagnostics start at zero and change only when their scheduled update runs. Their most recent values are reported between updates; they are not recomputed on every optimizer step.

BalancerState [source]

Store JIT-compatible loss weights and the most recent NTK diagnostics.

Attributes:

Name	Type	Description
`weights`	`Array`	Vector aligned with the balancer's fixed loss-name ordering.
`traces`	`Array`	Most recently computed NTK statistic for each loss, or zeros before the first update. Exact-NTK balancing stores the mean pointwise diagonal empirical NTK in this field.

BalancerUpdatePlan [source] `dataclass`

BalancerUpdatePlan(
    update: BalancerUpdate,
    every_n_steps: int,
    update_start_step: int,
    batch_sizes: Mapping[str, int] | None = None,
)

Describe one host-scheduled adaptive-balancer update.

Attributes:

Name	Type	Description
`update`	`BalancerUpdate`	Compiled functional update accepting model state, named batches, and balancer state.
`every_n_steps`	`int`	Positive optimizer-step interval between updates.
`update_start_step`	`int`	Nonnegative absolute optimizer step anchoring the update cadence.
`batch_sizes`	`Mapping[str, int] \| None`	Optional positive diagnostic batch size for each objective batch key. `None` reuses the current training batches instead of sampling a fixed diagnostic batch.

__post_init__ [source]

__post_init__() -> None

Validate host scheduling and fixed diagnostic-batch sizes.

Raises:

Type	Description
`TypeError`	If scheduling values or diagnostic batch names have invalid types.
`ValueError`	If an interval or diagnostic batch size is zero or negative.

LossBalancer [source]

Bases: Protocol

Define the functional contract for combining named scalar losses.

initialize [source]

initialize() -> Any

Create the JAX-compatible initial balancer state.

Returns:

Type	Description
`Any`	Initial balancer-state PyTree.

combine [source]

combine(
    losses: Mapping[str, Array], state: Any
) -> jax.Array

Combine named unweighted losses into one scalar objective.

Parameters:

Name	Type	Description	Default
`losses`	`Mapping[str, Array]`	Scalar losses keyed by every configured loss name.	required
`state`	`Any`	Current balancer-state PyTree.	required

Returns:

Type	Description
`Array`	Weighted scalar loss.

diagnostics [source]

diagnostics(state: Any) -> Mapping[str, jax.Array]

Expose scalar diagnostics without interpreting state externally.

Parameters:

Name	Type	Description	Default
`state`	`Any`	Current balancer-state PyTree.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar diagnostics with stable names.

AdaptiveBalancer [source]

Bases: Protocol

Define the assembly contract for a periodically updated loss balancer.

build_update_plan [source]

build_update_plan(
    module: Any, batch_keys: Sequence[str]
) -> BalancerUpdatePlan

Build a reusable update and declare its diagnostic data policy.

Parameters:

Name	Type	Description	Default
`module`	`Any`	Application module exposing the numerical streams required by the balancer.	required
`batch_keys`	`Sequence[str]`	Stable objective batch keys available from the training source.	required

Returns:

Type	Description
`BalancerUpdatePlan`	Functional update, scheduling, and diagnostic data plan consumed by training assembly.

Static weighting

StaticLossBalancer [source]

StaticLossBalancer(
    loss_names: Sequence[str],
    *,
    weights: Mapping[str, float] | None = None,
    default_weight: float = 1.0,
)

Apply immutable scalar weights to a fixed sequence of loss names.

Attributes:

Name	Type	Description
`initial_weights`		Configured `float32` weight vector.

Initialize stable loss ordering and configured weights.

Parameters:

Name	Type	Description	Default
`loss_names`	`Sequence[str]`	Unique non-empty sequence defining loss ordering.	required
`weights`	`Mapping[str, float] \| None`	Optional explicit weights keyed by loss name.	`None`
`default_weight`	`float`	Weight assigned to names absent from `weights`.	`1.0`

Raises:

Type	Description
`ValueError`	If `loss_names` is empty or contains duplicates.

initialize [source]

initialize() -> BalancerState

Create the fixed balancer state.

Returns:

Type	Description
`BalancerState`	State containing configured weights and zero diagnostic placeholders.

combine [source]

combine(
    losses: Mapping[str, Array], state: BalancerState
) -> jax.Array

Combine named scalar losses in stable order.

Parameters:

Name	Type	Description	Default
`losses`	`Mapping[str, Array]`	Scalar loss arrays keyed by every configured loss name.	required
`state`	`BalancerState`	Current balancer state containing the fixed weights.	required

Returns:

Type	Description
`Array`	Weighted scalar total.

Raises:

Type	Description
`KeyError`	If a configured loss name is absent from `losses`.

diagnostics [source]

diagnostics(
    state: BalancerState,
) -> Mapping[str, jax.Array]

Expose configured loss weights for logging.

Parameters:

Name	Type	Description	Default
`state`	`BalancerState`	Current fixed-weight state.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar weights keyed by loss name.

from phijax.balancers import StaticLossBalancer

balancer = StaticLossBalancer(
    module.loss_names,
    weights={"initial/u": 1.0, "pde/burgers": 1.0},
)

Pass the module's resolved loss_names; do not duplicate them in project settings.

Gradient-norm weighting

This policy balances inverse parameter-gradient magnitudes. It is not the training-rate-aware GradNorm multitask algorithm. Losses are differentiated individually to limit retained intermediate state. Each train/grad_norm/<loss_name> value is the L2 norm of one unweighted scalar loss gradient with respect to every model parameter. The update uses the current training batches.

GradNormBalancer [source]

GradNormBalancer(
    loss_names: Sequence[str],
    *,
    update_every_n_steps: int,
    update_start_step: int | None = None,
    eps: float = 1e-05,
    moving_average_coefficient: float = 0.9,
    initial_weights: Mapping[str, float] | None = None,
)

Balance named losses with parameter-gradient norm ratios.

The balancer computes the Euclidean parameter-gradient norm g_i of every unweighted scalar loss. Its raw weights are lambda_i = mean_j(g_j) / (g_i + eps * mean_j(g_j)), followed by an exponential moving average. Losses are differentiated one at a time so the update does not retain every loss Jacobian at once.

This is a direct gradient-norm weighting policy rather than the training-rate-aware GradNorm algorithm for multitask networks.

Attributes:

Name	Type	Description
`momentum`		Moving-average coefficient applied to previous weights.

Initialize gradient-norm balancing configuration.

Parameters:

Name	Type	Description	Default
`loss_names`	`Sequence[str]`	Unique non-empty sequence defining loss and diagnostic ordering.	required
`update_every_n_steps`	`int`	Positive optimizer-step interval between gradient-norm updates.	required
`update_start_step`	`int \| None`	Nonnegative absolute optimizer step of the first update. `None` starts after one update interval.	`None`
`eps`	`float`	Nonnegative multiplier regularizing each gradient-norm denominator.	`1e-05`
`moving_average_coefficient`	`float`	Previous-weight coefficient in `[0, 1)`.	`0.9`
`initial_weights`	`Mapping[str, float] \| None`	Optional initial weights keyed by loss name; unspecified names use `1.0`.	`None`

Raises:

Type	Description
`TypeError`	If update scheduling values have invalid types.
`ValueError`	If names are invalid, the update interval is not positive, `eps` is negative, or smoothing is outside `[0, 1)`.

initialize [source]

initialize() -> BalancerState

Create an initial balancer state.

Returns:

Type	Description
`BalancerState`	State with configured weights and zero gradient-norm diagnostics.

combine [source]

combine(
    losses: Mapping[str, Array], state: BalancerState
) -> jax.Array

Apply the most recently computed gradient-norm weights.

Parameters:

Name	Type	Description	Default
`losses`	`Mapping[str, Array]`	Scalar loss arrays keyed by every configured loss name.	required
`state`	`BalancerState`	Current weights and gradient-norm diagnostics.	required

Returns:

Type	Description
`Array`	Weighted scalar total.

Raises:

Type	Description
`KeyError`	If a configured loss name is absent from `losses`.

diagnostics [source]

diagnostics(
    state: BalancerState,
) -> Mapping[str, jax.Array]

Expose current loss weights and parameter-gradient norms.

Parameters:

Name	Type	Description	Default
`state`	`BalancerState`	Current gradient-norm balancer state.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar weights and gradient norms keyed by loss name.

update_from_grad_norms [source]

update_from_grad_norms(
    grad_norms: Array, state: BalancerState
) -> BalancerState

Update smoothed weights from ordered loss-gradient norms.

Parameters:

Name	Type	Description	Default
`grad_norms`	`Array`	Gradient-norm vector aligned with :attr:`loss_names`.	required
`state`	`BalancerState`	Previous weights and gradient-norm diagnostics.	required

Returns:

Type	Description
`BalancerState`	New stopped-gradient state. An all-zero norm vector leaves every weight unchanged.

update [source]

update(
    module: BasePhiModule,
    model_state: Any,
    batches: NamedBatches,
    state: BalancerState,
) -> BalancerState

Compute one loss-gradient norm at a time and update balancer state.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module producing the named unweighted scalar losses.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Current fixed-structure training batches.	required
`state`	`BalancerState`	Previous balancer state.	required

Returns:

Type	Description
`BalancerState`	Updated weights and gradient norms in :attr:`loss_names` order.

make_update [source]

make_update(
    module: BasePhiModule,
) -> Callable[
    [Any, NamedBatches, BalancerState], BalancerState
]

Create a reusable compiled gradient-norm update.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module producing named unweighted scalar losses.	required

Returns:

Type	Description
`Callable[[Any, NamedBatches, BalancerState], BalancerState]`	JIT-compiled function accepting model state, current batches, and balancer state.

build_update_plan [source]

build_update_plan(
    module: BasePhiModule, batch_keys: Sequence[str]
) -> BalancerUpdatePlan

Describe a gradient-norm refresh using the current training batches.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module producing named unweighted scalar losses.	required
`batch_keys`	`Sequence[str]`	Available objective batch keys, retained for the shared adaptive-balancer contract.	required

Returns:

Type	Description
`BalancerUpdatePlan`	Update plan that reuses the current training batches.

balancer = GradNormBalancer(
    module.loss_names,
    update_every_n_steps=100,
)

Exact NTK weighting

Exact NTK balancing computes the squared parameter-Jacobian norm at each point. It averages these values for each residual stream and normalizes the loss weights from those means. Named streams are evaluated one at a time. kernel_chunk_size=1 minimizes peak memory; None vectorizes over all samples.

Each train/ntk/<loss_name> value is the resulting mean for one named residual or output stream. The Trainer samples fixed diagnostic batches of kernel_size rows once for the fit, and each scheduled update evaluates the current model on those batches.

The method is motivated by the empirical NTK analysis and adaptive PINN loss weighting in Wang, Yu, and Perdikaris (2022). PhiJAX uses the mean-diagonal rule lambda_i = mean_j(mu_j) / mu_i. It supports moving-average smoothing and processes one named stream at a time to limit peak memory.

exact_ntk_trace [source]

exact_ntk_trace(
    stream_fn: Callable[[Any], Array], model_state: Any
) -> jax.Array

Compute an exact diagonal-NTK contribution without materializing a Jacobian.

The function differentiates one scalar stream entry at a time, flattens the parameter-gradient PyTree, and sums its squared norm. A multi-feature stream therefore contributes the sum of its diagonal empirical NTK entries. :class:ExactNTKBalancer maps this calculation over samples in bounded chunks before averaging the resulting pointwise diagonal values.

Parameters:

Name	Type	Description	Default
`stream_fn`	`Callable[[Any], Array]`	Pure function mapping explicit model state to one residual array.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required

Returns:

Type	Description
`Array`	Scalar `float32` sum of diagonal empirical-NTK entries across the flattened stream.

ExactNTKBalancer [source]

ExactNTKBalancer(
    loss_names: Sequence[str],
    *,
    update_every_n_steps: int,
    kernel_size: int,
    kernel_chunk_size: int | None = 1,
    update_start_step: int | None = None,
    eps: float = 1e-10,
    moving_average_coefficient: float = 0.9,
    initial_weights: Mapping[str, float] | None = None,
)

Balance named losses with exact diagonal-NTK ratios and optional smoothing.

For loss i, the balancer computes the exact diagonal empirical NTK independently at every sampled point and stores its pointwise mean as mu_i. The raw rule is lambda_i = mean_j(mu_j) / mu_i. A moving average blends each raw update with its previous value. Using empirical NTK statistics to diagnose and balance PINN loss terms follows Wang et al. (2022); PhiJAX uses mean-diagonal normalization and evaluates named residual streams sequentially to control peak memory.

Attributes:

Name	Type	Description
`momentum`		Moving-average coefficient applied to previous weights.

References

Wang, S., Yu, X., and Perdikaris, P. (2022). When and Why PINNs Fail to Train: A Neural Tangent Kernel Perspective. Journal of Computational Physics, 449, 110768.

Initialize exact-NTK balancing configuration.

Parameters:

Name	Type	Description	Default
`loss_names`	`Sequence[str]`	Unique non-empty sequence defining loss and trace ordering.	required
`update_every_n_steps`	`int`	Positive optimizer-step interval between NTK updates.	required
`kernel_size`	`int`	Positive diagnostic sample count requested for every objective batch.	required
`kernel_chunk_size`	`int \| None`	Positive number of diagnostic samples differentiated in parallel, or `None` for a full :func:`jax.vmap`.	`1`
`update_start_step`	`int \| None`	Nonnegative absolute optimizer step of the first update. `None` starts after one update interval.	`None`
`eps`	`float`	Positive mean-diagonal floor used to avoid division by zero.	`1e-10`
`moving_average_coefficient`	`float`	Previous-weight coefficient in `[0, 1)`.	`0.9`
`initial_weights`	`Mapping[str, float] \| None`	Optional initial weights keyed by loss name; unspecified names use `1.0`.	`None`

Raises:

Type	Description
`TypeError`	If update scheduling or kernel-size values have invalid types.
`ValueError`	If names are invalid, sizes are not positive, `eps` is not positive, or smoothing is outside `[0, 1)`.

initialize [source]

initialize() -> BalancerState

Create an initial balancer state.

Returns:

Type	Description
`BalancerState`	State with configured initial weights and zero mean-diagonal placeholders.

combine [source]

combine(
    losses: Mapping[str, Array], state: BalancerState
) -> jax.Array

Apply the most recently computed NTK weights.

Parameters:

Name	Type	Description	Default
`losses`	`Mapping[str, Array]`	Scalar loss arrays keyed by every configured loss name.	required
`state`	`BalancerState`	Current NTK weights and mean-diagonal diagnostics.	required

Returns:

Type	Description
`Array`	Weighted scalar total.

Raises:

Type	Description
`KeyError`	If a configured loss name is absent from `losses`.

diagnostics [source]

diagnostics(
    state: BalancerState,
) -> Mapping[str, jax.Array]

Expose current NTK weights and mean-diagonal estimates.

Parameters:

Name	Type	Description	Default
`state`	`BalancerState`	Current NTK balancer state.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar weights and traces keyed by loss name.

update_from_traces [source]

update_from_traces(
    traces: Array, state: BalancerState
) -> BalancerState

Update smoothed weights from ordered mean diagonal-NTK values.

Parameters:

Name	Type	Description	Default
`traces`	`Array`	Mean diagonal-NTK vector aligned with :attr:`loss_names`.	required
`state`	`BalancerState`	Previous weights and mean-diagonal diagnostics.	required

Returns:

Type	Description
`BalancerState`	New stopped-gradient state containing nonnegative mean diagonals and smoothed weights. Zero entries
`BalancerState`	preserve their previous weights; an all-zero vector leaves every weight unchanged.

update [source]

update(
    module: BasePhiModule,
    model_state: Any,
    ntk_batches: NamedBatches,
    state: BalancerState,
    *,
    kernel_chunk_size: int | None = 1,
) -> BalancerState

Compute pointwise diagonal NTKs one named stream at a time and update state.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module exposing `residual_stream(name, model_state, batches)` independently of this balancer.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`ntk_batches`	`NamedBatches`	Fixed-shape named batches with a shared leading sample size, used only for NTK estimation.	required
`state`	`BalancerState`	Previous balancer state.	required
`kernel_chunk_size`	`int \| None`	Positive number of kernel samples differentiated in parallel. `1` evaluates samples sequentially for minimum peak memory; `None` uses a full :func:`jax.vmap`.	`1`

Returns:

Type	Description
`BalancerState`	Updated weights and mean diagonal-NTK values in :attr:`loss_names` order.

make_update [source]

make_update(
    module: BasePhiModule,
    *,
    kernel_chunk_size: int | None = 1,
) -> Callable[
    [Any, NamedBatches, BalancerState], BalancerState
]

Create the reusable compiled NTK update for fixed batch structures.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module exposing one-at-a-time named residual streams.	required
`kernel_chunk_size`	`int \| None`	Positive number of kernel samples differentiated in parallel. `1` minimizes peak memory; `None` fully vectorizes across samples.	`1`

Returns:

Type	Description
`Callable[[Any, NamedBatches, BalancerState], BalancerState]`	JIT-compiled function accepting model state, NTK batches, and balancer state. New array values reuse the
`Callable[[Any, NamedBatches, BalancerState], BalancerState]`	executable when shapes, dtypes, and PyTree structures stay fixed.

Raises:

Type	Description
`ValueError`	If `kernel_chunk_size` is not positive or `None`.

build_update_plan [source]

build_update_plan(
    module: BasePhiModule, batch_keys: Sequence[str]
) -> BalancerUpdatePlan

Describe an exact-NTK refresh with fixed diagnostic samples.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module exposing one-at-a-time named residual streams.	required
`batch_keys`	`Sequence[str]`	Stable objective batch keys sampled for every NTK refresh.	required

Returns:

Type	Description
`BalancerUpdatePlan`	Update plan containing scheduling, the compiled NTK update, and fixed diagnostic batch sizes.

balancer = ExactNTKBalancer(
    module.loss_names,
    update_every_n_steps=100,
    kernel_size=256,
    kernel_chunk_size=1,
)

See Creating a custom loss balancer for implementation and testing requirements.

Loss balancers

Shared state and update contract

Logged metrics

BalancerState [source]

BalancerUpdatePlan [source] dataclass

__post_init__ [source]

LossBalancer [source]

initialize [source]

combine [source]

diagnostics [source]

AdaptiveBalancer [source]

build_update_plan [source]

Static weighting

StaticLossBalancer [source]

initialize [source]

combine [source]

diagnostics [source]

Gradient-norm weighting

GradNormBalancer [source]

initialize [source]

combine [source]

diagnostics [source]

update_from_grad_norms [source]

update [source]

make_update [source]

build_update_plan [source]

Exact NTK weighting

exact_ntk_trace [source]

ExactNTKBalancer [source]

initialize [source]

combine [source]

diagnostics [source]

update_from_traces [source]

update [source]

make_update [source]

build_update_plan [source]

BalancerUpdatePlan [source] `dataclass`