Generalize noise_corrected_bce to multiclass (num_out > 1)

[jayhesselberth]

Motivation

--loss noise_corrected_bce (#279, NoiseCorrectedBCEWithLogitsLoss) is strictly binary — a single scalar flip rate rho per source_group, sigmoid-based (Patrini et al. 2017's 2x2 forward correction). training.py dispatches it as a separate branch from cross_entropy with no num_out > 1 path at all.

rnabioco/2026-aa-trna-models's production-v4 corpus is exactly the motivating case this loss was built for, except it's multiclass: a 20+1-class charge-aware classifier (20 amino acids + uncharged) where a chunk's "charged" label is measured to be wrong for a known fraction of reads in a given arm — a matched no-enzyme/no-amino-acid control's own cl-classifier false-positive rate gives a per-arm purity directly (already computed in that repo's charging_elevation_precheck.json, e.g. GlnRS_gln_b3 purity 0.959, ThrRS_thr_b3 purity 0.052). The noise is structured, not a dense confusion matrix: for amino-acid class i with purity π_i, a (1 - π_i) fraction of reads labeled i should actually be labeled uncharged — there is exactly one "sink" class the noise mass flows to, the same shape noise_corrected_bce already handles for binary.

Right now the only way to use the existing loss on that corpus is to decompose it into N binary one-vs-all classifiers (one per amino acid), losing the unified multiclass model and its shared representation.

Proposed mechanism

Generalize the existing forward-correction approach to num_out > 1:

• Model emits softmax probabilities p over C classes (an estimate of the true-label posterior).
• A per-source_group noise structure maps p to the observed-label distribution: P(observed = j) = sum_i T[i, j] * p_i, and the loss is NLL against the observed label under that corrected distribution.
• For this project's use case, T is sparse and simple: each class i has one scalar purity π_i and one designated sink class s (here, uncharged) that absorbs all its noise mass — i.e. T[i, i] = π_i, T[i, s] = 1 - π_i, T[i, j] = 0 otherwise, and the sink class's own row is identity (T[s, s] = 1, clean). This is a direct generalization of the existing binary 2x2 matrix, not a new noise model.
• A fully dense per-group C x C matrix would be more general but is not needed to unblock this project — see non-goals.

Suggested CLI surface

• Either extend --loss noise_corrected_bce to accept --num-out > 1, or add a new value (e.g. noise_corrected_ce) that shares the same --label-noise-rate group=rate[,group=rate,...] parsing already in place (parse_label_noise_rate).
• A new flag to name the sink class the noise mass flows to for the multiclass case, e.g. --noise-sink-class uncharged (or an index) — no equivalent needed in the binary case since there's only one other class.
• Same "measured upstream, not estimated by leech" contract as today; unmapped source_group values still default to rate 0 (no correction).

Acceptance criteria

• leech model train --num-out N --loss noise_corrected_bce --label-noise-rate group=rate,... --noise-sink-class LABEL (N > 2) trains without error and produces a valid multiclass checkpoint.
• With every group's rate at 0 (or no --label-noise-rate passed at all), output matches plain cross_entropy bit-for-bit or to floating-point precision — mirroring the existing binary short-circuit (noise_rate is None fast path) in NoiseCorrectedBCEWithLogitsLoss.
• Unit tests analogous to the existing binary loss's: a synthetic multiclass dataset with a known, injected class-conditional flip rate into a single sink class recovers the true posterior better under the corrected loss than under plain CE, measured on a held-out clean-labeled validation set.
• config.json records the resolved noise-rate mapping and sink class the same way the binary loss's config recording works today, so a checkpoint is auditable after the fact.
• Documented in CLI help and CHANGELOG, same shape as training: a label-noise-aware loss with per-sample noise rates #279's entry.

Non-goals

• No automatic noise-rate estimation — rates are supplied by the caller, same as today.
• No fully dense per-group C x C transition matrix on day one — the single-sink-class structure above is sufficient for this project's use case; a dense matrix can be a fast-follow if there's demand elsewhere.
• No changes to data prepare or the source_group chunk-metadata field/lookup semantics — reuse exactly what noise_corrected_bce already reads.

Reference

Motivated by rnabioco/2026-aa-trna-models production-v4 (issue #31 Phase 2, PR #36) and its version-bump retrain (PR #39) — happy to share the full per-arm purity table (charging_elevation_precheck.json's per_class_purity) if useful for building/validating the synthetic test case.

[jayhesselberth]

Deprioritizing this for now. It was motivated by 2026-aa-trna-models' per-arm purity table as a source of known label-noise rates, but a deep-dive there (rnabioco/2026-aa-trna-models#40) found we can't actually characterize the noise cleanly: the purity spread doesn't track the published aminoacyl-ester half-life ranking in the expected direction, and we can't rule out that it's partly `cl` generalizing unevenly across tRNA bodies rather than honest label noise. Feeding an uncharacterized rate into a forward-correction loss isn't better than not correcting at all. Leaving this open in case the noise gets properly characterized some other way later, but not actively planning to implement it.