diff --git a/changelog.d/321.added.md b/changelog.d/321.added.md new file mode 100644 index 00000000..6bde5c96 --- /dev/null +++ b/changelog.d/321.added.md @@ -0,0 +1,10 @@ +**`--loss noise_corrected_bce` generalized to multiclass (`--num-out N > 1`).** +`losses.NoiseCorrectedCrossEntropyLoss` extends the binary forward correction +(#279) with a single-sink noise structure: each non-sink class has its own +measured purity, and every class's label noise flows only to one designated +sink class named by `--noise-sink-class` (resolved via `label_map`, or a raw +index). At `num_out > 1`, `--label-noise-rate` keys are resolved as class +labels rather than per-sample `source_group` lookups; all-zero rates reduce +to plain cross-entropy bit-for-bit, matching the binary loss's own +no-correction fast path. The resolved sink class/index and rates are +recorded in `config.json` for auditability. diff --git a/docs/api/training.md b/docs/api/training.md index 78362d63..7839e246 100644 --- a/docs/api/training.md +++ b/docs/api/training.md @@ -108,3 +108,29 @@ noisy regime, judge it by **callable yield at a fixed precision floor** (e.g. AUROC-adjacent metrics can look flat or even move against it while the fraction of confidently-called reads at 99% precision improves) rather than by AUROC alone. + +### Multiclass (`--num-out > 1`) + +The same `--loss noise_corrected_bce` value trains a multiclass forward +correction at `--num-out N` for `N > 1`, for a corpus where every class's +label noise flows to one designated **sink** class -- e.g. a 20+1-class +charge-aware classifier (20 amino acids + `uncharged`) where a known fraction +of each amino acid's labeled chunks are secretly `uncharged`. Name the sink +with `--noise-sink-class uncharged` (a class label resolved against the +corpus's `label_map.json`, or a raw index); `--label-noise-rate` keys are then +class labels rather than arbitrary `source_group` values (e.g. +`--label-noise-rate Gln=0.041,Thr=0.948`), one purity per non-sink class. As +in the binary case, a class not named gets rate 0 and an all-zero +`--label-noise-rate` (or none at all) reproduces `--loss cross_entropy` +bit-for-bit. + +The two losses differ in what the correction needs: the binary loss looks up +one flip rate per *sample* from its `source_group`, but the multiclass sink +structure needs the *global* per-class rate vector to correctly weigh how +much of a sink-labeled chunk's evidence should be credited to each other +class -- so rates are resolved once from `label_map`, not read per chunk. +Only samples observed as the sink actually get a different gradient from +plain cross-entropy; a non-sink-observed sample's correction is a per-class +constant that does not move the gradient at all, since the sink is the only +class more than one other class can leak into. See +`leech.losses.NoiseCorrectedCrossEntropyLoss` for the full derivation. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 12f53c9e..f3298eea 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -289,7 +289,8 @@ from a fresh optimizer and logs a warning. |--------|---------|-------------| | `--loss STR` | `bce` | Loss function: `bce`, `focal`, `cross_entropy`, or `noise_corrected_bce` | | `--focal-gamma FLOAT` | `2.0` | Focal loss gamma (only with `--loss focal`) | -| `--label-noise-rate STR` | -- | `group=rate[,group=rate,...]` per-`source_group` label-flip probability for `--loss noise_corrected_bce` (e.g. `gold=0.09,enzymatic=0.17`). Unmapped groups get rate 0. See [Label-noise-aware loss](../api/training.md#label-noise-aware-loss). | +| `--label-noise-rate STR` | -- | `group=rate[,group=rate,...]` label-flip probability for `--loss noise_corrected_bce` (e.g. `gold=0.09,enzymatic=0.17`). At `--num-out 1`, keys are per-`source_group`; at `--num-out > 1`, keys are class labels (see `--noise-sink-class`). Unmapped values get rate 0. See [Label-noise-aware loss](../api/training.md#label-noise-aware-loss). | +| `--noise-sink-class STR` | -- | Class label (resolved via `label_map`) or raw index that `--loss noise_corrected_bce` treats as the sink every other class's label noise flows to. Required by that loss at `--num-out > 1`; unused for binary. | | `--focal-neg-gamma FLOAT` | none (symmetric) | Separate gamma for negative-labeled examples, making the focal loss asymmetric (only with `--loss focal`). Larger than `--focal-gamma` down-weights easy negatives harder, shaping the loss for a low-FPR operating regime | **Data augmentation:** @@ -387,7 +388,7 @@ leech model optimize --train-data FILE --output-dir DIR --context-grid VALUES [O | `--base-justify STR` | `center` | Signal chunk centering | | `--parallel INT` | `1` | Grid points to train concurrently | -Training options (`--model`, `--epochs`, `--batch-size`, `--learning-rate`, `--device`, `--seed`, `--early-stopping`) work the same as in `model train`. So do the loss and augmentation options above `--early-stopping` in that reference (`--loss`, `--label-noise-rate`, `--focal-gamma`/`--focal-neg-gamma`, `--scheduler*`, `--warmup-epochs`, `--weight-decay`, `--max-grad-norm`, every `--augment-*` flag, `--label-smoothing`, `--mixed-precision`) and `--motif`/`--motif-offset` from provenance -- both decorators (`training_hyperparams`, `model_provenance`) are shared between `train` and `optimize`. `--selection-metric` accepts the same names as `model train`'s `--checkpoint-metric` (see the "Output files" note under `model train` above: `auto`/`val_acc`/`val_f1`/`val_auc`, plus the parametric `tpr_at_fpr:` and `callable_at_precision:

`). +Training options (`--model`, `--epochs`, `--batch-size`, `--learning-rate`, `--device`, `--seed`, `--early-stopping`) work the same as in `model train`. So do the loss and augmentation options above `--early-stopping` in that reference (`--loss`, `--label-noise-rate`, `--noise-sink-class`, `--focal-gamma`/`--focal-neg-gamma`, `--scheduler*`, `--warmup-epochs`, `--weight-decay`, `--max-grad-norm`, every `--augment-*` flag, `--label-smoothing`, `--mixed-precision`) and `--motif`/`--motif-offset` from provenance -- both decorators (`training_hyperparams`, `model_provenance`) are shared between `train` and `optimize`. `--selection-metric` accepts the same names as `model train`'s `--checkpoint-metric` (see the "Output files" note under `model train` above: `auto`/`val_acc`/`val_f1`/`val_auc`, plus the parametric `tpr_at_fpr:` and `callable_at_precision:

`). **Output:** diff --git a/src/leech/cli.py b/src/leech/cli.py index 8a5f7414..6ecd5dc7 100644 --- a/src/leech/cli.py +++ b/src/leech/cli.py @@ -846,10 +846,24 @@ def merge( type=str, default=None, help=( - "Per-source_group label-flip probability for --loss noise_corrected_bce, " - "as 'group=rate[,group=rate,...]' (e.g. 'gold=0.09,enzymatic=0.17'). Rates " - "are measured upstream, not estimated by leech; unmapped source_group " - "values get rate 0 (no correction)." + "Label-flip probability for --loss noise_corrected_bce, as " + "'group=rate[,group=rate,...]' (e.g. 'gold=0.09,enzymatic=0.17'). At " + "--num-out 1, keys are per-source_group. At --num-out > 1, keys are " + "class labels (resolved via label_map, or a raw class index) naming " + "each class's own noise-flip rate into --noise-sink-class. Rates are " + "measured upstream, not estimated by leech; unmapped values get rate 0 " + "(no correction)." + ), +) +@click.option( + "--noise-sink-class", + type=str, + default=None, + help=( + "Class label (resolved via label_map) or raw index that " + "--loss noise_corrected_bce treats as the sink every other class's " + "label noise flows to (issue #321). Required by that loss at " + "--num-out > 1; unused for binary (--num-out 1)." ), ) @click.option( @@ -953,6 +967,7 @@ def train( confound, confound_config, label_noise_rate, + noise_sink_class, cl_regression, cl_lambda, signal_mode, @@ -1032,6 +1047,7 @@ def train( adversarial_anneal_epochs=adversarial_anneal_epochs, confound=confound, label_noise_rates=label_noise_rates, + noise_sink_class=noise_sink_class, cl_regression=cl_regression, cl_lambda=cl_lambda, signal_mode=signal_mode, @@ -1529,10 +1545,24 @@ def fetch(name, model_version, tag, output_dir, repo): type=str, default=None, help=( - "Per-source_group label-flip probability for --loss noise_corrected_bce, " - "as 'group=rate[,group=rate,...]' (e.g. 'gold=0.09,enzymatic=0.17'). Rates " - "are measured upstream, not estimated by leech; unmapped source_group " - "values get rate 0 (no correction)." + "Label-flip probability for --loss noise_corrected_bce, as " + "'group=rate[,group=rate,...]' (e.g. 'gold=0.09,enzymatic=0.17'). At " + "--num-out 1, keys are per-source_group. At --num-out > 1, keys are " + "class labels (resolved via label_map, or a raw class index) naming " + "each class's own noise-flip rate into --noise-sink-class. Rates are " + "measured upstream, not estimated by leech; unmapped values get rate 0 " + "(no correction)." + ), +) +@click.option( + "--noise-sink-class", + type=str, + default=None, + help=( + "Class label (resolved via label_map) or raw index that " + "--loss noise_corrected_bce treats as the sink every other class's " + "label noise flows to (issue #321). Required by that loss at " + "--num-out > 1; unused for binary (--num-out 1)." ), ) @click.option( @@ -1611,6 +1641,7 @@ def optimize( confound, confound_config, label_noise_rate, + noise_sink_class, cl_regression, cl_lambda, signal_mode, @@ -1673,6 +1704,7 @@ def optimize( adversarial_anneal_epochs=adversarial_anneal_epochs, confound=confound, label_noise_rates=label_noise_rates, + noise_sink_class=noise_sink_class, cl_regression=cl_regression, cl_lambda=cl_lambda, signal_mode=signal_mode, diff --git a/src/leech/cli_options.py b/src/leech/cli_options.py index c45a05ae..c4038c48 100644 --- a/src/leech/cli_options.py +++ b/src/leech/cli_options.py @@ -260,7 +260,9 @@ def training_hyperparams(f): type=click.Choice(["bce", "focal", "cross_entropy", "noise_corrected_bce"]), default=DEFAULT_LOSS_TYPE, help="Loss function type. 'noise_corrected_bce' applies a forward " - "correction for known, per-source_group label noise; see --label-noise-rate.", + "correction for known label noise; see --label-noise-rate. At " + "--num-out > 1 this trains a multiclass forward correction with a " + "single sink class (see --noise-sink-class) instead of the binary one.", )(f) f = click.option( "--warmup-epochs", diff --git a/src/leech/commands/optimize.py b/src/leech/commands/optimize.py index 412fbe15..4e34eca5 100644 --- a/src/leech/commands/optimize.py +++ b/src/leech/commands/optimize.py @@ -61,6 +61,7 @@ def handle_optimize( adversarial_anneal_epochs: int = 0, confound: str | None = None, label_noise_rates: dict[str, float] | None = None, + noise_sink_class: str | None = None, cl_regression: bool = False, cl_lambda: float = 1.0, signal_mode: str = "both", @@ -103,7 +104,12 @@ def handle_optimize( num_workers: DataLoader workers balance_groups: Balance sampling across source groups label_noise_rates: ``{source_group: flip_rate}`` for - ``loss_type="noise_corrected_bce"``; unmapped groups get rate 0 + ``loss_type="noise_corrected_bce"``; unmapped groups get rate 0. + At ``num_out > 1`` keys are class labels instead -- see + ``leech.losses.build_class_noise_rates``. + noise_sink_class: Sink class name/index for + ``loss_type="noise_corrected_bce"`` at ``num_out > 1``; unused + for binary Returns: Path to grid search summary file @@ -158,6 +164,7 @@ def handle_optimize( oversample_minority=oversample_minority, confound=confound, label_noise_rates=label_noise_rates, + noise_sink_class=noise_sink_class, signal_mode=signal_mode, optim=OptimConfig( learning_rate=learning_rate, diff --git a/src/leech/commands/train.py b/src/leech/commands/train.py index 4dbd43f1..751fbf47 100644 --- a/src/leech/commands/train.py +++ b/src/leech/commands/train.py @@ -68,6 +68,7 @@ def handle_train( adversarial_anneal_epochs: int = 0, confound: str | None = None, label_noise_rates: dict[str, float] | None = None, + noise_sink_class: str | None = None, cl_regression: bool = False, cl_lambda: float = 1.0, signal_mode: str = "both", @@ -128,7 +129,12 @@ def handle_train( dataset the first time it happens. balance_groups: Balance sampling across source groups label_noise_rates: ``{source_group: flip_rate}`` for - ``loss_type="noise_corrected_bce"``; unmapped groups get rate 0 + ``loss_type="noise_corrected_bce"``; unmapped groups get rate 0. + At ``num_out > 1`` keys are class labels instead -- see + ``leech.losses.build_class_noise_rates``. + noise_sink_class: Sink class name/index for + ``loss_type="noise_corrected_bce"`` at ``num_out > 1``; unused + for binary sample_weight_field: Chunk metadata field to inverse-frequency weight sampling by (mutually exclusive with balance_groups and oversample_minority), e.g. "junction_indel" to over-sample the @@ -217,6 +223,7 @@ def handle_train( "adversarial_anneal_epochs", "confound", "label_noise_rates", + "noise_sink_class", "cl_regression", "cl_lambda", "signal_mode", @@ -291,6 +298,7 @@ def handle_train( label_map=label_map, confound=confound, label_noise_rates=label_noise_rates, + noise_sink_class=noise_sink_class, optim=OptimConfig( learning_rate=learning_rate, weight_decay=weight_decay, diff --git a/src/leech/configs.py b/src/leech/configs.py index 2f0e5df7..65fc1354 100644 --- a/src/leech/configs.py +++ b/src/leech/configs.py @@ -292,8 +292,16 @@ class TrainConfig: focal_gamma: float = 2.0 #: ``{source_group: flip_rate}`` for ``loss_type="noise_corrected_bce"``, #: parsed from ``--label-noise-rate`` by ``leech.losses.parse_label_noise_rate``. - #: Unmapped groups get rate 0. Recorded verbatim in config.json. + #: Unmapped groups get rate 0. Recorded verbatim in config.json. At + #: ``num_out > 1`` keys are resolved as class labels (via ``label_map``), + #: not per-sample source_group lookups -- see + #: ``leech.losses.build_class_noise_rates``. label_noise_rates: dict[str, float] | None = None + #: Sink class name/index for ``loss_type="noise_corrected_bce"`` at + #: ``num_out > 1`` -- the class every other class's noise mass flows to + #: (issue #321). Resolved by ``leech.losses.resolve_noise_sink_index``. + #: Unused (and not required) for the binary loss. + noise_sink_class: str | None = None # Asymmetric focal loss (--focal-neg-gamma, issue #280): None keeps the # symmetric loss bit-for-bit; see FocalBCEWithLogitsLoss for why. focal_neg_gamma: float | None = None diff --git a/src/leech/losses.py b/src/leech/losses.py index e41b8cb8..6156459b 100644 --- a/src/leech/losses.py +++ b/src/leech/losses.py @@ -4,11 +4,23 @@ from __future__ import annotations +import logging + import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Function +logger = logging.getLogger("leech.losses") + +#: Stand-in for log(0) in the multiclass noise correction below. Real -inf +#: differentiates logsumexp to nan whenever every summand in a row is masked +#: out (-inf - (-inf) = nan in the max-subtraction); a large finite floor +#: keeps the same masking effect (exp(-1e30) underflows to exactly 0) while +#: staying differentiable. leech.crf's CTC-CRF loss hits the identical +#: landmine -- see its own ``_UNREACHABLE``. +_NEG_INF = -1e30 + class GradientReversalFunction(Function): """Reverse gradients during backward pass (Ganin et al. 2016).""" @@ -132,6 +144,216 @@ def parse_label_noise_rate(token: str | None) -> dict[str, float] | None: return rates or None +def resolve_noise_sink_index( + noise_sink_class: str, label_map: dict[str, int] | None, num_out: int +) -> int: + """Resolve ``--noise-sink-class`` to a class index. + + Accepts either a raw integer index (as a string) or a class name looked + up in ``label_map`` (``{class_name: index}``, the same mapping + ``leech.confounds`` uses for label-keyed confounds -- normally loaded + from a corpus's ``label_map.json`` sidecar). Raises rather than silently + training with an unresolved sink: a wrong sink index corrupts the + corrected loss on every sink-observed sample with no shape error to + catch it (see :class:`NoiseCorrectedCrossEntropyLoss`). + """ + try: + index = int(noise_sink_class) + except (TypeError, ValueError): + if label_map is None: + raise ValueError( + f"--noise-sink-class {noise_sink_class!r} is not an integer index " + "and no label_map is available to resolve it by name (expected a " + "label_map.json sidecar next to the training data)." + ) from None + if noise_sink_class not in label_map: + raise ValueError( + f"--noise-sink-class {noise_sink_class!r} not found in label_map " + f"(known classes: {sorted(label_map)})" + ) from None + index = label_map[noise_sink_class] + if not (0 <= index < num_out): + raise ValueError( + f"--noise-sink-class {noise_sink_class!r} resolves to index {index}, " + f"out of range for num_out={num_out}" + ) + return index + + +def build_class_noise_rates( + label_noise_rates: dict[str, float] | None, + label_map: dict[str, int] | None, + sink_index: int, + num_out: int, +) -> torch.Tensor | None: + """Build the per-class flip-rate vector for :class:`NoiseCorrectedCrossEntropyLoss`. + + Unlike the binary loss's per-*sample* ``noise_rate`` (looked up per chunk + from its own ``source_group``), the multiclass correction needs a + per-*class* rate: computing the corrected probability of an + observed-as-sink sample requires knowing every OTHER class's leak rate, + not just the rate of whichever group this one sample happens to belong + to -- see the class docstring. ``label_noise_rates`` keys (from + ``--label-noise-rate``, :func:`parse_label_noise_rate`) are therefore + resolved as class labels here, the same way :func:`resolve_noise_sink_index` + resolves ``--noise-sink-class``: an integer index, or a name via + ``label_map``. + + An unresolvable key (no ``label_map``, or a name it does not contain) is + dropped with a warning rather than raising -- the same lenient "unmapped + defaults to 0 (no correction)" contract the binary loss's per-sample + lookup already has. The sink class's own rate is always forced to 0 + (``T[sink, sink] = 1``, clean by construction), regardless of what + ``label_noise_rates`` says about it. + + Returns ``None`` when ``label_noise_rates`` is empty/``None`` -- the + short-circuit that keeps :class:`NoiseCorrectedCrossEntropyLoss` + bit-for-bit equal to plain cross-entropy for the common no-correction + case. + """ + if not label_noise_rates: + return None + rates = torch.zeros(num_out, dtype=torch.float32) + for key, rate in label_noise_rates.items(): + try: + index = int(key) + except ValueError: + if label_map is None or key not in label_map: + logger.warning( + "--label-noise-rate group '%s' does not match a class name in " + "label_map; ignoring (no correction applied for it)", + key, + ) + continue + index = label_map[key] + if not (0 <= index < num_out): + logger.warning( + "--label-noise-rate group '%s' resolves to index %d, out of range " + "for num_out=%d; ignoring", + key, + index, + num_out, + ) + continue + if index == sink_index: + if rate != 0: + logger.warning( + "--label-noise-rate gives the sink class ('%s') a nonzero rate " + "(%.4g); ignoring -- the sink class is always treated as clean " + "(T[sink, sink] = 1)", + key, + rate, + ) + continue + rates[index] = rate + return rates + + +class NoiseCorrectedCrossEntropyLoss(nn.Module): + """Multiclass forward-corrected cross-entropy for known, class-conditional + label noise flowing to a single sink class (Patrini et al. 2017) -- the + C-class generalization of :class:`NoiseCorrectedBCEWithLogitsLoss` + (issue #321). + + Each non-sink class ``i`` has a measured purity ``pi_i = 1 - rho_i`` + (``rho_i`` from ``--label-noise-rate``, resolved to class indices by + :func:`build_class_noise_rates`); the noise transition matrix is + + T[i, i] = pi_i # true=i (i != sink) -> observed i + T[i, sink] = rho_i # true=i (i != sink) -> observed sink + T[sink, sink] = 1 # true=sink -> observed sink, always (clean) + + every other entry 0: a class only ever leaks into the sink, never into + another class, and the sink never leaks at all. The model's softmax + output ``p`` estimates the *true*-label posterior; the loss is NLL of the + *observed* label ``y`` under the corrected distribution ``q = T^T p``: + + q_j = p_j * pi_j for j != sink + q_sink = p_sink + sum_{i != sink} p_i * rho_i + + Because every non-sink column of ``T`` has exactly one nonzero entry (its + own diagonal), ``q_y = p_y * pi_y`` for an observed non-sink label -- a + single term, not a mixture. ``log(pi_y)`` does not depend on the model, + so for a sample observed as a non-sink class this loss has *the same + gradient* as plain cross-entropy (a per-class constant shift in the loss + value only, invisible to the optimizer). All of the correction's effect + is on samples observed as the sink: there, ``q_sink`` mixes in every + other class's leaked mass, weighted by the model's own current belief in + that class -- crediting the model for recognizing contamination instead + of forcing it to explain that signal as "sink" (issue #321's motivating + case: sink-labeled chunks a known fraction of which are secretly some + other, contaminating class). + + ``class_rates=None`` (no ``--label-noise-rate`` at all) short-circuits to + plain ``F.cross_entropy``, bit-for-bit -- exactly mirroring + ``NoiseCorrectedBCEWithLogitsLoss``'s ``noise_rate is None`` fast path, + but checked once at construction rather than per forward call: + ``class_rates`` is a fixed vector here, not a per-batch tensor, so there + is no per-step CUDA host-sync to avoid by deferring the check. + + Args: + sink_index: Class index the noise mass flows to (resolved from + ``--noise-sink-class`` by :func:`resolve_noise_sink_index`). + class_rates: ``(num_out,)`` per-class flip rate; entry + ``sink_index`` is always 0 (:func:`build_class_noise_rates` + enforces this when building it from ``--label-noise-rate``). + ``None`` or all-zero takes the plain cross-entropy fast path. + weight: Optional per-class weight, same convention as + ``nn.CrossEntropyLoss(weight=...)``. + """ + + def __init__( + self, + sink_index: int, + class_rates: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + ) -> None: + super().__init__() + self.sink_index = sink_index + self.weight = weight + self._has_noise = class_rates is not None and bool(torch.any(class_rates > 0).item()) + self._log_rho: torch.Tensor | None = None + self._log_pi: torch.Tensor | None = None + if self._has_noise: + assert class_rates is not None + self._log_rho = torch.where( + class_rates > 0, + torch.log(class_rates.clamp_min(1e-38)), + torch.full_like(class_rates, _NEG_INF), + ) + # log(pi) = log(1 - rho); 0 (pi=1, clean) wherever rate is 0, + # including at sink_index (build_class_noise_rates forces that). + self._log_pi = torch.where( + class_rates > 0, torch.log1p(-class_rates), torch.zeros_like(class_rates) + ) + + def log_probs(self, logits: torch.Tensor) -> torch.Tensor: + """``log(T^T softmax(logits))``, i.e. the corrected log-probabilities + ``forward`` takes NLL of. Exposed separately so a caller normalizing + by something other than element count (``Trainer._weighted_ce_global``, + for the DDP weighted-mean case) can build its own reduction over the + same corrected distribution instead of reimplementing it. + """ + if not self._has_noise: + return F.log_softmax(logits, dim=-1) + + assert self._log_rho is not None and self._log_pi is not None + log_p = F.log_softmax(logits, dim=-1) + log_rho = self._log_rho.to(dtype=log_p.dtype, device=log_p.device) + log_pi = self._log_pi.to(dtype=log_p.dtype, device=log_p.device) + + log_q = log_p + log_pi.unsqueeze(0) + leak = torch.logsumexp(log_p + log_rho.unsqueeze(0), dim=-1) + sink_col = torch.logsumexp(torch.stack([log_p[:, self.sink_index], leak]), dim=0) + log_q[:, self.sink_index] = sink_col + return log_q + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + if not self._has_noise: + return F.cross_entropy(logits, targets, weight=self.weight) + return F.nll_loss(self.log_probs(logits), targets, weight=self.weight) + + class NoiseCorrectedBCEWithLogitsLoss(nn.Module): """Forward-corrected BCE for class-conditional label noise with known, per-sample flip rates (forward correction; Patrini et al. 2017). diff --git a/src/leech/model_loading.py b/src/leech/model_loading.py index 4fe6caf7..ada94c60 100644 --- a/src/leech/model_loading.py +++ b/src/leech/model_loading.py @@ -132,6 +132,8 @@ def load_model_from_checkpoint( "adversarial_anneal_epochs", "confound", "label_noise_rates", + "noise_sink_class", + "noise_sink_index", "cl_lambda", "checkpoint_metric", "signal_mode", diff --git a/src/leech/training.py b/src/leech/training.py index 5d62aeb7..ec344f3f 100644 --- a/src/leech/training.py +++ b/src/leech/training.py @@ -57,7 +57,10 @@ AdversarialHead, FocalBCEWithLogitsLoss, NoiseCorrectedBCEWithLogitsLoss, + NoiseCorrectedCrossEntropyLoss, RegressionHead, + build_class_noise_rates, + resolve_noise_sink_index, ) from leech.metrics import ( PARAMETRIC_SELECTION_METRICS, @@ -487,6 +490,8 @@ def __init__( checkpoint_metric: str = "auto", dist: DistContext | None = None, cfg: TrainConfig | None = None, + noise_class_rates: torch.Tensor | None = None, + noise_sink_index: int | None = None, ): # The recipe as one object (#270). Every caller above still passes # loose kwargs (43 test call sites plus train_model's own kwarg path) @@ -562,6 +567,17 @@ def __init__( # (defaulted to 1), so unpacking it would turn "not yet known" into # "1" for every caller that passes num_out=None. # + # noise_class_rates/noise_sink_index are ALSO not resolved from + # cfg.label_noise_rates/cfg.noise_sink_class here, for the same + # reason num_out isn't: resolving a class name against label_map + # needs the corpus's label_map, which (like num_out) may only be + # known to the caller (train_model loads it from a label_map.json + # sidecar, possibly after auto-detecting num_out itself) and not yet + # written back into cfg by the time Trainer is constructed. The + # caller resolves both and passes them in already-built, the same + # way pos_weight arrives as a Tensor rather than a class-frequency + # dict for Trainer to compute itself. + # # self.cfg is currently write-only: nothing in this class reads it # back, and it does not enter a checkpoint payload. It is also not # re-synced after construction, so self.cfg.pos_weight (float-typed @@ -689,22 +705,28 @@ def __init__( self._num_out = num_out if num_out is not None else 1 self.label_smoothing = label_smoothing pw = pos_weight.to(device) if pos_weight is not None else None + + def _resolve_ce_weight(pw: torch.Tensor | None) -> torch.Tensor | None: + """``pos_weight`` -> a CrossEntropyLoss-style per-class weight. + + Binary (num_out<=2): a single positive-class weight becomes the + 2-entry [1.0, pw] CE convention. Multiclass (num_out>2): pw is + already the per-class weight tensor. Shared by the + 'cross_entropy' and multiclass 'noise_corrected_bce' branches so + the two stay identical rather than drifting. + """ + if pw is None: + return None + if self._num_out <= 2: + return torch.tensor([1.0, pw.item()], dtype=torch.float32).to(device) + return pw.to(device) + if loss_type == "cross_entropy": # CrossEntropyLoss expects (B, num_classes) logits and (B,) integer labels - if pw is not None and self._num_out <= 2: - # Convert pos_weight to per-class weights for CE (binary case) - ce_weights = torch.tensor([1.0, pw.item()], dtype=torch.float32).to(device) - self.criterion = nn.CrossEntropyLoss( - weight=ce_weights, label_smoothing=label_smoothing - ) - elif pw is not None and self._num_out > 2: - # Multiclass: pw is already a per-class weight tensor - self.criterion = nn.CrossEntropyLoss( - weight=pw.to(device), label_smoothing=label_smoothing - ) + ce_weight = _resolve_ce_weight(pw) + self.criterion = nn.CrossEntropyLoss(weight=ce_weight, label_smoothing=label_smoothing) + if ce_weight is not None and self._num_out > 2: logger.info(f"Using class weights for {self._num_out}-class CE") - else: - self.criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing) logger.info(f"Using CrossEntropyLoss ({self._num_out}-class)") elif loss_type == "focal": self.criterion = FocalBCEWithLogitsLoss( @@ -718,6 +740,33 @@ def __init__( ) else: logger.info(f"Using focal loss (gamma={focal_gamma}{neg_gamma_note})") + elif loss_type == "noise_corrected_bce" and self._num_out > 1: + # Multiclass forward correction (issue #321): a fixed, per-class + # rate vector baked in at construction, resolved by the caller + # (see the comment above self.cfg) -- not the per-sample + # 'noise_rate' batch field the binary branch below reads, which + # this loss has no use for (NoiseCorrectedCrossEntropyLoss's own + # docstring explains why a per-class vector is required instead). + if label_smoothing > 0: + raise ValueError( + "label_smoothing is not supported with loss_type='noise_corrected_bce' " + "at num_out>1; pass --label-smoothing 0 (the default) or use " + "--loss cross_entropy instead." + ) + if noise_sink_index is None: + raise ValueError( + "loss_type='noise_corrected_bce' with num_out>1 requires a resolved " + "sink class -- pass --noise-sink-class on the CLI." + ) + self.criterion = NoiseCorrectedCrossEntropyLoss( + sink_index=noise_sink_index, + class_rates=noise_class_rates.to(device) if noise_class_rates is not None else None, + weight=_resolve_ce_weight(pw), + ) + logger.info( + f"Using noise-corrected cross-entropy ({self._num_out}-class, forward " + f"correction); sink_index={noise_sink_index}" + ) elif loss_type == "noise_corrected_bce": self.criterion = NoiseCorrectedBCEWithLogitsLoss(pos_weight=pw) logger.info( @@ -731,12 +780,14 @@ def __init__( self.criterion = nn.BCEWithLogitsLoss() logger.info("Training without class weighting") # Weighted CE normalizes by the summed weight of the samples it sees, - # not by their count, so it is the one loss here that does not - # decompose over shards. See _weighted_ce_global. + # not by their count, so it is the one loss family here that does not + # decompose over shards. See _weighted_ce_global. NoiseCorrectedCrossEntropyLoss + # shares this exactly: its forward ends in F.nll_loss(..., weight=...), + # the same weighted-mean reduction as nn.CrossEntropyLoss(weight=...). self._ce_needs_global_norm = ( self.dist.enabled - and loss_type == "cross_entropy" - and getattr(self.criterion, "weight", None) is not None + and isinstance(self.criterion, (nn.CrossEntropyLoss, NoiseCorrectedCrossEntropyLoss)) + and self.criterion.weight is not None ) if label_smoothing > 0 and loss_type != "cross_entropy": logger.info(f"Label smoothing={label_smoothing} (applied to binary targets)") @@ -1116,8 +1167,13 @@ def _compute_batch_loss( # Forward pass (wrapper handles moving tensors and calling model correctly) logits = self.model_wrapper.forward_batch(batch, self.device) - if self.loss_type == "cross_entropy": - # CrossEntropyLoss wants (B,) integer class labels + is_multiclass_noise_corrected = ( + self.loss_type == "noise_corrected_bce" and self._num_out > 1 + ) + if self.loss_type == "cross_entropy" or is_multiclass_noise_corrected: + # Both want (B,) integer class labels; the multiclass correction + # bakes its per-class rates in at construction (self.criterion), + # so unlike the binary branch below it reads nothing from batch. ce_targets = labels.squeeze(-1).long() if self._ce_needs_global_norm: main_loss = self._weighted_ce_global(logits, ce_targets) @@ -1159,7 +1215,8 @@ def _batch_noise_rate(self, batch: dict[str, Any]) -> torch.Tensor | None: return noise_rate.to(self.device, non_blocking=True) def _weighted_ce_global(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: - """Weighted cross-entropy normalized by the GLOBAL summed weight. + """Weighted cross-entropy (or its noise-corrected multiclass sibling) + normalized by the GLOBAL summed weight. Every other loss in this trainer reduces by element count, which equal shards make exact: DDP's gradient average over N shards of n/N samples @@ -1169,7 +1226,14 @@ def _weighted_ce_global(self, logits: torch.Tensor, targets: torch.Tensor) -> to one only when the shards happen to draw the same classes. On a fixture whose shards differ that is an 8% error in the loss, and nothing raises: the run converges to a slightly different recipe than the same - command line asked for at ``--gpus 1``. + command line asked for at ``--gpus 1``. ``NoiseCorrectedCrossEntropyLoss`` + shares this exactly (its own ``forward`` ends in the same + ``F.nll_loss(..., weight=...)`` reduction), which is why this method + takes the log-probabilities from ``self.criterion.log_probs`` when + available rather than assuming plain ``nn.CrossEntropyLoss`` -- + label_smoothing is the one thing that diverges between the two, and + the noise-corrected branch refuses a nonzero one at construction, so + reusing plain ``F.nll_loss`` here (no smoothing term) is exact for it. Each rank instead forms ``world_size * (its weighted sum) / (the global weighted sum)``, whose average over ranks telescopes back to exactly @@ -1178,13 +1242,19 @@ def _weighted_ce_global(self, logits: torch.Tensor, targets: torch.Tensor) -> to on this path. """ weight = self.criterion.weight - numerator = nn.functional.cross_entropy( - logits, - targets, - weight=weight, - label_smoothing=self.label_smoothing, - reduction="sum", - ) + log_probs_fn = getattr(self.criterion, "log_probs", None) + if log_probs_fn is not None: + numerator = nn.functional.nll_loss( + log_probs_fn(logits), targets, weight=weight, reduction="sum" + ) + else: + numerator = nn.functional.cross_entropy( + logits, + targets, + weight=weight, + label_smoothing=self.label_smoothing, + reduction="sum", + ) denominator = all_reduce_tensor_sum(weight.detach()[targets].sum(), self.dist) return self.dist.world_size * numerator / denominator @@ -1259,7 +1329,9 @@ def train_epoch( tally.add("cl_loss", cl_loss, num_splits) labels_flat = labels.detach().flatten() - if self.loss_type == "cross_entropy": + if self.loss_type == "cross_entropy" or ( + self.loss_type == "noise_corrected_bce" and self._num_out > 1 + ): # argmax(logits) is exactly the softmax > 0.5 threshold # (softmax is monotonic per-row in the logits), for # both binary and multi-class CE -- so this reads the @@ -1374,8 +1446,12 @@ def validate( # Move labels to device labels = batch["label"].to(self.device, non_blocking=True) - # Adapt labels for CrossEntropyLoss - if self.loss_type == "cross_entropy": + # Adapt labels for CrossEntropyLoss (and its multiclass + # noise-corrected sibling, which takes the same integer + # targets and reads nothing from the batch itself). + if self.loss_type == "cross_entropy" or ( + self.loss_type == "noise_corrected_bce" and self._num_out > 1 + ): ce_labels = labels.squeeze(-1).long() else: ce_labels = None @@ -1436,11 +1512,16 @@ def validate( tally.add("loss", loss) tally.add("loss_sum", loss * labels.shape[0]) val_seen += int(labels.shape[0]) - if self.loss_type == "cross_entropy" and self._num_out > 2: + # A softmax head either way -- plain cross_entropy or its + # multiclass noise-corrected sibling (ce_labels is not None + # is exactly that pair, see above) -- reads probabilities the + # same way; only the sigmoid (num_out=1) losses differ. + uses_softmax_head = ce_labels is not None + if uses_softmax_head and self._num_out > 2: probs_mc = torch.softmax(logits, dim=-1).cpu().numpy() all_probs_mc.append(probs_mc) all_preds.append(probs_mc.argmax(axis=-1).ravel()) - elif self.loss_type == "cross_entropy": + elif uses_softmax_head: all_preds.append(torch.softmax(logits, dim=-1)[:, 1].cpu().numpy().ravel()) else: all_preds.append(torch.sigmoid(logits).cpu().numpy().ravel()) @@ -1976,6 +2057,7 @@ def train_model( adversarial_anneal_epochs: int = 0, confound: str | None = None, label_noise_rates: dict[str, float] | None = None, + noise_sink_class: str | None = None, cl_regression: bool = False, cl_lambda: float = 1.0, signal_mode: str = "both", @@ -2103,6 +2185,7 @@ def train_model( label_map=label_map, confound=confound, label_noise_rates=label_noise_rates, + noise_sink_class=noise_sink_class, optim=OptimConfig( learning_rate=learning_rate, weight_decay=weight_decay, @@ -2163,6 +2246,7 @@ def train_model( label_map = cfg.label_map confound = cfg.confound label_noise_rates = cfg.label_noise_rates + noise_sink_class = cfg.noise_sink_class learning_rate = cfg.optim.learning_rate weight_decay = cfg.optim.weight_decay max_grad_norm = cfg.optim.max_grad_norm @@ -2727,11 +2811,40 @@ def train_model( except Exception as e: logger.warning(f"torch.compile failed, falling back to eager mode: {e}") - # Auto-detect cross-entropy models (num_out > 1) - if num_out > 1 and loss_type != "cross_entropy": + # Auto-detect cross-entropy models (num_out > 1). noise_corrected_bce is + # its own multiclass-capable family (issue #321, dispatched on num_out + # inside Trainer) and must not be clobbered here -- doing so silently + # discarded an explicit `--loss noise_corrected_bce --num-out N` request. + if num_out > 1 and loss_type not in ("cross_entropy", "noise_corrected_bce"): loss_type = "cross_entropy" logger.info(f"Model {model_name} has num_out={num_out}, switching to cross_entropy loss") + # Resolve the multiclass noise-corrected sink/rates now that num_out and + # loss_type are both final (issue #321). label_map may still be None here + # even at num_out > 1: the sidecar load above only runs when num_out was + # auto-detected from the data, not when the caller passed --num-out + # explicitly, so this repeats that lookup rather than trusting it ran. + noise_class_rates = None + noise_sink_index = None + if loss_type == "noise_corrected_bce" and num_out > 1: + if label_map is None and train_data_path is not None: + _sink_lm_path = train_data_path.parent / "label_map.json" + if not _sink_lm_path.exists(): + _sink_lm_path = train_data_path.parent.parent / "label_map.json" + if _sink_lm_path.exists(): + with open(_sink_lm_path) as f: + label_map = json.load(f) + logger.info(f"Loaded label_map from {_sink_lm_path}: {label_map}") + if noise_sink_class is None: + raise ValueError( + "--loss noise_corrected_bce with num_out>1 requires --noise-sink-class " + "to name the class the noise mass flows to." + ) + noise_sink_index = resolve_noise_sink_index(noise_sink_class, label_map, num_out) + noise_class_rates = build_class_noise_rates( + label_noise_rates, label_map, noise_sink_index, num_out + ) + # Introspect feature_start/feature_end from the training corpus. # # Resolved from the whole corpus (`train_dataset.chunks`, a ChunkTable @@ -2848,6 +2961,12 @@ def train_model( # Recorded verbatim (not resolved against the corpus) -- predict never # reads this back; it is provenance for what the training run applied. "label_noise_rates": label_noise_rates, + # noise_sink_class as given on the CLI, plus the index it resolved to + # (None unless loss_type="noise_corrected_bce" at num_out>1) -- so a + # checkpoint is auditable after the fact without re-resolving it + # against a label_map that may since have changed (issue #321). + "noise_sink_class": noise_sink_class, + "noise_sink_index": noise_sink_index, "cl_regression": cl_regression, "cl_lambda": cl_lambda, "signal_mode": signal_mode, @@ -2910,7 +3029,9 @@ def train_model( # Create trainer. cfg carries the recipe (learning_rate, scheduler, # augmentation, aux-head lambdas, ...); the rest are runtime objects cfg # doesn't describe, plus num_out, which may have been auto-detected from - # the data above and so can differ from cfg.num_out. + # the data above and so can differ from cfg.num_out -- noise_class_rates/ + # noise_sink_index bypass cfg for the same reason (see the comment in + # Trainer.__init__). trainer = Trainer( model=model, model_type=model_name, @@ -2924,6 +3045,8 @@ def train_model( adversarial_num_classes=adversarial_num_classes, dist=dist_ctx, cfg=cfg, + noise_class_rates=noise_class_rates, + noise_sink_index=noise_sink_index, ) # Train diff --git a/tests/test_training_advanced.py b/tests/test_training_advanced.py index 45aa2054..db821daa 100644 --- a/tests/test_training_advanced.py +++ b/tests/test_training_advanced.py @@ -13,11 +13,15 @@ from torch.utils.data import DataLoader import leech.training +from leech.chunking import save_chunks from leech.dataset import LeechDataset, collate_fn from leech.losses import ( FocalBCEWithLogitsLoss, NoiseCorrectedBCEWithLogitsLoss, + NoiseCorrectedCrossEntropyLoss, + build_class_noise_rates, parse_label_noise_rate, + resolve_noise_sink_index, ) from leech.models import get_model from leech.training import Trainer, train_model @@ -321,6 +325,431 @@ def test_batch_noise_rate_reaches_the_criterion(self, temp_chunks_file, model_co assert not torch.allclose(main_loss, plain) +class TestNoiseSinkResolution: + """Test resolve_noise_sink_index / build_class_noise_rates (issue #321).""" + + LABEL_MAP = {"Gln": 0, "Thr": 1, "uncharged": 2} + + def test_resolves_by_name(self): + assert resolve_noise_sink_index("uncharged", self.LABEL_MAP, 3) == 2 + + def test_resolves_by_raw_index(self): + assert resolve_noise_sink_index("2", None, 3) == 2 + assert resolve_noise_sink_index("2", self.LABEL_MAP, 3) == 2 + + def test_unknown_name_without_label_map_raises(self): + with pytest.raises(ValueError, match="no label_map"): + resolve_noise_sink_index("uncharged", None, 3) + + def test_unknown_name_not_in_label_map_raises(self): + with pytest.raises(ValueError, match="not found in label_map"): + resolve_noise_sink_index("nonexistent", self.LABEL_MAP, 3) + + def test_out_of_range_index_raises(self): + with pytest.raises(ValueError, match="out of range"): + resolve_noise_sink_index("5", self.LABEL_MAP, 3) + + def test_build_class_noise_rates_none_when_empty(self): + assert build_class_noise_rates(None, self.LABEL_MAP, 2, 3) is None + assert build_class_noise_rates({}, self.LABEL_MAP, 2, 3) is None + + def test_build_class_noise_rates_resolves_by_name(self): + rates = build_class_noise_rates({"Gln": 0.1, "Thr": 0.3}, self.LABEL_MAP, 2, 3) + assert torch.equal(rates, torch.tensor([0.1, 0.3, 0.0])) + + def test_build_class_noise_rates_resolves_by_index(self): + rates = build_class_noise_rates({"0": 0.1, "1": 0.3}, self.LABEL_MAP, 2, 3) + assert torch.equal(rates, torch.tensor([0.1, 0.3, 0.0])) + + def test_sink_own_rate_always_forced_to_zero(self): + """Even if the caller's ``--label-noise-rate`` names the sink class, + its own rate is ignored -- T[sink, sink] = 1 by construction.""" + rates = build_class_noise_rates({"Gln": 0.1, "uncharged": 0.9}, self.LABEL_MAP, 2, 3) + assert rates[2].item() == 0.0 + + def test_unmapped_group_dropped_not_raised(self): + """Same lenient contract as the binary loss's per-sample lookup: + an unresolvable key is ignored (rate 0), not fatal.""" + rates = build_class_noise_rates({"Gln": 0.1, "Nonexistent": 0.5}, self.LABEL_MAP, 2, 3) + assert torch.equal(rates, torch.tensor([0.1, 0.0, 0.0])) + + +class TestNoiseCorrectedCrossEntropyLoss: + """Test NoiseCorrectedCrossEntropyLoss, the multiclass generalization of + NoiseCorrectedBCEWithLogitsLoss (issue #321).""" + + def test_no_rates_matches_plain_ce_bit_for_bit(self): + torch.manual_seed(0) + logits = torch.randn(8, 4) + targets = torch.randint(0, 4, (8,)) + + corrected = NoiseCorrectedCrossEntropyLoss(sink_index=3)(logits, targets) + plain = torch.nn.functional.cross_entropy(logits, targets) + assert torch.equal(corrected, plain) + + def test_all_zero_rates_matches_plain_ce_bit_for_bit(self): + """An explicit all-zero rate vector also takes the fast path (checked + once at construction, not per forward call -- see the class + docstring).""" + torch.manual_seed(1) + logits = torch.randn(8, 4) + targets = torch.randint(0, 4, (8,)) + zero_rates = torch.zeros(4) + + corrected = NoiseCorrectedCrossEntropyLoss(sink_index=3, class_rates=zero_rates)( + logits, targets + ) + plain = torch.nn.functional.cross_entropy(logits, targets) + assert torch.equal(corrected, plain) + + def test_sink_own_rate_ignored(self): + """A nonzero rate at the sink's own index must not change anything -- + build_class_noise_rates forces it to 0 before it ever reaches here, + but the loss itself must also not rely on that (T[sink,sink]=1 is + asserted by construction: log_pi is 0 wherever rate<=0, and the + forward pass never reads class_rates[sink] for anything but the + (masked-out) leak term).""" + torch.manual_seed(2) + logits = torch.randn(8, 4) + targets = torch.randint(0, 4, (8,)) + rates_clean_sink = torch.tensor([0.2, 0.3, 0.1, 0.0]) + + loss = NoiseCorrectedCrossEntropyLoss(sink_index=3, class_rates=rates_clean_sink)( + logits, targets + ) + assert torch.isfinite(loss) + + def test_nonzero_rates_diverge_from_plain_ce(self): + torch.manual_seed(3) + logits = torch.randn(16, 4) + targets = torch.randint(0, 4, (16,)) + rates = torch.tensor([0.2, 0.3, 0.1, 0.0]) + + corrected = NoiseCorrectedCrossEntropyLoss(sink_index=3, class_rates=rates)(logits, targets) + plain = torch.nn.functional.cross_entropy(logits, targets) + assert not torch.allclose(corrected, plain) + assert corrected.item() >= 0 + assert torch.isfinite(corrected) + + def test_gradient_is_finite(self): + torch.manual_seed(4) + logits = torch.randn(16, 4, requires_grad=True) + targets = torch.randint(0, 4, (16,)) + rates = torch.tensor([0.2, 0.3, 0.1, 0.0]) + + loss = NoiseCorrectedCrossEntropyLoss(sink_index=3, class_rates=rates)(logits, targets) + loss.backward() + assert torch.isfinite(logits.grad).all() + + def test_non_sink_gradient_matches_plain_ce(self): + """The single-sink structure has no cross-leakage between non-sink + classes: every non-sink observed column has exactly one nonzero T + entry (its own diagonal), so log(pi_y) is a per-class *constant* + that drops out of the gradient. All of the correction's effect is on + samples observed as the sink -- see the class docstring.""" + torch.manual_seed(5) + logits = torch.randn(16, 4, requires_grad=True) + targets = torch.randint(0, 3, (16,)) # never the sink (index 3) + rates = torch.tensor([0.2, 0.3, 0.1, 0.0]) + + logits_a = logits.detach().clone().requires_grad_() + logits_b = logits.detach().clone().requires_grad_() + NoiseCorrectedCrossEntropyLoss(sink_index=3, class_rates=rates)( + logits_a, targets + ).backward() + torch.nn.functional.cross_entropy(logits_b, targets).backward() + assert torch.allclose(logits_a.grad, logits_b.grad, atol=1e-6) + + def test_weighted_reduces_like_cross_entropy_loss(self): + """weight= must reduce the same way nn.CrossEntropyLoss(weight=...) + does (normalized by the weighted sum, not element count) -- this is + the property Trainer._weighted_ce_global relies on to stay correct + under DDP for this loss too.""" + torch.manual_seed(6) + logits = torch.randn(10, 3) + targets = torch.randint(0, 3, (10,)) + weight = torch.tensor([1.0, 2.0, 0.5]) + rates = torch.tensor([0.3, 0.0, 0.0]) + + corrected = NoiseCorrectedCrossEntropyLoss(sink_index=1, class_rates=rates, weight=weight)( + logits, targets + ) + log_q = NoiseCorrectedCrossEntropyLoss( + sink_index=1, class_rates=rates, weight=weight + ).log_probs(logits) + expected = torch.nn.functional.nll_loss(log_q, targets, weight=weight) + assert torch.allclose(corrected, expected) + + def test_synthetic_class_conditional_noise_recovers_better_than_plain_ce(self): + """A multiclass softmax classifier fit on labels where a known + fraction of each non-sink class's examples were relabeled to a sink + class, given the true (measured) per-class rates, must recover + better held-out accuracy against the TRUE (clean) labels than plain + cross-entropy trained on the same noisy labels -- the concrete claim + behind the multiclass generalization of --loss noise_corrected_bce + (issue #321).""" + torch.manual_seed(0) + n = 6000 + n_classes = 3 + sink = 2 + means = torch.tensor([[2.0, 0.0], [-2.0, 0.0], [0.0, -2.0]]) + true_y = torch.randint(0, n_classes, (n,)) + x = means[true_y] + 0.7 * torch.randn(n, 2) + + rho = torch.tensor([0.4, 0.6, 0.0]) + flip = torch.bernoulli(rho[true_y]).bool() + obs_y = torch.where(flip, torch.full_like(true_y, sink), true_y) + + def fit(*, corrected: bool) -> tuple[torch.Tensor, torch.Tensor]: + w = torch.zeros(2, n_classes, requires_grad=True) + b = torch.zeros(n_classes, requires_grad=True) + opt = torch.optim.Adam([w, b], lr=0.05) + loss_fn = NoiseCorrectedCrossEntropyLoss( + sink_index=sink, class_rates=rho if corrected else None + ) + for _ in range(400): + opt.zero_grad() + loss = loss_fn(x @ w + b, obs_y) + loss.backward() + opt.step() + return w.detach(), b.detach() + + w_c, b_c = fit(corrected=True) + w_p, b_p = fit(corrected=False) + + x_test = means.repeat_interleave(2000, dim=0) + 0.7 * torch.randn(6000, 2) + y_test = torch.arange(n_classes).repeat_interleave(2000) + + def acc(w: torch.Tensor, b: torch.Tensor) -> float: + preds = (x_test @ w + b).argmax(-1) + return (preds == y_test).float().mean().item() + + assert acc(w_c, b_c) > acc(w_p, b_p) + # Not just marginally better -- plain CE is badly attenuated toward + # the sink class here (rho up to 0.6), so the gap should be large. + assert acc(w_c, b_c) - acc(w_p, b_p) > 0.1 + + +class TestNoiseCorrectedCrossEntropyLossTrainer: + """Test the multiclass noise_corrected_bce dispatch in Trainer (issue #321).""" + + def _multiclass_model_and_loader(self, model_config, temp_chunks_file): + """A tiny 3-class ConvLSTMDwell + a matching 1-batch DataLoader, + reusing the binary fixtures' chunks file (label values don't matter + for these tests -- they only exercise loss dispatch, not correctness + of the corpus).""" + model = get_model("ConvLSTMDwell", **{**model_config, "num_out": 3}) + dataset = LeechDataset( + chunk_path=temp_chunks_file, + model_type="ConvLSTMDwell", + signal_len=model_config["signal_len"], + kmer_len=model_config["kmer_len"], + seq_encoding="base_onehot", + ) + loader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, collate_fn=collate_fn) + return model, loader + + def test_missing_sink_index_raises(self, model_config, temp_chunks_file): + model, loader = self._multiclass_model_and_loader(model_config, temp_chunks_file) + with pytest.raises(ValueError, match="noise-sink-class"): + Trainer( + model=model, + model_type="ConvLSTMDwell", + train_loader=loader, + device="cpu", + loss_type="noise_corrected_bce", + num_out=3, + ) + + def test_label_smoothing_raises(self, model_config, temp_chunks_file): + model, loader = self._multiclass_model_and_loader(model_config, temp_chunks_file) + with pytest.raises(ValueError, match="label_smoothing"): + Trainer( + model=model, + model_type="ConvLSTMDwell", + train_loader=loader, + device="cpu", + loss_type="noise_corrected_bce", + num_out=3, + label_smoothing=0.1, + noise_sink_index=2, + ) + + def test_criterion_is_multiclass_noise_corrected(self, model_config, temp_chunks_file): + model, loader = self._multiclass_model_and_loader(model_config, temp_chunks_file) + trainer = Trainer( + model=model, + model_type="ConvLSTMDwell", + train_loader=loader, + device="cpu", + loss_type="noise_corrected_bce", + num_out=3, + noise_sink_index=2, + noise_class_rates=torch.tensor([0.2, 0.3, 0.0]), + ) + assert isinstance(trainer.criterion, NoiseCorrectedCrossEntropyLoss) + assert trainer.criterion.sink_index == 2 + + def test_no_rates_still_builds_the_multiclass_criterion(self, model_config, temp_chunks_file): + """A sink index alone (no --label-noise-rate at all) is valid -- the + criterion just takes its own no-correction fast path, exactly the + binary loss's precedent.""" + model, loader = self._multiclass_model_and_loader(model_config, temp_chunks_file) + trainer = Trainer( + model=model, + model_type="ConvLSTMDwell", + train_loader=loader, + device="cpu", + loss_type="noise_corrected_bce", + num_out=3, + noise_sink_index=2, + ) + assert isinstance(trainer.criterion, NoiseCorrectedCrossEntropyLoss) + + def test_train_epoch_runs(self, model_config, temp_chunks_file): + model, loader = self._multiclass_model_and_loader(model_config, temp_chunks_file) + trainer = Trainer( + model=model, + model_type="ConvLSTMDwell", + train_loader=loader, + device="cpu", + loss_type="noise_corrected_bce", + num_out=3, + noise_sink_index=2, + noise_class_rates=torch.tensor([0.2, 0.3, 0.0]), + ) + loss, acc = trainer.train_epoch() + assert loss >= 0 + + def test_binary_loss_unaffected_by_num_out_default(self, sample_model, sample_dataloader): + """num_out=1 (the default) must still build the original binary loss + -- the new branch is keyed on num_out>1, not on loss_type alone.""" + trainer = Trainer( + model=sample_model, + model_type="ConvLSTMDwell", + train_loader=sample_dataloader, + device="cpu", + loss_type="noise_corrected_bce", + ) + assert isinstance(trainer.criterion, NoiseCorrectedBCEWithLogitsLoss) + + +def _multiclass_chunks(sample_leech_read, class_names=("classA", "classB", "sink")): + """Chunks cycling through len(class_names) labels/source_groups, for + exercising the multiclass noise_corrected_bce path end-to-end.""" + chunks = [] + num_bases = len(sample_leech_read.dwells) + kmer_context = 5 + start_idx = kmer_context + end_idx = min(15, num_bases - kmer_context) + + for i, base_idx in enumerate(range(start_idx, end_idx)): + chunk = sample_leech_read.get_chunk( + base_idx, signal_context=(200, 200), kmer_context=kmer_context + ) + if chunk is not None: + label = i % len(class_names) + chunk["read_id"] = sample_leech_read.read_id + chunk["label_int"] = label + chunk["label"] = class_names[label] + chunk["source_group"] = class_names[label] + chunks.append(chunk) + return chunks + + +class TestMulticlassNoiseCorrectedTrainModel: + """train_model end-to-end with --loss noise_corrected_bce at num_out>1 + (issue #321 acceptance criteria).""" + + @pytest.fixture + def multiclass_data(self, sample_leech_read, tmp_path): + class_names = ("classA", "classB", "sink") + chunks = _multiclass_chunks(sample_leech_read, class_names) + data_dir = tmp_path / "data" + data_dir.mkdir() + chunks_file = data_dir / "chunks.npz" + save_chunks(chunks, chunks_file) + label_map = {name: i for i, name in enumerate(class_names)} + with open(data_dir / "label_map.json", "w") as f: + json.dump(label_map, f) + return chunks_file, label_map + + def test_trains_without_error_and_records_config(self, multiclass_data, tmp_path): + chunks_file, label_map = multiclass_data + output_dir = tmp_path / "training" + history = train_model( + train_data_path=chunks_file, + val_data_path=chunks_file, + model_name="ConvLSTMDwell", + output_dir=output_dir, + epochs=1, + batch_size=2, + device="cpu", + motif="CCAGGC", + num_out=3, + loss_type="noise_corrected_bce", + label_noise_rates={"classA": 0.2, "classB": 0.3}, + noise_sink_class="sink", + ) + assert len(history["train_loss"]) == 1 + + with open(output_dir / "config.json") as f: + config = json.load(f) + # Not silently clobbered to cross_entropy by the num_out>1 auto-detect. + assert config["loss_type"] == "noise_corrected_bce" + assert config["num_out"] == 3 + assert config["noise_sink_class"] == "sink" + assert config["noise_sink_index"] == label_map["sink"] + assert config["label_noise_rates"] == {"classA": 0.2, "classB": 0.3} + + def test_missing_sink_class_raises(self, multiclass_data, tmp_path): + chunks_file, _label_map = multiclass_data + with pytest.raises(ValueError, match="noise-sink-class"): + train_model( + train_data_path=chunks_file, + val_data_path=chunks_file, + model_name="ConvLSTMDwell", + output_dir=tmp_path / "training", + epochs=1, + batch_size=2, + device="cpu", + motif="CCAGGC", + num_out=3, + loss_type="noise_corrected_bce", + ) + + def test_zero_rates_matches_plain_cross_entropy_checkpoint(self, multiclass_data, tmp_path): + """No --label-noise-rate at all (rates all default to 0) must train + an identical run to --loss cross_entropy -- the end-to-end version of + NoiseCorrectedCrossEntropyLoss's own bit-for-bit fast-path guarantee.""" + chunks_file, _label_map = multiclass_data + + def run(loss_type, output_dir, **extra): + torch.manual_seed(0) + return train_model( + train_data_path=chunks_file, + val_data_path=chunks_file, + model_name="ConvLSTMDwell", + output_dir=output_dir, + epochs=1, + batch_size=2, + device="cpu", + motif="CCAGGC", + num_out=3, + loss_type=loss_type, + seed=0, + **extra, + ) + + history_nc = run( + "noise_corrected_bce", + tmp_path / "nc", + noise_sink_class="sink", + ) + history_ce = run("cross_entropy", tmp_path / "ce") + assert history_nc["train_loss"] == pytest.approx(history_ce["train_loss"]) + + class TestAsymmetricFocalLoss: """Test FocalBCEWithLogitsLoss's --focal-neg-gamma asymmetry (#280)."""