sage.architecture.custom_losses.consistency_loss

Per-detector consistency loss.

A standalone, masked heteroscedastic Gaussian negative-log-likelihood for the multi-detector consistency heads. Completely separate from BCEWithPEsigmaLoss (the existing classification + merged-PE loss is untouched); the consistency training loop adds this term on top.

The per-detector tc and mchirp heads are supervised with

nll(mu, sigma, y) = 0.5 * (mu - y)^2 / sigma^2 + log(sigma)

masked per detector and averaged over the supervised entries. Detectors that are not supervised in a given sample (mask = 0) contribute nothing, so their sigma is free to grow — exactly the desired behaviour for noise detectors and for the faint member of a faint coherent coincidence.

Masking convention (per-detector mask supplied by the caller, shape (B, D)):

• matched coherent : 1 for both detectors

• mismatch signal+noise : 1 for the signal detector only

• mismatch different-signal : 1 for both (each toward its own truth)

• pure noise : 0 for both

The same mask gates both the tc and mchirp terms.

Classes

ConsistencyNLLLoss

Masked per-detector Gaussian NLL for tc and mchirp.

Module Contents

class ConsistencyNLLLoss(tc_weight=1.0, mc_weight=1.0, beta=0.5, eps=1e-08)

Bases: torch.nn.Module

Masked per-detector Gaussian NLL for tc and mchirp.

Parameters:

• tc_weight (float) – Relative weights of the tc and mchirp NLL terms.

• mc_weight (float) – Relative weights of the tc and mchirp NLL terms.

• eps (float) – Stabiliser for the mask-count denominator.

• ``(B ()`` /) –

• broadcast (D)``; targets)

• ``(B –

• ``(B –

• 1)``)

• -------

• mu_tc (predicted means / standard deviations)

• sigma_tc (predicted means / standard deviations)

• mu_mc (predicted means / standard deviations)

• sigma_mc (predicted means / standard deviations)

• tc_target (per-detector arrival times (window-normalised))

• mc_target (per-detector chirp mass (standardised))

• mask ((B, D) per-detector supervision mask)

• beta (float)

Returns:

• torch.Tensor, shape ``(3,)`` ([total, tc_term, mchirp_term] (index 0 is)

• the value to backpropagate, matching the convention of the other losses).

• Initialize internal Module state, shared by both nn.Module and ScriptModule.

tc_weight

mc_weight

beta

eps

num_components = 3

forward(mu_tc, sigma_tc, mu_mc, sigma_mc, tc_target, mc_target, mask)