General Multi-label Image Classification with Transformers

In plain words

A photo often contains several things at once, and the labels that describe it are not independent: some co-occur and some rarely appear together. A common way to label such photos is to train a network that scores each possible label on its own. That throws away the knowledge that labels depend on each other.

This paper's model, C-Tran, keeps that knowledge. It gives each possible label its own small token, and it feeds those tokens together with the image's visual features into a Transformer. The Transformer lets every label look at every other label and at every part of the image. Each label token also has a tag saying whether that label is unknown, known to be absent, or known to be present. During training, the model is shown random subsets of labels with their tags and asked to predict the hidden ones. This teaches it how labels relate to each other.

The payoff is flexibility. With no label information, the model predicts from the image alone. If someone already knows some of the labels, for example from a caption or from related metadata that is not itself a target label, they can tag those labels, and the model updates its predictions for the rest in a single pass.

TL;DR

• Multi-label image classifier: a Transformer encoder over ResNet-101 grid features plus one learned token per label, each with a state embedding (unknown, negative, positive) §3.1 §3.2.
• Label mask training hides a random share of labels in each training image and computes the loss only on hidden labels §3.4.
• Author-reported: 85.1 mAP on COCO-80 and 38.4 mAP on VG-500, above the compared methods tab. 1 tab. 2. Also works with known partial or extra labels at test time tab. 3 tab. 4.
• The official code is MIT-licensed and mostly matches the paper. Differences to check: crop augmentation, the attention scale, the unknown-state vector, and label masking that is switched off for VG-500 and NEWS-500 in the released loader [see Paper vs code].

Problem

Multi-label classification asks for a set of labels per image. Labels are not independent: some co-occur and some exclude each other §1. Earlier label-dependency methods relied on label graphs built from co-occurrence statistics, and they could not take partial evidence (labels known at test time) into account §5.

The paper also states that prior partial-label inference (feedback-based iterative optimisation) needs iterative updates at test time and was not trained with partial evidence §5 §4.2.

Two inference settings are studied: partial labels (some target labels known) and extra labels (some related labels outside the target set known, such as location or weather) §2.

Method

C-Tran has three parts: the image and label inputs, a Transformer encoder that mixes them, and a per-label readout. Each part is described below. Equation numbers follow the order in the compiled paper §3.1 §3.2 §3.3.

Inputs §3.1. The image goes through a CNN that outputs a feature grid Z with h × w × d entries. Each spatial vector z_i is one image token. For every label i there is a learned label embedding l_i of size d. A state embedding s_i is added to each label embedding to give l̃_i = l_i + s_i eq. 1. The state is one of three: unknown (U), negative (N) or positive (P). The unknown vector is fixed to zero by definition in the paper §3.1; see Paper vs code for what the released code does.

Transformer encoder §3.2. Image tokens and label tokens form one set H. Each layer computes scaled dot-product attention over the whole set eq. 2, averages the value vectors with those weights eq. 3, and passes the result through a two-layer ReLU feed-forward network eq. 4. Residual connections and layer norm wrap each update §3.5. Layers do not share weights. The paper uses L = 3 layers §3.5.

Readout §3.3. After the last layer, each label's output token is passed through its own linear map to a probability: ŷ_i = σ(w_i · l'_i + b_i) eq. 5. Each label therefore has its own classifier weight vector and bias.

Label mask training §3.4. For each training image the method chooses how many labels will be hidden. The number n of unknown labels is drawn between 0.25ℓ and ℓ, where ℓ is the number of labels §3.4. The hidden labels get the unknown state. The other labels keep their ground-truth state, negative or positive, as input. Binary cross-entropy is computed only on the hidden labels §3.4. The paper writes the training objective as an expectation over the known-label distribution eq. 6. Hiding at least a quarter of labels is justified by comparison with typical masked-language-model rates and by the assumption that inputs with more than 75% of labels known are unrealistic §3.4.

Inference §4.2 §4.3. At test time all labels are hidden for regular classification. For partial-label inference, a chosen fraction of labels is given its ground-truth state. For extra-label inference, the known states come from the extra labels. Predictions are one forward pass with no iterative refinement §4.2. The paper thresholds probabilities at 0.5 for the threshold metrics §4.1.

