Reference / API reference / dorothy / dorothy.losses / dorothy.losses.heteroscedastic

dorothy.losses.heteroscedastic¶

Heteroscedastic loss function for stellar parameter inference.

The loss is the Gaussian negative log-likelihood of each label under the prediction, up to a constant and a factor of 2:

s = sqrt(exp(2 * ln_s) + s_0^2)
loss = (mu - y)^2 / s^2 + log(s^2)

Where:

mu: predicted stellar parameter value
ln_s: predicted log-scatter. The model computes it from the reported label error or a learned missing vector (see LabelAwareOutput), so s is the whole expected scatter between mu and the label, label noise included.
y: label value
s_0: scatter floor to prevent numerical instability

Functions

predicted_scatter(ln_s, scatter_floor)

Convert predicted log-scatter to scatter: sqrt(exp(2 ln_s) + s_0^2).

Classes

`HeteroscedasticLoss`([scatter_floor, ...])	Heteroscedastic loss for regression with predicted uncertainties.
`MaskedMSELoss`([n_parameters, reduction])	Masked mean squared error on predicted means, for `loss: mse`.

dorothy.losses.heteroscedastic.predicted_scatter(ln_s, scatter_floor)[source]¶

Convert predicted log-scatter to scatter: sqrt(exp(2 ln_s) + s_0^2).

This is the single definition used by the loss, training validation, evaluation and inference, so predicted uncertainties agree everywhere. ln_s is clamped to [LN_S_MIN, LN_S_MAX] first.

Parameters:

ln_s (Tensor) – Predicted log-scatter (any shape).
scatter_floor (float) – Scatter floor s_0.

Returns:

Predicted scatter, same shape as ln_s.

Return type:

Tensor

class dorothy.losses.heteroscedastic.HeteroscedasticLoss(scatter_floor=0.01, n_parameters=11, reduction='mean')[source]¶

Bases: Module

Heteroscedastic loss for regression with predicted uncertainties.

loss = (mu - y)^2 / s^2 + log(s^2), where s is the predicted scatter (predicted_scatter()). The squared error is divided by the predicted variance, and the log-variance term penalizes underconfidence.

The model output has shape (batch, 2, n_parameters): - output[:, 0, :]: predicted means (mu) - output[:, 1, :]: predicted log-scatter (ln_s)

The target has 3-channel format (batch, 3, n_parameters): - target[:, 0, :]: true labels (y) - target[:, 1, :]: label errors, not used here (the model reads them) - target[:, 2, :]: mask, valid where > 0, used for loss computation

scatter_floor¶: Minimum scatter value (s_0) to prevent division by zero.

n_parameters¶: Number of stellar parameters being predicted.

reduction¶: How to reduce the loss (‘mean’, ‘sum’, or ‘none’).

Examples

>>> loss_fn = HeteroscedasticLoss(scatter_floor=0.01, n_parameters=11)
>>> output = model(x)  # Shape: (batch, 2, n_parameters)
>>> target = labels  # (batch, 3, n_parameters) = [y, label errors, mask]
>>> loss = loss_fn(output, target)

element_losses(output, target, beta=0.0)[source]¶

Unreduced loss per (sample, parameter), zero where masked.

loss = sg(s^2)^beta * ((mu - y)^2 / s^2 + log(s^2)), the β-NLL of Seitzer et al. (2022), where sg passes its value but no gradient.

Masked entries are zeroed with torch.where (not multiplied by the mask), so non-finite values in masked entries cannot leak NaN into the loss.

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
target (Tensor) – Targets of shape (batch_size, 3, n_parameters) holding [labels, label errors, mask].
beta (float) – Exponent of the variance weight, in [0, 1].

Returns:

Tensor of shape (batch_size, n_parameters).

Raises:

ValueError – If output or target shapes don’t match expectations, or beta is outside [0, 1].

Return type:

Tensor

forward(output, target)[source]¶

