feat(losses): generalize noise_corrected_bce to multiclass

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.

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.

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:<f>` and `callable_at_precision:<p>`).
+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:<f>` and `callable_at_precision:<p>`).
 
 **Output:**
 

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,

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",

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,

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,

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