Pseudocode (training step). The steps below follow the paper's description §3.4 and the released loader and loss data_utils.py run_epoch.py. The paper has no algorithm box, so the code is cited instead.

# Label mask training, one image  [§3.4]
# Input:  image (3, H, W); labels (L,) with 0 (absent) or 1 (present)
# Output: loss value for this image
def training_step(image, labels):
    # Draw how many labels stay known, between 0 and 0.75 L  [§3.4]
    num_known = random_int(0, floor(0.75 * L))
    # Hide all other labels at random  [§3.4]
    hidden = random_subset(L, size=L - num_known)
    # Build the ternary state: -1 unknown, 0 negative, 1 positive  [§3.1]
    mask = labels.copy()
    mask[hidden] = -1
    # Predict every label; the model sees the image and the mask  [§3.2]
    logits = forward(image, mask)
    # Binary cross-entropy, counted only where the label was hidden  [§3.4]
    bce = binary_cross_entropy(logits, labels)
    return sum(bce[mask == -1])

Pseudocode (forward pass). The steps below follow the paper's architecture description §3.1 §3.2 §3.3 and the released model CTran.py.

# C-Tran forward pass  [§3.1] [§3.2] [§3.3]
# Input:  image (3, H, W); mask (L,) with -1 unknown, 0 negative, 1 positive
# Output: logits (L,)
def forward(image, mask):
    # Image grid features, flattened to tokens  [§3.5]
    feature_map = resnet101_grid(image)            # (2048, 18, 18)
    image_tokens = flatten_to_tokens(feature_map)  # (324, 2048)
    # One learned embedding per label index  [§3.1]
    label_tokens = label_embedding(arange(L))      # (L, 2048)
    # Add the state embedding of each label  [§3.1] [eq. 1]
    state_ids = mask + 1                           # U=0, N=1, P=2
    label_tokens = label_tokens + state_table[state_ids]
    # Joint sequence, normalised before the first layer  [§3.2]
    tokens = layer_norm(concat(image_tokens, label_tokens))
    # Full self-attention over image and label tokens  [§3.2] [eq. 2-4]
    for layer in transformer_layers:
        tokens = layer(tokens)
    # One logit per label, from that label's own output token  [§3.3]
    label_out = tokens[-L:]                        # (L, 2048)
    return per_label_dot(label_out)                # (L,)

Pseudocode (decision). The steps below follow the paper's evaluation settings §4.1 §4.2.

# Predict the labels of one image  [§4.1] [§4.2]
# Input:  image (3, H, W); known: dict label -> 0 or 1 (may be empty)
# Output: boolean vector (L,)
def predict(image, known):
    # All labels start unknown; known labels get their state  [§4.2]
    mask = full(L, -1)
    for label, value in known.items():
        mask[label] = value
    # One forward pass, no iterative refinement  [§4.2]
    probs = sigmoid(forward(image, mask))
    # Threshold at 0.5 for the threshold metrics  [§4.1]
    return probs >= 0.5

Hyperparameters

Values are as the paper states them unless the "Where" column says otherwise. The released code's defaults are listed in the notes where they differ.

Results

Author-reported. Setup: four multi-label datasets (COCO-80, VG-500, COCO-1000, NEWS-500) and one multi-class extra-label dataset (CUB-312), each with a ResNet-101 backbone except CUB (Inception-v3) §3.5 §4. No standard deviations or repeated-run numbers are reported. The paper does not say whether baseline rows were rerun or taken from their original papers §4.1.

Table 1: regular inference on COCO-80 (%), threshold 0.5 for precision, recall and F1 tab. 1. The paper's best value per column is marked in bold in the source. "All" uses every label above 0.5. "Top 3" uses the three most probable labels per image §4.1.

Headline: mAP 85.1 against 84.3 for KGGR, an absolute gain of 0.8 §4.1. The paper calls mAP, OF1 and CF1 the most important metrics §4.1.

Table 2: regular inference on VG-500 (%), same metrics as Table 1 tab. 2.

Headline: mAP 38.4 against 37.4 for KGGR, an absolute gain of 1.0 §4.1.

Table 3: inference with partial labels, mAP (%) on the unknown labels tab. 3. ε is the share of labels known at test time.

