Calibrated Disambiguation for Partial Multi-label Learning

In plain words

Some image datasets are labelled by crowd workers who tick every tag that might apply. Many ticks are wrong, so a model trained on them learns the mistakes too. A common fix is to trust the model's confident predictions more than its unsure ones, and to slowly let it learn from harder examples. The open question is how much to trust each tick, and the usual answer is a hand-picked rule.

This paper says that the right amount of trust for a tick is the chance that the tick is actually correct, given how confident the model is about it. Turning that into a rule means estimating the chance from the model's own confidence statistics. The authors train a small helper network on one dataset (VOC) at several noise rates to make that estimate, then reuse it elsewhere. They add a second training signal that pulls each class embedding toward a learned prototype for its tag. The reported result is better accuracy than several earlier methods on four image datasets with added label noise.

TL;DR

• Partial multi-label learning: each image carries candidate tags that include wrong ones. The method weights each candidate by a calibrated estimate of being correct §1 §3.
• A pre-trained two-layer calibrator turns a 20-bin confidence histogram into per-bin reliabilities, and those reliabilities are the loss weights §4.1 alg. 1.
• A prototype-cosine loss is added in the second phase §4.2 eq. 13.
• Author-reported: best mAP in 15 of 16 noise settings against nine baselines §5.2 tab. 1.
• The official code differs from the paper in several places (weight rescaling, prototype range, frozen query embeddings, smoothing of calibrator targets). See Paper vs code.

Problem

Crowd-sourced multi-label datasets contain irrelevant labels that look relevant. Existing methods use the memorisation effect and small-loss selection §1 §2, but their weighting functions are heuristic (hard threshold, soft threshold, temperature annealing) and have no stated optimum §1 §2. Those methods also use a fixed rule, although the best threshold moves with dataset, noise rate and epoch §1. The paper's premise is that without an optimal weighting function the methods cannot be compared, and that the weighting itself has a provable target §1 §3.

Method

Weighting target

Treat each candidate positive as a label that is correct with some probability given its confidence t. The paper defines the curriculum loss as a per-entry weight times the binary cross-entropy eq. 1. Choosing the weight that minimises the expected loss gives the probability that the candidate is a true positive eq. 3 eq. 4. The paper states this optimum is equivalent to calibrating the model's confidences, and argues it gives an upper bound for small-loss disambiguation §3. Because the weight needs the unknown ground truth y, the method has to estimate it §3.

\ell_\text{curr}(\hat y, \tilde y) = w \cdot \ell_\text{BCE}(\hat y, \tilde y) eq. 1

w^*(t) = P(y = \tilde y \mid t) eq. 3 eq. 4

Here (\tilde y) is the candidate label, (y) the true label, and (\hat y) the predicted confidence §3 eq. 1.

Two-phase training

1. Warm-up. Train with BCE on the raw candidate labels until early stopping. The prototype loss is not used here §4 eq. 12 alg. 1.
2. Disambiguation. Each epoch, collect the confidences on the training set, compute the weights from the calibrator, then train with the weighted classification loss plus the prototype loss, until early stopping §4 eq. 13 alg. 1.

Network

Features come from a pre-trained CNN backbone (ResNet-101 in the experiments) §5.1. A transformer decoder with one query vector per class turns these features into one representation per class, in the style of ML-Decoder §4. Two heads read each class representation: a per-class linear classifier eq. 8 and a prototype cosine eq. 10. The paper gives no dimensions for the decoder or its query vectors §4. Figure 3 shows the overall layout fig. 3.

\hat y^\text{clf}_{ij} = \theta_j^\top z_{ij} + b_j eq. 8

\hat y^\text{proto}_{ij} = \frac{\langle z_{ij}, \bar z_j\rangle}{\lVert z_{ij}\rVert \, \lVert \bar z_j\rVert} eq. 10

Here (z_{ij}) is class j's representation for image i, (\theta_j, b_j) the classifier parameters for class j, and (\bar z_j) the prototype for class j §4 §4.2.

Losses

\mathcal{L}_\text{clf} = \sum_i \sum_j \ell_\text{curr}(\hat y^\text{clf}_{ij}, \tilde y_{ij}) eq. 9

\mathcal{L}_\text{proto} = \sum_i \sum_j \ell_\text{curr}(\hat y^\text{proto}_{ij}, \tilde y_{ij}) eq. 11

