Weighting Training Loss by Annotation Confidence

A per-annotation confidence score is only worth computing if something consumes it. Routing low-confidence labels to review is one consumer; the other is the training loss, where a confidence of 0.55 becomes a smaller gradient than a confidence of 0.95. The implementation is a few lines, and all of the difficulty is in three choices: the mapping from confidence to weight, the floor that stops the mapping becoming deletion, and the normalisation that keeps the effective learning rate steady as batch composition changes. This guide covers all three, the segmentation variant that needs a per-pixel weight map, and the ablation that says whether any of it helped.

# Why This Matters in Geospatial Pipelines

Geospatial labels are noisy in a structured way. The uncertain ones cluster: shadowed valleys, haze, class boundaries that the taxonomy never resolved, tiles annotated in the first week before the guide settled. Training as though every label is equally true lets those regions pull the model as hard as clean ground, and because the noise is spatially clustered rather than random, it does not average out.

The naive fix — delete anything below a threshold — is worse. Low confidence marks hard examples, and a model trained only on easy ones is confidently wrong exactly where a human would have hesitated.

# Step-by-Step Implementation

# Step 1 — Carry the Score Into the Sample

pip install torch==2.3.1 geopandas==0.14.4 rasterio==1.3.10 numpy==1.26.4

The loader must return the confidence alongside the image and target, so the training step never needs a second lookup.

from dataclasses import dataclass
import torch
from torch.utils.data import Dataset

@dataclass(frozen=True)
class Sample:
    image: torch.Tensor       # C × H × W
    target: torch.Tensor      # H × W for masks, or a label index
    confidence: float         # the per-annotation score, 0–1

class WeightedTileDataset(Dataset):
    def __init__(self, records, transform=None):
        self.records, self.transform = records, transform

    def __getitem__(self, i: int) -> Sample:
        r = self.records[i]
        image, target = load_tile(r["image_path"]), load_target(r["target_path"])
        if self.transform:
            image, target = self.transform(image, target)
        return Sample(image, target, float(r["confidence"]))

    def __len__(self) -> int:
        return len(self.records)

Storing confidence in the record rather than recomputing it at load time matters for reproducibility: the score depends on annotator agreement and model calibration at the time the batch was built, both of which move.

# Step 2 — Map Confidence to a Weight, With a Floor

def confidence_to_weight(conf: torch.Tensor, floor: float = 0.2,
                         c_min: float = 0.5, c_max: float = 0.95) -> torch.Tensor:
    """Linear map from [c_min, c_max] to [floor, 1.0], clamped at both ends.

    floor > 0 keeps hard examples in the batch instead of silently deleting them.
    """
    scaled = (conf - c_min) / max(c_max - c_min, 1e-6)
    return torch.clamp(scaled, 0.0, 1.0) * (1.0 - floor) + floor

Three properties are deliberate. The map is linear, because anything steeper amounts to a soft threshold and reintroduces the deletion behaviour. It saturates at c_max, since the difference between 0.95 and 0.99 confidence is not meaningful. And the floor is well above zero, so a batch always has the same effective size.

# Step 3 — Apply Per Sample and Normalise

The common mistake is multiplying an already-reduced loss, which scales the whole batch rather than its members.

import torch.nn.functional as F

def weighted_loss(logits: torch.Tensor, target: torch.Tensor,
                  weight: torch.Tensor) -> torch.Tensor:
    """Per-sample weighted cross-entropy, normalised by the weight sum."""
    per_sample = F.cross_entropy(logits, target, reduction="none")   # N
    if per_sample.dim() > 1:                        # segmentation: N × H × W
        per_sample = per_sample.flatten(1).mean(1)
    return (per_sample * weight).sum() / weight.sum().clamp_min(1e-6)

Dividing by weight.sum() rather than by N is what keeps the effective learning rate steady. Without it, a batch that happens to contain many low-confidence samples produces a smaller total loss and therefore a smaller step, so the optimiser’s behaviour depends on batch composition — a source of training instability that is very hard to attribute later.

# Step 4 — Rasterise Weights for Segmentation

For mask tasks a scalar per tile is too coarse: one tile can hold a confidently drawn warehouse and an uncertain field boundary. Burn each feature’s weight into a per-pixel map.

import numpy as np
from rasterio.features import rasterize

def weight_map(features, transform, out_shape, floor: float = 0.2,
               background: float = 1.0) -> np.ndarray:
    """Per-pixel weights: each feature burns its own weight, background keeps `background`."""
    shapes = [(f.geometry, confidence_to_weight(torch.tensor(f.confidence), floor=floor).item())
              for f in features]
    burned = rasterize(shapes, out_shape=out_shape, transform=transform,
                       fill=background, dtype="float32", all_touched=False)
    return burned

Two decisions worth stating. Background keeps a weight of 1.0, because “no object here” is usually a confident statement even when the objects present are uncertain. And all_touched=False matches whatever the label rasterisation used — a weight map built with a different rule than the mask it weights is misaligned at every boundary, which is the subtlest version of the problem rasterizing vector labels for segmentation masks describes.

# Step 5 — Ablate It

def ablation_report(runs: dict[str, dict]) -> str:
    """Compare weighted against unweighted on the same frozen evaluation set."""
    lines = ["arm            mIoU   noisy-class IoU   clean-class IoU"]
    for name, m in runs.items():
        lines.append(f"{name:<14} {m['miou']:.3f}  {m['noisy']:.3f}            {m['clean']:.3f}")
    return "\n".join(lines)

The expected shape of the result is a gain concentrated on classes with noisy labels and nothing much elsewhere:

arm            mIoU   noisy-class IoU   clean-class IoU
unweighted     0.681  0.512             0.774
weighted       0.698  0.561             0.776

That is a real but modest effect, which is exactly why it needs a control arm rather than an impression. Run both on the same blocked split — comparing a weighted run on one split against an unweighted run on another measures the splits.

# Parameters and Thresholds Reference

# Common Errors and Fixes

Training loss drops but validation gets worse Root cause: the loss was reduced before weighting, so the weight scaled the whole batch and effectively lowered the learning rate. Fix: use reduction="none" and normalise by the weight sum, as in Step 3.

Loss becomes nan after a few hundred steps Root cause: a batch in which every weight is the floor, combined with a division by a near-zero weight sum. Fix: the clamp_min(1e-6) above, plus a floor high enough that a full-floor batch is still a reasonable denominator.

Segmentation weights are misaligned with the mask Root cause: the weight map was rasterised with a different all_touched setting than the label mask. Fix: rasterise both with one function and one setting; assert their shapes and transforms match before training.

A rare class disappears from the model’s output Root cause: class weighting and confidence weighting multiplied, and that class is both rare and uncertain. Fix: log the mean effective weight per class for one epoch and rebalance; the product, not either factor, is what the optimiser sees.

# Frequently Asked Questions

# Related

• Confidence Scoring for Geospatial Labels — where the score comes from, and the QA routing that consumes it alongside this
• Calibrating Confidence Scores with Temperature Scaling — why an uncalibrated score makes this weighting arbitrary
• Rasterizing Vector Labels for Segmentation Masks — the rasterisation settings the weight map must match exactly
• Reproducible Train/Validation Splits for Spatial Data — the split both ablation arms have to share for the comparison to mean anything

This technique is one consumer of the scores produced in Confidence Scoring for Geospatial Labels, part of Geospatial Annotation Fundamentals & Architecture.