The paper states that improvement over Feedbackprop grows with ε and that C-Tran needs no hyperparameter tuning and one pass at inference §4.2.

Table 4: inference with extra labels, accuracy (%) on CUB-312, 200 target classes tab. 4. ε is the share of extra-label groups (28 groups in total) known.

The Standard, Multi-task and ConceptBottleneck rows are the baselines of Koh et al. as cited in the paper §4.3 tab. 4. The paper states that the multi-task model cannot use known extra labels, which is why its row is flat §4.3. The paper does not explain the flat Standard row.

Table 5: ablation, mAP (%) at 0% and 50% known labels tab. 5.

Further analysis in the paper: with 2, 3 and 4 layers the COCO-80 mAP was 85.0, 85.1 and 84.3 §4.4. The paper estimates a practical upper bound of about 2,000 labels on four Titan X GPUs, which is a limit of the memory cost of the full token set §4.4.

Paper vs code

The official code is at QData/C-Tran, commit 560cf6f47a7021b9993d63c78f61565840786c24, licensed MIT (copyright Jack Lanchantin, 2021) LICENSE. The owner is the research group of co-author Yanjun Qi at the University of Virginia, and the README links the paper README.md. The code is also checked against the paper's architecture and training description; discrepancies are listed below, most significant first.

• Known-label masking is off for VG-500 and NEWS-500 as shipped. The training datasets for VG-500 and NEWS-500 are built with known_labels=0, so every label is hidden in every training image, regardless of --use_lmt load_data.py load_data.py data_utils.py. The paper describes label mask training for all datasets §3.4. The VG-500 and NEWS-500 results in Tables 2, 3 and 5 may therefore have been trained differently from the description; the paper does not say so. Check the flags used for those runs before reproducing them.
• Unknown state vector is not zero. The paper says the unknown state vector is fixed to zero §3.1. The code builds the state table as Embedding(3, d, padding_idx=0), which zeros row 0 at construction. It then applies weights_init to the same table, which overwrites all rows with uniform random values in ±1/√d CTran.py CTran.py utils.py. Row 0 receives no gradient (padding index), so the unknown embedding stays at that random value for the whole run. This is read from the code and not run.
• Augmentation differs. The paper describes a 576 × 576 random crop after resizing to 640 × 640 §3.5. The released training transform picks one crop from five sizes (640, 576, 512, 384, 320) at random, then resizes to 576 load_data.py. The resize to a square 640 × 640 also ignores aspect ratio in both the training and test transforms load_data.py load_data.py.
• Attention scale. The paper divides scores by √d with d = 2048 §3.2. The code uses nn.MultiheadAttention with 4 heads, which divides by the per-head size √512 transformer_layers.py transformer_layers.py. This is the standard multi-head form; the paper's notation does not show it.
• Loss is a sum. The loss adds the per-label binary cross-entropy over the batch and over the hidden labels, with no division by the count run_epoch.py. The paper writes the objective without a normalisation constant eq. 6. Learning rate 1e-5 applies to this summed loss.
• Input and output normalisation. The code applies a LayerNorm to the concatenated image and label tokens before the first layer, and then uses post-norm blocks (normalise after each residual add) CTran.py transformer_layers.py. The paper states only "residual and layer norm" §3.2. The output readout computes a full label-by-label matrix and then keeps only the diagonal CTran.py. The result is the paper's per-label classifier, but the cost is quadratic in the number of labels.
• Optimiser schedule. The paper states no learning-rate schedule and no number of epochs §3.5. The code defaults to 100 epochs and reduces the learning rate by 10× after 5 epochs without improvement in the validation unknown-label loss config_args.py main.py main.py.
• Backbone trained end to end. The paper does not say whether the CNN is frozen §3.5. The code fine-tunes it by default; the optimiser includes all parameters that require gradients main.py.
• Model selection and reporting use one split for COCO-80. For COCO-80 the released loader has a validation set (val_test.data) and no separate test set. The "test" metrics are computed on that same split. The best epoch is chosen by validation mAP, and the test metrics at that epoch are logged as the result load_data.py main.py logger.py. The paper does not describe the selection rule. The reported numbers are therefore selected on the evaluation split.
• Known-label count. The appendix says ε = 25% gives 21 known labels for COCO-80 app. A.1. The released formula gives int(0.25 × 80) = 20 config_args.py. The two disagree by one label.
• Top-3 metric. The paper evaluates "the Top-3 predicted labels" §4.1. The code keeps a label in the top 3 only if its probability is also at least 0.5, so an image with fewer than three confident labels gets fewer than three positives evaluate.py. Ties at the third-highest score are not broken.
• Positional encoding is off by default. The released model has an optional 2-D positional encoding for image tokens, disabled unless --pos_emb is passed CTran.py config_args.py. The paper does not mention it.
• CUB path. CUB uses a separate model class with an Inception-v3 backbone and an extra embedding for the type of each token CTran_cub.py. This path was only skimmed while compiling this card.

