Supervision Losses

NeuralPOM implements multiple supervision strategies in supervision.py for comparing predicted ocean state trajectories against reference (ground truth) trajectories. All losses support masking (land/sea masks) to exclude non-ocean grid points, and all losses are computed per prognostic field before being combined into a total loss via weighted summation.

Supported Supervision Modes

Mode	Identifier	Description
Pointwise Baseline	pointwise_baseline	Standard MSE in full grid space (no mask, no filtering)
Multiscale Lowpass	multiscale_lowpass	Gaussian-filtered multi-scale MSE with horizon decay and land masks

The mode is selected via the --supervision_mode argument. Internally, build_supervision_loss() dispatches to the appropriate loss function:

def build_supervision_loss(preds, refs, *, supervision_mode, masks, sigmas,
                           scale_weights, horizon_decay, eps):
    if supervision_mode == "pointwise_baseline":
        return compute_pointwise_losses(preds, refs)
    if supervision_mode == "multiscale_lowpass":
        return compute_multiscale_lowpass_loss(
            preds, refs, masks, sigmas=sigmas, scale_weights=scale_weights,
            horizon_decay=horizon_decay, eps=eps)
    raise ValueError(f"Unknown supervision_mode '{supervision_mode}'")

Pointwise Baseline Loss

The simplest supervision mode: standard mean squared error over all grid points for each prognostic field.

Formulation

For each field $f$, the pointwise MSE is:

$${\mathcal{L}}_f^{\text{point}}=\frac{1}{N}\sum _{i,j,k}{({\mathbf{pred}}_f[i,j,k]-{\mathbf{ref}}_f[i,j,k])}^2$$

where $N$ is the total number of grid points (including both land and ocean), averaged over the batch and time dimensions.

Implementation

def compute_pointwise_losses(preds, refs):
    return {
        field: F.mse_loss(preds[field], refs[ref_field])
        for field, ref_field in PRED_TO_REF_KEY.items()
    }

This is the simplest and fastest loss, useful for initial debugging and rapid prototyping. It does not apply land masks, so land grid points contribute equally to the loss. It also does not apply any temporal weighting or spatial filtering.

Pred-to-Ref Field Mapping

The prediction field names differ slightly from the reference field names for some fields. The mapping is:

PRED_TO_REF_KEY = {
    "u": "u", "ub": "ub", "v": "v", "vb": "vb",
    "t": "t", "s": "s",
    "el": "eta",     # el → eta in reference
    "ua": "ua", "uab": "uab", "va": "va", "vab": "vab",
}

Note that el (sea surface height in predictions) maps to eta in the reference data.

Multiscale Lowpass Loss

The primary and most sophisticated supervision mode. It combines three mechanisms:

1. Land/sea masking: Exclude land grid points from the loss
2. Gaussian multi-scale filtering: Compute MSE at multiple spatial scales
3. Horizon decay: Weight earlier time steps more heavily than later ones

Formulation

The full multiscale lowpass loss for field $f$ is:

$${\mathcal{L}}_f^{\text{ms}}=\sum _{t=0}^{T-1}{w}_t^{\text{horizon}}\sum _{s=1}^S{\alpha }_s\cdot \text{MSE}(G_{{\sigma }_s}\star (\mathbf{m}\odot {\mathbf{pred}}_{f,t}),G_{{\sigma }_s}\star (\mathbf{m}\odot {\mathbf{ref}}_{f,t}))$$

where:

• $T$ is the number of time steps
• $S$ is the number of scale levels
• $\mathbf{m}$ is the land/sea mask (1 for ocean, 0 for land)
• $G_{{\sigma }_s}$ is a 2D Gaussian kernel with standard deviation ${\sigma }_s$
• $\star$ denotes 2D convolution
• $\odot$ denotes element-wise multiplication
• ${\alpha }_s$ is the scale weight (normalized so $\sum _s{\alpha }_s=1$)
• $w_t^{\text{horizon}}$ is the horizon weight (normalized so $\sum _t{w}_t^{\text{horizon}}=1$)

1. Land/Sea Masking

Ocean models have land grid points where the physics is ill-defined. Including these points in the loss would penalize the model for mismatches in regions that are physically meaningless. NeuralPOM uses static loss masks derived from the physics core's geometry.

Three mask types are computed:

Mask	Source Attribute	Applied To
dum	core.dum	Zonal velocity fields (u, ub, ua, uab)
dvm	core.dvm	Meridional velocity fields (v, vb, va, vab)
fsm	core.fsm	Scalar fields (t, s, el)

The masks are built once at initialization:

def build_static_loss_masks(core):
    base_masks = {
        "fsm": core.fsm.detach().float(),  # free-surface mask
        "dum": core.dum.detach().float(),  # u-velocity mask
        "dvm": core.dvm.detach().float(),  # v-velocity mask
    }
    return {field: base_masks[attr] for field, attr in MASK_ATTR_BY_FIELD.items()}

The Arakawa C-grid used by POM places velocity components at staggered grid points, which is why separate masks are needed for u-velocity and v-velocity fields.

2. Gaussian Multi-Scale Filtering

The Gaussian kernel is defined as:

$$G_{\sigma }[u,v]=\frac{1}{Z}\exp (-\frac{u^2+v^2}{2{\sigma }^2})$$

where $Z=\sum _{u,v}{G}_{\sigma }[u,v]$ normalizes the kernel to unit sum, and the kernel radius is $r=\lceil 3\sigma \rceil$.

The masked Gaussian lowpass operation applies the kernel only over ocean points, normalizing by the convolved mask to prevent land points from contaminating the filtered values:

$$\text{Lowpass}(\mathbf{x},\mathbf{m},\sigma )=\frac{G_{\sigma }\star (\mathbf{m}\odot \mathbf{x})}{G_{\sigma }\star \mathbf{m}+\epsilon }\odot \mathbf{m}$$

where $\epsilon$ (multiscale_eps) is a small constant to avoid division by zero.

Why multi-scale? Different physical processes in the ocean operate at different spatial scales. A single-scale loss might over-emphasize small-scale noise or under-emphasize large-scale circulation patterns. By computing losses at multiple Gaussian filter scales, the model is encouraged to match both the large-scale circulation and the finer mesoscale features.

Typical configuration:

multiscale_sigmas = [0.0, 1.0, 2.0, 4.0]     # σ = 0 means no filtering
multiscale_weights = [0.1, 0.2, 0.3, 0.4]    # more weight on larger scales

When $\sigma =0$, the "Gaussian filter" is the identity — this corresponds to pointwise MSE over ocean points only.

3. Horizon Decay

The horizon decay weighting biases training toward short-term accuracy:

$$w_t^{\text{horizon}}=\frac{{\gamma }^t}{\sum _{\tau =0}^{T-1}{\gamma }^{\tau }}$$

where $\gamma \in (0,1]$ is the horizon_decay parameter. When $\gamma <1$, earlier time steps contribute more to the loss:

• $\gamma =1.0$: All time steps weighted equally
• $\gamma =0.8$: Each successive step is weighted 80% of the previous
• $\gamma =0.5$: Each successive step is weighted 50% of the previous

Rationale: In chaotic systems like the ocean, long-term predictions are inherently uncertain. Without horizon decay, the model might sacrifice short-term accuracy in an attempt to improve long-term predictions — but due to chaos, those long-term improvements may be impossible. Horizon decay ensures the model prioritizes getting the immediate future right, while still providing a learning signal for longer horizons.

Full Multiscale Lowpass Implementation

def compute_multiscale_lowpass_loss(preds, refs, masks, *, sigmas,
                                     scale_weights, horizon_decay, eps):
    sample_field = preds["u"]
    time_steps = sample_field.shape[1]

    # Normalize weights
    scale_w = _normalize_weights(scale_weights, ...)
    horizon_raw = [horizon_decay ** step for step in range(time_steps)]
    horizon_w = _normalize_weights(horizon_raw, ...)

    losses = {}
    for field, ref_field in PRED_TO_REF_KEY.items():
        total = 0.0
        mask = masks[field]
        for step in range(time_steps):
            pred_step = preds[field][:, step:step+1, ...]
            ref_step = refs[ref_field][:, step:step+1, ...]
            step_loss = 0.0
            for sigma, sigma_weight in zip(sigmas, scale_w):
                pred_filtered = masked_gaussian_lowpass(pred_step, mask, sigma, eps)
                ref_filtered = masked_gaussian_lowpass(ref_step, mask, sigma, eps)
                step_loss += sigma_weight * masked_mse(pred_filtered, ref_filtered, mask, eps)
            total += horizon_w[step] * step_loss
        losses[field] = total
    return losses

Masked MSE

The masked MSE applies spatial masking to the standard MSE:

def masked_mse(pred, ref, mask, eps):
    weights = mask.expand_as(pred)       # broadcast mask
    diff_sq = (pred - ref).square() * weights
    return diff_sq.sum() / weights.sum().clamp_min(eps)

This is equivalent to computing MSE only over ocean points:

$$\text{MaskedMSE}(\mathbf{p},\mathbf{r},\mathbf{m})=\frac{\sum _{i,j}\mathbf{m}[i,j]\cdot (\mathbf{p}[i,j]-\mathbf{r}[i,j]{)}^2}{\sum _{i,j}\mathbf{m}[i,j]}$$