Compute the heteroscedastic loss.

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters). output[:, 0, :] contains predicted means (mu). output[:, 1, :] contains predicted log-scatter (ln_s).
target (Tensor) – Target values of shape (batch_size, 3, n_parameters). target[:, 0, :] contains true labels (y). target[:, 1, :] contains label errors (not used). target[:, 2, :] contains mask (valid where > 0).

Returns:

The computed loss, reduced according to self.reduction.

Raises:

ValueError – If output or target shapes don’t match expectations.

Return type:

Tensor

get_predicted_scatter(output)[source]¶

Extract the predicted scatter (standard deviation) from model output.

This is useful for post-processing predictions to get uncertainty estimates.

Parameters:: output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
Returns:: Predicted scatter of shape (batch_size, n_parameters), computed as sqrt(exp(2 * ln_s) + s_0^2).
Return type:: Tensor

forward_detailed(output, target)[source]¶

Compute loss with detailed per-parameter and per-component breakdown.

Returns both the total loss and a breakdown into: - Mean component per parameter: (mu - y)² / s² - Scatter component per parameter: log(s²)

Both components are averaged over the valid entries of each parameter.

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
target (Tensor) – Target values of shape (batch_size, 3, n_parameters). target[:, 0, :] contains true labels (y). target[:, 1, :] contains label errors (not used). target[:, 2, :] contains mask (valid where > 0).

Returns:

Dictionary containing – - ‘loss’: Total scalar loss (reduced) - ‘mean_component’: Per-parameter mean loss, shape (n_parameters,) - ‘scatter_component’: Per-parameter scatter loss, shape (n_parameters,) - ‘residuals’: mu - y, shape (batch, n_parameters), 0 where masked - ‘y_pred’: Predicted means, shape (batch, n_parameters) - ‘pred_scatter’: Predicted scatter, shape (batch, n_parameters)

Raises:

ValueError – If output or target shapes don’t match expectations.

Return type:

dict[str, Tensor]

extra_repr()[source]¶

Return a string representation of the loss configuration.

class dorothy.losses.heteroscedastic.MaskedMSELoss(n_parameters=11, reduction='mean')[source]¶

Bases: Module

Masked mean squared error on predicted means, for loss: mse.

Uses the same (batch, 2, n_params) output / (batch, 3, n_params) target format as HeteroscedasticLoss and ignores label uncertainties and the predicted log-scatter channel (which receives no gradient).

element_losses(output, target, beta=0.0)[source]¶

Squared error per (sample, parameter), zero where masked.

beta has no effect: the squared error is the negative log-likelihood with unit variance, whose β-NLL weight is 1.

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
target (Tensor) – Targets of shape (batch_size, 3, n_parameters) holding [labels, label errors, mask].
beta (float) – Ignored.

Returns:

Tensor of shape (batch_size, n_parameters).

Raises:

ValueError – If output or target shapes don’t match expectations.

Return type:

Tensor

forward(output, target)[source]¶

Compute the masked mean squared error.

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
target (Tensor) – Targets of shape (batch_size, 3, n_parameters) holding [labels, label errors, mask].

Returns:

The loss, reduced according to self.reduction.

Raises:

ValueError – If output or target shapes don’t match expectations.

Return type:

Tensor

forward_detailed(output, target)[source]¶

Per-parameter MSE as the mean component (scatter component is 0).

Parameters:

output (Tensor) – Model predictions of shape (batch_size, 2, n_parameters).
target (Tensor) – Targets of shape (batch_size, 3, n_parameters).

Returns:

Dictionary with the keys ‘loss’ (reduced), ‘mean_component’ and ‘scatter_component’ (both shape (n_parameters,)), ‘residuals’ (mu - y, shape (batch, n_parameters), 0 where masked) and ‘y_pred’ (predicted means). Unlike HeteroscedasticLoss it has no ‘pred_scatter’.

Raises:

ValueError – If output or target shapes don’t match expectations.

Return type:

dict[str, Tensor]