Consistent with the paper: Adam with betas (0.9, 0.999) and no weight decay §3.5 main.py (the config's weight_decay default of 1e-4 is applied only to SGD main.py); 3 layers, 4 heads, dropout 0.1 §3.5 config_args.py; the 18 × 18 grid from a 576 px input §3.5 CTran.py; the unknown-count range from 0.25 ℓ to ℓ §3.4 data_utils.py; center crop at test time §3.5 load_data.py.

Pitfalls

• Sum reduction and learning rate. The loss is summed over images and hidden labels run_epoch.py, so its scale grows with batch size and with the number of hidden labels. The paper gives a learning rate of 1e-5 and batch size 16 but no loss normalisation §3.5. When you change the batch size, rescale the learning rate or the loss.
• Two sources of truth for training masks. The mask at training time depends on the dataset class, not only on the --use_lmt flag load_data.py load_data.py. Check the training configuration of each dataset before comparing to the paper's tables.
• Hidden-label count is inclusive and truncated. The known count is drawn uniformly from 0 to int(0.75 ℓ), both ends included data_utils.py. For small ℓ the truncation matters; for ℓ = 80 the upper bound is 60 known labels.
• Test masks are seeded per image. At test time the random choice of known labels is seeded by a hash of the image pixels data_utils.py. The same image always gets the same known labels, but changing the pixels, for example by a different resize, changes the known set. Compare results only with the same preprocessing.
• Known labels are trusted. The paper assumes that partial and extra labels are correct §2 §3.4. It does not test noisy known labels, so the gains in Table 3 are only shown for accurate known labels.
• Quadratic cost in labels. Attention runs over 324 image tokens plus ℓ label tokens §3.2. The paper reports a practical limit near 2,000 labels on four Titan X GPUs, and suggests reducing d or using efficient Transformer variants §4.4. The released code does not implement those variants.
• Full L × L readout. The released readout computes the entire label-by-label output and masks the diagonal CTran.py. This allocates a (batch, L, L) tensor. Replace it with a per-label dot product when L is large.
• Results are single runs. No variance, seeds or repeated runs are reported §4. Differences of a few tenths of a point, such as the 0.8 mAP gain on COCO-80 §4.1, may be within run-to-run noise.
• Selection on evaluation data. The COCO-80 reporting protocol selects the epoch on the reported split (see Paper vs code). Expect small optimistic bias in the reported numbers.
• Layer and label-set claims are narrow. Layer sensitivity is reported only for COCO-80 mAP §4.4. The extra-label result is for CUB-312 only §4.3.
• Metric definitions. The paper uses per-class (C) and overall (O) precision, recall and F1 at threshold 0.5 app. A.3. CF1 and OF1 are the F1 of the averaged precision and recall, not the mean of per-label F1 scores app. A.3.

Lineage

• Builds on: the Transformer encoder and self-attention §3.2; masked prediction as in BERT and the cloze task §3.4; ResNet-101 as the feature extractor §3.5. These works are named in the paper and are not compiled as cards here.
• Compared against: ML-GCN, SSGRL, KGGR, SRN, Multi-Evidence, ML-ZSL, Order-Free RNN, RNN-Attention, CNN-RNN, ResNet-101 baseline tab. 1 tab. 2; Feedbackprop for partial labels tab. 3; Concept Bottleneck Models for extra labels tab. 4.
• Superseded by: not stated in this card.