\mathcal{L}_\text{warmup} = \mathcal{L}_\text{clf}; \qquad \mathcal{L}_\text{disam} = \mathcal{L}_\text{clf} + \mathcal{L}_\text{proto} eq. 12 eq. 13

The prototypes are updated by gradient; the paper gives no separate update rule and no initialisation §4.2. The paper describes the prototype term as making embeddings cluster around their class prototype, and the classifier term as making them linearly separable §4.2.

Calibrator

Only candidate positives ((\tilde y = 1)) are used. Their confidences are split into true positives (TP) and false positives (FP), and each is counted in K = 20 bins §4.1 eq. 5:

q^P_k = \frac{|\{t\}^{[k]} \cap \{t\}^P|}{|\{t\}^P|}, \quad q^{TP}_k = \frac{|\{t\}^{[k]} \cap \{t\}^{TP}|}{|\{t\}^P|}, \quad q^{FP}_k = \frac{|\{t\}^{[k]} \cap \{t\}^{FP}|}{|\{t\}^P|} eq. 5

The reliability of bin k is the TP share of that bin eq. 6:

r_k = \frac{q^{TP}_k}{q^{TP}_k + q^{FP}_k} eq. 6

A confidence t is then mapped to a weight by interpolating between neighbouring reliabilities eq. 7:

P(y = \tilde y \mid t) = \left(\tfrac{k+1}{K} - t\right) r_k + \left(t - \tfrac{k}{K}\right) r_{k+1}, \quad \tfrac{k}{K} \le t < \tfrac{k+1}{K} eq. 7

The calibrator is a learned function (g: [0,1]^K \to [0,1]^K), from the histogram (q^P) to the reliabilities. It is a two-layer MLP with LeakyReLU §4.1. The paper says the calibrator is applied to each class separately §4.1. It is pre-trained on VOC 2007 with the identity-weighted baseline run at noise rates 10% to 90%. Each epoch of each run gives one ((q^P, r)) pair per class, and the calibrator is trained with MSE using Adam at learning rate 1e-2 §5.1. Once trained, it is used unchanged on other datasets and noise rates as a plug-in §4.1 §5.2.

Figure 1 shows existing weighting shapes, confidence histograms and the calibrator's prediction on unseen data fig. 1. Figure 2 shows an example of optimal weighting next to its ECE fig. 2.

Pseudocode

The following follows the paper's algorithm. The weights are only set on candidate positives, starting at 1 alg. 1 eq. 1.

# PML-CD training  [alg. 1] [§4]
# Input:  images (N, 3, H, W); partial labels (N, C) in {0, 1}
# Output: network with classifier and prototype heads

def train_pml_cd(images, partial_labels, calibrator):
    # 1. Warm-up: BCE on the raw partial labels  [alg. 1] [eq. 12]
    weights = ones_like(partial_labels)            # w_ij = 1  [alg. 1]
    while not early_stopping():
        clf_score, proto_score = network(images)   # (N, C) each
        loss = weighted_bce(clf_score, partial_labels, weights)  # [eq. 9]
        update(network, loss)

    # 2. Disambiguation: recompute weights every epoch  [alg. 1] [§4.1]
    while not early_stopping():
        # Confidences on the training set, without gradient  [alg. 1]
        clf_score = no_grad(network_clf(images))   # (N, C)
        # Histogram of confidences on candidate positives  [eq. 5]
        hist = positive_histogram(clf_score, partial_labels)  # (C, K)
        # Calibrator gives a reliability per bin  [eq. 6] [§4.1]
        reliability = calibrator(hist)             # (C, K)
        # Weights only where partial_labels == 1  [eq. 7] [alg. 1]
        weights = interpolate(clf_score, reliability, partial_labels)
        # Weighted classification and prototype losses  [eq. 13]
        clf_score, proto_score = network(images)   # (N, C) each
        loss = weighted_bce(clf_score, partial_labels, weights)
        loss = loss + weighted_bce(proto_score, partial_labels, weights)
        update(network, loss)
    return network

# Interpolated weight for one confidence t  [eq. 7]
# reliability: (K,) values, bins indexed in order  [§4.1]

def interpolate(t, reliability):
    # Bin index k with k/K <= t < (k+1)/K  [eq. 7]
    k = floor(t * K)                               # K = 20  [§4.1]
    frac = t * K - k                               # position inside the bin
    # Convex combination; see Pitfalls for the 1/K factor  [eq. 7]
    return (1 - frac) * reliability[k] + frac * reliability[k + 1]

