sage.architecture.network.consistency

Multi-detector consistency heads.

Per-detector parameter heads that hang off each frontend (pre-merge), used to build an uncertainty-weighted coherence statistic between detectors. A real GW is coherent — its arrival times differ by at most the inter-detector light-travel time and its chirp mass is shared — so disagreement that is large relative to the predicted uncertainty is evidence against a real coincidence.

Components

• AttentionPool1d : learned temporal soft-attention pooling. Returns the pooled feature, the attention weights, the raw scores, and the attention entropy (a confidence signal: low entropy = peaked = confident).

• GlobalAvgMaxPool1d : concat of temporal mean and max.

• PerDetHead : per-detector tc (soft-argmax over physical time, attention-entropy-driven sigma) and mchirp (attention + avg/max pooled, MLP mean and sigma) heads, with their attention entropies.

All tensors follow the (B, C, T) convention of the Sage frontend output (channels then time). Everything here is torch.compile-safe (no data-dependent control flow; config flags are static Python booleans).

Classes

AttentionPool1d	Learned temporal soft-attention pooling over the time axis.
GlobalAvgMaxPool1d	Concatenate temporal mean and max pooling: (B, C, T) -> (B, 2C).
PerDetOutput	Per-detector head outputs (all shape (B,)).
PerDetHead	Per-detector tc and mchirp heads applied to one frontend output.

Functions

consistency_statistic(out_a, out_b, light_travel_time)	Uncertainty-weighted inter-detector consistency statistics for a pair.
corroboration_features(out_a, out_b, s_tc, s_mc)	Build the (B, 8) corroboration feature block for the ranking head.

Module Contents

class AttentionPool1d(in_ch, eps=1e-08)

Bases: torch.nn.Module

Learned temporal soft-attention pooling over the time axis.

Parameters:

• in_ch (int) – Number of input channels.

• eps (float) – Stabiliser inside the entropy logarithm.

• Forward

• -------

• x (torch.Tensor, shape (B, C, T))

Returns:

• pooled ((B, C) attention-weighted feature)

• attn ((B, T) softmax attention over time)

• scores ((B, T) raw (pre-softmax) scores)

• entropy ((B,) attention entropy -sum attn*log(attn+eps)) – (low = peaked/confident, high = diffuse/uncertain)

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

score

eps = 1e-08

forward(x)

Parameters:

x (torch.Tensor)

class GlobalAvgMaxPool1d(*args, **kwargs)

Bases: torch.nn.Module

Concatenate temporal mean and max pooling: (B, C, T) -> (B, 2C).

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

Parameters:

• args (Any)

• kwargs (Any)

forward(x)

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

class PerDetOutput

Bases: NamedTuple

Per-detector head outputs (all shape (B,)).

sigma_* are the actual standard deviations (> 0), produced by a softplus parameterisation, not log-sigmas — softplus grows linearly so it cannot blow up the 1/sigma^2 precision the way exp(-2 log_sigma) can.

mu_tc: torch.Tensor

sigma_tc: torch.Tensor

mu_mc: torch.Tensor

sigma_mc: torch.Tensor

entropy_tc: torch.Tensor

entropy_mc: torch.Tensor

embedding: torch.Tensor | None = None

class PerDetHead(in_ch, hidden=128, sigma_min=0.001, sigma_max=10.0, ensemble_tc=False, eps=1e-08, dropout=0.0)

Bases: torch.nn.Module

Per-detector tc and mchirp heads applied to one frontend output.

The same module is applied to each detector’s frontend output (weights may be shared across detectors; only the inputs differ). t_position is the physical (within-window) time in seconds of each of the T steps — the multirate/frontend downsampling means index != time, so the mapping must be supplied.

tc head (localisation-preserving): a saliency conv produces a temporal softmax whose soft-argmax over t_position gives mu_tc. A separate attention pool drives sigma_tc from the pooled feature plus the attention entropy (diffuse attention -> larger sigma).

mchirp head (time-tolerant): attention + global avg/max features feed an MLP for mu_mc; a second MLP (plus attention entropy) gives sigma_mc.

Both sigmas use a softplus parameterisation (sigma = softplus(raw) + sigma_min, clamped to [sigma_min, sigma_max]) rather than exp of a clamped log-sigma — softplus is linear in the tail, so it cannot drive the 1/sigma^2 precision to overflow.

Parameters:

• in_ch (int) – Channels of the frontend output.

• hidden (int) – Hidden width of the sigma/mean MLPs.

• sigma_min (float) – Lower / upper bound on the predicted standard deviations (both sides bounded so the precision 1/sigma^2 and the variance sigma^2 stay finite).

• sigma_max (float) – Lower / upper bound on the predicted standard deviations (both sides bounded so the precision 1/sigma^2 and the variance sigma^2 stay finite).

• ensemble_tc (bool) – If True, mu_tc is the mean of the soft-argmax and the attention-weighted time estimate. Default False (soft-argmax only).

• eps (float) – Entropy logarithm stabiliser.

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

• dropout (float)

eps = 1e-08

sigma_min

sigma_max

ensemble_tc = False

tc_saliency

tc_attn

tc_log_sigma

mc_attn

mc_pool

mc_mu

mc_log_sigma

forward(x, t_position, return_embedding=False)

Parameters:

• x (torch.Tensor)

• t_position (torch.Tensor)

• return_embedding (bool)

Return type:

PerDetOutput

consistency_statistic(out_a, out_b, light_travel_time, eps=1e-08)

Uncertainty-weighted inter-detector consistency statistics for a pair.

s_tc is the chi-square-like arrival-time disagreement in excess of the light-travel time, scaled by the combined per-detector variance — so it is tolerant of a faint (large-sigma) detector but strict when both are confident. s_mc is the analogous statistic for chirp mass (no light-travel allowance: a real signal shares mchirp exactly).

Parameters:

• out_a (PerDetOutput) – Per-detector head outputs for detectors A and B.

• out_b (PerDetOutput) – Per-detector head outputs for detectors A and B.

• light_travel_time (torch.Tensor) – Scalar light-travel time (seconds) between the two detectors.

• eps (float) – Denominator stabiliser.

Returns:

s_tc, s_mc

Return type:

torch.Tensor, each (B,) (always float32)

corroboration_features(out_a, out_b, s_tc, s_mc)

Build the (B, 8) corroboration feature block for the ranking head.

``[log1p(s_tc), log1p(s_mc), sigma_tc_A, sigma_tc_B, sigma_mc_A, sigma_mc_B,

entropy_tc_A, entropy_tc_B]`` — the disagreement statistics together with

the per-detector uncertainties and attention entropies, so the learned combiner can require small disagreement AND small sigma AND peaked (low entropy) attention on both detectors, rather than a hard gate on s.

The chi-square statistics are log1p-compressed: s is an unbounded squared z-score (it can reach ~1e7 for a confident disagreement), and feeding it raw would overflow the fp16 ranking head and let a single feature dominate the classification logits. log1p is monotonic (large s still reads as “inconsistent”) and bounded, and a log scale is friendlier to a linear head. Returned in float32; the caller casts to the ranking-head dtype.

Parameters:

• out_a (PerDetOutput)

• out_b (PerDetOutput)

• s_tc (torch.Tensor)

• s_mc (torch.Tensor)

Return type:

torch.Tensor