Static Loss Masks

The land/sea masks are static — computed once from the physics core geometry and reused throughout training. They are tensors of shape (H, W) with values 1.0 (ocean) or 0.0 (land).

For 5D tensors (when the time dimension is kept as a channel for $B\times T\times C\times H\times W$), the mask is expanded to match:

def _expand_mask(mask, target):
    if target.dim() == 4:
        return mask.view(1, 1, H, W).expand_as(target)
    if target.dim() == 5:
        return mask.view(1, 1, H, W, 1).expand_as(target)

Per-Field Loss Weights

Each prognostic field has a weight in the composite loss function. These weights are hyperparameters set in the trainer, not in supervision.py. The default weights (defined in NPOMTrainerDDP) are:

Field	Default Weight	Physical Magnitude	Justification
u	12.0	~0.1-1.0 m/s	Primary prognostic
ub	18.0	~0.01-0.1 m/s	Barotropic, larger scale
v	10.0	~0.1-1.0 m/s	Primary prognostic
vb	12.0	~0.01-0.1 m/s	Barotropic, larger scale
t	1.0	~0-30 C (after mean subtraction: ~±5)	Wide range but physically constrained
s	6.0	~30-40 PSU (small range, but important for density)	Critical for buoyancy
el	30.0	~±1 m	Small magnitude, needs upweighting
ua	60.0	~0.01-0.1 m/s	Small magnitude, important for transport
uab	60.0	~0.01 m/s	Very small magnitude
va	60.0	~0.01-0.1 m/s	Small magnitude, important for transport
vab	60.0	~0.01 m/s	Very small magnitude

Composite Loss Function

The composite loss is built by combine_loss_components() in the trainer:

def combine_loss_components(self, losses):
    w = self.loss_weights
    return (
        w["u"]  * losses["u"]   + w["ub"]  * losses["ub"] +
        w["v"]  * losses["v"]   + w["vb"]  * losses["vb"] +
        w["t"]  * losses["t"]   + w["s"]   * losses["s"]  +
        w["el"] * losses["el"]  + w["ua"]  * losses["ua"] +
        w["uab"]* losses["uab"] + w["va"]  * losses["va"] +
        w["vab"]* losses["vab"]
    ) * self.loss_display_scale

Formally, the composite loss is:

$$\mathcal{L}={\lambda }_d\cdot \sum _{f\in \mathcal{F}}{w}_f\cdot {\mathcal{L}}_f$$

where $\mathcal{F}$ is the set of all 12 prognostic fields, $w_f$ are the per-field weights, ${\mathcal{L}}_f$ is the (possibly multi-scale, horizon-decayed) per-field loss, and ${\lambda }_d$ is loss_display_scale.

loss_display_scale

The loss_display_scale parameter (default: 1.0) is a numerical convenience factor that scales all displayed and logged loss values. It does not affect the optimization — it is applied when logging, not during the backward pass:

# In combine_loss_components:
total_loss = weighted_sum * self.loss_display_scale

# In _reduce_losses: individual field losses are also scaled
losses["u"] * self.loss_display_scale

This is useful when the raw loss values are very small (e.g., ${10}^{-6}$) and you want more readable log output. The effective optimization loss is $\mathcal{L}/{\lambda }_d$ since the loss_display_scale factor is a constant multiplied uniformly.

Design Considerations

Why Masked Losses?

The POM model uses an Arakawa C-grid where land points have degenerate physics (zero velocity, fixed temperature/salinity). Including these points in the loss would:

1. Artificially lower the apparent loss (many points have near-zero error)
2. Penalize the model for "errors" in regions where physics is intentionally disabled
3. Bias gradient updates toward correcting land points rather than ocean dynamics

Why Multi-Scale?

Ocean dynamics span a wide range of spatial scales, from basin-scale gyres ($\sim$10,000km)tomesoscaleeddies($\sim$10-100 km). A single-scale MSE treats all scales equally, but from a practical forecasting perspective, getting the large-scale circulation right is more important than matching every small-scale eddy exactly.

The Gaussian multi-scale approach provides a principled way to balance these scales. Setting $\sigma =0$ captures all scales (including grid-scale noise), while larger $\sigma$ values emphasize the basin-scale patterns.

Why Horizon Decay?

Due to the chaotic nature of ocean dynamics (sensitive dependence on initial conditions), the practical predictability horizon is limited. After a few weeks to months, even a perfect model would diverge from the true trajectory due to the butterfly effect. Horizon decay acknowledges this fundamental limitation and focuses the training signal on the time range where skill is achievable.