# Prototype cosine per class  [eq. 10] [§4.2]
# class_embeds: (N, C, D); prototypes: (C, D)

def proto_score(class_embeds, prototypes):
    # Cosine between each sample's class embedding and its prototype
    return cosine(class_embeds, prototypes)        # (N, C), in [-1, 1]

Hyperparameters

Results

Author-reported. Mean average precision (mAP, %), ResNet-101 backbone, 224 × 224 inputs. Each cell is the mAP for one dataset and one noise rate. The paper's bold-face highlight marks the best per column, shown here in bold tab. 1.

Headline claim: PML-CD is best in 15 of 16 columns; the exception is VOC at 20% noise, where LL-R is higher §5.2 tab. 1. The paper attributes the VOC 20% result to simpler ResNet-only models being better suited to low noise §5.2.

Weighting policies on MS-COCO at 80% noise (epoch 1, ECE and mAP in %) tab. 2:

Prototype alignment ablation, mAP (%) tab. 3:

The paper concludes that the prototype term helps mainly at high noise §5.3 tab. 3. The "2nd phase only" COCO row matches the PML-CD row of tab. 1; the VOC row does not (see Pitfalls).

Paper vs code

Repository: https://github.com/lee-plus-plus/PML-CD (MIT licence, linked from the paper's own links block), at commit 6d578dd3834d0166b8b73b90390ebe25af099f61. The LICENSE file is MIT LICENSE. The code is not a reimplementation, since the paper links it directly. It contains the training script, the calibrator script, and a shipped calibrator checkpoint weighter.ckpt whose training data and provenance are not described in the paper PML_CD_train.py calibrator_train.py readme.md.

