Scale the loss

Scale the loss#

The default training loss combines likelihood and priors. A minibatch likelihood is scaled up to represent its full training group; priors are not scaled up. Each group gets its own factor, so different batch sizes do not change the relative weight of the groups.

Validation and test likelihoods include only the split branches and are scaled to the corresponding training size. Unsplit observed likelihoods contribute to training only, including observations explicitly marked as passthrough. LieselOptim then divides losses by the total training sample size by default; use scale_loss=False to keep the sum. Sample size counts likelihood terms, which need not equal the number of array elements.

Validation leaves out priors by default. Use validation_strategy="log_prob" to include them. See NegLogProbLoss for details.

Open the likelihood-scaling overview in a separate page.

Set sample sizes#

Observed distributions with per_obs=False sum their likelihood before returning it, preventing automatic sample-size inference. Build a small Normal model with this setting and construct its split explicitly:

import logging

import jax.numpy as jnp
import optax
import tensorflow_probability.substrates.jax.distributions as tfd

import liesel.model as lsl
import liesel.optim as opt

logging.getLogger("liesel").setLevel(logging.WARNING)

loc = lsl.Var.new_param(0.0, name="loc")
y = lsl.Var.new_obs(
    jnp.array([1.0, 2.0, 3.0]), lsl.Dist(tfd.Normal, loc, 1.0), name="y"
)
y.dist_node.per_obs = False
model = lsl.Model(y)

split = opt.PositionSplit.from_model(
    model, infer_sample_sizes=False, multi_size="manager", shuffle=False
)
optim = opt.LieselOptim(
    model,
    split=split,
    optimizers=optax.adam(0.01),
    loss_monitor="train_full_data",
    show_progress=False,
)
split.train_sample_size
3.0
result = optim.fit()
result.position_min_monitor
{'loc': Array(1.9974802, dtype=float32, weak_type=True)}

This chooses split-axis counts for scaling. Alternatively, supply effective sample_sizes to the split factory. Setting scale_loss=False on LieselOptim only disables final loss normalization; it does not disable split inference or specify batch scaling.

Use a custom loss#

Custom aggregate likelihood, prior, or probability nodes require a custom Loss and an explicit split. A manual split specifies data grouping and scaling; it does not change the objective used by NegLogProbLoss. The built-in loss accepts the standard sums of observed distribution factors and parameter priors, including weak observed variables. Other distribution factors must be classified appropriately or handled by a custom loss.

For the data setup, see Split data and Create minibatches. Sample with weights explains the additional correction for unequal sampling probabilities.