1. Extra weight rescaling. The paper has no such step. The code rescales every weight after the calibrator, with clamp(1.2 * w - 0.1, 0, 1), which sends weights below about 0.083 to 0 PML_CD_train.py. Negatives start at weight 1 and stay there, since the clamp maps 1 to 1 PML_CD_train.py PML_CD_train.py.
2. Prototype score range. The paper uses the raw cosine, which lies in [-1, 1] eq. 10. The code maps it to 0.5 + 0.5 · cos, so it lies in [0, 1] before the BCE PML_CD_train.py. The paper does not say how a value in [-1, 1] is passed to the BCE; the code's mapping makes it a valid probability.
3. Classifier output. The paper writes (\hat y^{clf} = \theta^\top z + b) eq. 8 and feeds it to BCE, which needs a probability. The code applies a sigmoid PML_CD_train.py.
4. Hard threshold. Paper: 0.6 §5.3. Code: 0.7 PML_CD_train.py.
5. Frozen class queries. The paper describes C query vectors §4. The code creates the query embedding and then calls requires_grad_(False) on it, so the embedding stays at its random initialisation. Only the projection layers around it are trained PML_CD_train.py.
6. Prototype initialisation. Not stated in the paper. The code uses random Gaussian vectors normalised to unit length, trainable PML_CD_train.py.
7. Calibrator target. Paper: r_k from eq. 6. The code's training target is the same ratio, plus a 3-tap smoothing and a cumulative maximum calibrator_train.py. The smoothing and monotone envelope are not in the paper.
8. Calibrator output and input. The code feeds the logit of each histogram bin into the MLP, and applies softmax followed by a cumulative sum to the output, so the output is monotone PML_CD_train.py. The paper describes a plain 2-layer MLP §4.1.
9. Interpolation geometry. The code places each reliability value at the bin centre and adds anchors (0, 0) and (1, 1) before linear interpolation PML_CD_train.py PML_CD_train.py. Eq. 7 places reliabilities at bin edges.
10. Calibrator collection run. The paper says the baseline used to collect calibrator data has no prototype alignment loss §5.1. The readme command passes --weighting identity but not --coef-proto 0, and the default is 1.0, so the prototype loss is on [readme.md#L12-L13] PML_CD_train.py.
11. Calibrator training. Code: 2000 full-batch epochs with weight decay 5e-5 calibrator_train.py. The paper gives only the optimiser, learning rate and loss.
12. Model selection on the test split. Early stopping and best-epoch selection use the validation loader, which is built from the test split PML_CD_train.py PML_CD_train.py PML_CD_train.py. The paper does not say which split was used.
13. Early stopping and epoch caps. The code uses patience 3 on mAP, capped at 20 epochs per phase PML_CD_train.py PML_CD_train.py. The paper gives neither number.
14. Learning rate by dataset. The paper gives 1e-5 for COCO and NUS-WIDE §5.1. The code has one default of 1e-4 and no per-dataset configuration PML_CD_train.py. The readme runs use the default.
15. Augmentation for weights. The code computes weights from the weakly augmented image and trains on the strongly augmented one PML_CD_train.py PML_CD_train.py. The paper does not mention a weak view.
16. Decoder details. The code uses a single decoder layer without self-attention. Its first residual adds dropout to the input rather than to an attention output PML_CD_train.py. The paper's description (ML-Decoder-style) does not state this.

Where code and paper agree: the warm-up uses only the classification loss with weight 1 PML_CD_train.py; the disambiguation loss uses both terms with coefficient 1 PML_CD_train.py; the noise model flips only zeros utils.py; the optimiser, batch size and image size match PML_CD_train.py PML_CD_train.py.

Pitfalls

• Printed Eq. 7 is off by a factor of 1/K. As printed, the two terms have coefficients summing to 1/K, so at t = k/K the weight is r_k / K. Implement it as a convex combination of the two neighbouring reliabilities, which is what the pseudocode above does eq. 7.
• Index convention is inconsistent. Eq. 5 defines bin k as (k−1)/K ≤ t < k/K. Eq. 7 uses k/K ≤ t < (k+1)/K, with r_k at the left edge. Edge values at t = 1 are not defined. The paper also gives no value for r_k when q^TP + q^FP is zero, which occurs for empty bins eq. 5 eq. 6 eq. 7. The code avoids division by zero with a 1e-4 term calibrator_train.py.
• Eq. 1 vs Algorithm 1. Eq. 1 writes the weight on every entry, but Algorithm 1 and the code only set weights on candidate positives eq. 1 alg. 1 PML_CD_train.py. Negatives keep weight 1.
• Weights need the ground truth in the optimal version. The "optimal" weighting in tab. 2 uses true labels, so it is an upper bound and not a deployable method §3 tab. 2 PML_CD_train.py.
• Calibrator pre-training data. The calibrator is trained on VOC with the identity baseline and then reused on all other datasets. Its transfer is shown indirectly through the main results and, for weight ECE, on COCO only §4.1 §5.2 tab. 2.
• Table 3 and Table 1 disagree on VOC. The "2nd phase only" VOC row of tab. 3 reads 88.77 / 88.12 / 85.07 / 71.70. The PML-CD VOC row of tab. 1 reads 88.83 / 87.12 / 83.87 / 71.70. Only the 80% value matches. The COCO rows match. The paper does not explain the gap.
• Table 2 "calibrator" and "optimal" have identical weight ECE (3.74). The paper does not discuss this. It may be coincidence or a copy error; treat the ECE comparison with care.
• Section numbers. The compiled PDF has section numbering switched off (secnumdepth set to 0), so no section numbers appear in the layout. The § numbers on this card count sections in source order: §1 Introduction, §2 Related Work, §3 Preliminaries, §4 Methods (§4.1 Calibrator, §4.2 Prototype Alignment Regularization), §5 Experiments (§5.1 Settings, §5.2 Comparison Results, §5.3 Ablation Study), §6 Conclusion.
• Equation numbers. The paper has 13 numbered equations. They were counted from the source in order, and not checked against the compiled PDF.
• Compute. Each disambiguation epoch runs a full pass over the training set to collect confidences, and the code also builds a per-class histogram each epoch PML_CD_train.py. Warm-up and disambiguation are each capped at 20 epochs in the code.
• Run count and variance. The paper does not state the number of runs or seeds behind each table cell, and gives no variance §5.1.

Lineage

• Builds on: the memorisation effect and small-loss selection §1 §2; the curriculum formulation eq. 1.
• Architecture borrowed from: ML-Decoder, for the transformer decoder and group linear projection §4.
• Compared against: BCE, ASL, ML-Decoder (multi-label); PML-NI, PML-LRS, PLAIN (traditional PML); LL-R, PML-CDCR, UNM (curriculum-based PML) §5.1 tab. 1.
• Prototype term: the paper cites earlier partial-label work using prototypes for disambiguation §4.2.