ClipCap: CLIP Prefix for Image Captioning

In plain words

Image captioning means writing a sentence that describes a picture. Many captioners are heavy: they first detect objects, then train a decoder on large datasets for a long time. ClipCap takes a shortcut. A pretrained model called CLIP already turns a picture into a vector that carries its meaning. A small network turns that vector into a few "pseudo-word" vectors, which are not real words, and a pretrained text generator (GPT-2) reads them as if they were the first words of the sentence. Training only teaches the system how to continue from those pseudo-words into a caption.

The surprising part is that both CLIP and the text generator can stay frozen, and only the small translator is trained. That keeps the trainable part small. The other variant lets the text generator change too, which gives a bigger model. Which one scores higher depends on the dataset. Because the translator is the only part that has to learn something new, training takes a few hours to a few days on one GPU, and no extra labels are needed.

TL;DR

• Captions are generated from a CLIP image embedding that is mapped to k prefix vectors and placed before the caption embeddings of GPT-2 §3.1.
• Two variants: an MLP mapper with GPT-2 fine-tuned, and a transformer mapper with CLIP and GPT-2 frozen so only the mapper trains §3.2, §3.3.
• Author-reported: matches or beats Oscar and VLP on nocaps and Conceptual Captions, with far fewer parameters and training hours; Oscar stays best on COCO, where it uses object tags §4 tab. 1.
• Worth reading if you want a cheap captioner built from frozen pretrained models; check the Pitfalls first, because the paper and code differ in several details.

Problem

Captioners that reach strong scores need a detector or extra object tags, large datasets, and long training runs. The authors argue that this makes them hard to retrain often or to adapt for several datasets §1. They want a lightweight captioner that needs only images and captions §1.

Method

Training objective. Given N image-caption pairs, the captioner maximises the likelihood of each caption, token by token, conditioned on the image eq. 1 eq. 2. The image is first encoded by the visual encoder of a pretrained CLIP model. A mapping network F turns that embedding into k vectors, each with the size of a word embedding §3.1 eq. 3. These vectors are placed in front of the caption's word embeddings to form one sequence §3.1 eq. 4. The language model is trained on that sequence, and the loss is the cross-entropy over the caption tokens only, conditioned on the prefix and the earlier caption tokens §3.1 eq. 5.

\mathcal{L}_{X} = -\sum_{i=1}^{N}\sum_{j=1}^{\ell} \log p_\theta\left(c^i_j \mid p^i_1, \ldots, p^i_k, c^i_1, \ldots, c^i_{j-1}\right) eq. 5

Variant A: GPT-2 fine-tuned, MLP mapper §3.2 §3.3. The language model is trained together with F. The authors use a single hidden layer MLP, because with the language model fine-tuned the mapping is less demanding §3.3. Fine-tuning CLIP was tried and did not help, so CLIP stays frozen §3.2.

Variant B: GPT-2 frozen, transformer mapper §3.2 §3.3. Only F is trained, so the language model is frozen. The mapper is a transformer that receives the CLIP embedding and a set of learned constant inputs. The constants pull information out of the CLIP embedding through attention and also adapt the frozen language model to the new data §3.3. The paper does not say which transformer outputs form the prefix; the code takes the outputs at the constant positions train.py.

Inference §3.4. Compute the prefix from the image, then generate the caption one token at a time from the language model's output. Choose each token greedily or with beam search §3.4.

The pseudocode below follows the training objective §3.1 eq. 3-5. The code computes the same step with a batch-wise mean instead of the paper's sum train.py.

# One training step, ClipCap  [§3.1] [eq. 3-5]
# Input:  images (B, 3, H, W); caption_tokens (B, L) of GPT-2 ids
# Output: scalar loss

def training_step(images, caption_tokens):
    # Encode each image with the frozen CLIP encoder  [§3.1]
    clip_embedding = clip_image_encoder(images)     # (B, 512)
    # Map the CLIP embedding to k prefix vectors  [§3.1] [eq. 3]
    prefix = mapping_network(clip_embedding)        # (B, k, d_word)
    # Embed the caption tokens with the language model's embedding table
    caption_embedding = language_model.embed(caption_tokens)  # (B, L, d_word)
    # Put the prefix in front of the caption embeddings  [§3.1] [eq. 4]
    inputs = concat(prefix, caption_embedding, axis=1)
    # Run the language model over the whole sequence  [§3.1]
    logits = language_model(inputs)                 # (B, k+L, vocab)
    # Keep only positions that predict caption tokens  [§3.1] [eq. 5]
    # The last prefix position predicts the first caption token
    caption_logits = logits[:, k - 1 : -1]          # (B, L, vocab)
    # Cross-entropy over caption tokens only  [§3.1] [eq. 5]
    return cross_entropy(caption_logits, caption_tokens)

Mapper, MLP variant §3.3 train.py train.py.

# MLP mapping network, variant A  [§3.3]
# Input:  clip_embedding (B, 512)
# Output: prefix (B, k, d_word)

def mlp_mapper(clip_embedding):
    # One hidden layer; the code makes it half the output width  [§3.3]
    hidden = tanh(linear_in(clip_embedding))        # (B, k * d_word / 2)
    # Project to k word-sized vectors  [§3.3]
    prefix = linear_out(hidden)                     # (B, k * d_word)
    return reshape(prefix, (B, k, d_word))

Mapper, transformer variant §3.3 train.py.

# Transformer mapping network, variant B  [§3.3]
# Input:  clip_embedding (B, 512)
# Output: prefix (B, k, d_word)

def transformer_mapper(clip_embedding):
    # Project the CLIP vector to c tokens  [code: train.py#L206]
    clip_tokens = reshape(linear(clip_embedding), (B, c, d_word))
    # Learned constant inputs, shared across the batch  [§3.3]
    constants = tile(learned_constants, B)          # (B, k, d_word)
    # Self-attention over clip tokens and constants  [§3.3]
    # 8 layers, 8 heads  [§4 Implementation details]
    sequence = concat(clip_tokens, constants, axis=1)
    sequence = transformer_layers(sequence)
    # Keep the outputs at the constant positions as the prefix  [code]
    return sequence[:, c:]                          # (B, k, d_word)

Decoding, greedy §3.4 predict.py. The code's generate2 filters with top-p and then takes the argmax, so the chosen token is always the most likely one. This makes it greedy decoding.

# Greedy caption decoding  [§3.4] [code: predict.py#L240-L302]
# Input:  prefix (1, k, d_word); max_tokens (67 in code, not stated in paper)
# Output: caption token ids

def greedy_decode(prefix, max_tokens):
    # Start from the prefix alone  [§3.4]
    embeddings = prefix
    caption = []
    for step in range(max_tokens):
        # Score the next token from the last position  [§3.4]
        logits = language_model(embeddings)[:, -1]  # (1, vocab)
        # Pick the most likely token  [§3.4] [code: predict.py#L288]
        next_token = argmax(logits)
        caption.append(next_token)
        # Stop after the period token  [code: predict.py#L295-L296]
        if next_token == period_token:
            break
        # Append the chosen token's embedding and continue  [§3.4]
        embeddings = concat(embeddings, embed(next_token), axis=1)
    return caption

Beam search is the other option §3.4. The code's generate_beam keeps 5 candidates by default and ranks them by length-normalised log-probability predict.py. The paper gives no beam width.

Hyperparameters

Results

Author-reported. Setup: COCO-captions, nocaps and Conceptual Captions, with the baselines BUTD, VLP and Oscar §4 Datasets §4 Baselines. Metrics: BLEU-4, METEOR, CIDEr and SPICE on COCO; CIDEr and SPICE on nocaps; ROUGE-L, CIDEr and SPICE on Conceptual Captions §4 Evaluation metrics. All numbers below come from tab. 1. Training times are GPU hours, and "Ours" variants are the paper's own. Oscar uses object tags §4 Baselines.

(A) Conceptual Captions tab. 1 (A)

(B) nocaps tab. 1 (B)

(C) COCO tab. 1 (C)

(D) Ablation tab. 1 (D)

Qualitative results and the prefix-length curves are in §4 and are not reproduced in numbers here §4 Qualitative evaluation §4 Prefix length.

Paper vs code

Code: https://github.com/rmokady/CLIP_prefix_caption at commit 1ad805a844a62ab2e5480479aa021bccf0d4d12a, MIT license. The repository is the authors' own; the paper's abstract links it.

• Language model size. Paper says GPT-2 (large) §3.1. Code loads gpt2 and uses 768-dim embeddings train.py predict.py. The 156M reported for the MLP variant matches GPT-2 small plus the MLP mapper, by our arithmetic from the code's shapes: about 124M + 31.5M ≈ 156M. This is our calculation, not stated by the paper.
• Transformer prefix length. Paper sets K=10 for the transformer §4 Implementation details. The README's transformer command uses --prefix_length 40 --prefix_length_clip 40 README.md README.md. The paper's ablation goes up to 80 §4 Prefix length. The paper does not say which length produced tab. 1 (A)-(D).
• Loss normalisation. Eq. 5 sums over samples and positions eq. 5. Code uses PyTorch cross-entropy with its default mean reduction over non-ignored tokens train.py.
• Learning rate schedule. Paper gives 2e-5 and 5000 warmup steps §4 Implementation details. Code decays the rate linearly to zero over all steps train.py. The paper does not state a decay.
• Weight decay. Paper says "weight decay fix" §4 Implementation details. Code passes no weight-decay argument train.py.
• MLP architecture. Paper says single hidden layer §3.3. Code uses tanh with hidden width half the output width train.py train.py. The paper does not state the width or activation.
• Inference MLP for long prefixes. train.py always builds the MLP mapper when mapping_type is mlp train.py. predict.py uses a single linear layer when prefix_length > 10 predict.py. Checkpoints trained with long MLP prefixes therefore won't load into predict.py without a change.
• Transformer layer details. Paper gives 8 layers and 8 heads §4 Implementation details; code matches train.py train.py. Code also sets the MLP ratio to 2.0 train.py and uses no bias in attention projections train.py; the paper says nothing about either.
• Captions from Conceptual Captions. Code appends a period to every caption that does not already end in one parse_conceptual.py. The paper does not mention this.
• Prefix normalisation. Code has an optional prefix-normalisation flag, off by default train.py. The paper does not mention it.
• Dropout during frozen-LM training. Code puts the frozen language model in eval mode during training, so its dropout is off train.py. The paper does not mention dropout.

Pitfalls

• Pad id 0 is also a real token. The loss ignores id 0 train.py, and padding is set to id 0 train.py. A genuine id 0 in a caption is therefore dropped from the loss too, and the code cannot tell the two apart.
• Captions are truncated. The maximum length is the mean plus 10 standard deviations, capped at the longest caption train.py. Longer captions are cut train.py. The paper does not state a maximum length.
• Decoding is greedy in the code. The top-p filter at 0.8 predict.py keeps the most likely token, so the argmax is unchanged predict.py. Beam search runs only when asked predict.py. Decoding stops at the first "." predict.py or after 67 tokens predict.py. The paper names neither limit.
• Prefix interpretation depends on training. Paper reports that the nearest-word interpretation of the prefix is meaningful only when GPT-2 is also trained §4 Prefix interpretability. With frozen GPT-2 the prefix is mostly unreadable, and the authors say it is shared across images §4 Prefix interpretability. Do not read the frozen-LM prefix as a caption.
• Transformer prefix length is costly. Attention memory is quadratic in prefix length, which the authors give as the reason for stopping at 80 §4 Prefix length. The MLP variant at prefix 40 has over 450M parameters §4 Prefix length.
• Warmup share depends on run length. The code decays the rate linearly to zero over all steps train.py, while the paper gives only the 5000 warmup steps §4 Implementation details. Changing the epoch count or batch size changes how much of training the warmup covers, so recompute the schedule when you change either.
• README flag typo. The README transformer commands pass --num_layres README.md README.md, but the parser defines --num_layers train.py. Argparse will reject the misspelt flag, so fix it before running those commands.
• GPT-2 size is unresolved. The paper says "large" §3.1, while the code loads gpt2 train.py. The reported 156M matches the smaller model in our arithmetic (see Paper vs code), so follow the code if you want to match the table.
• Speed numbers are not comparable. Oscar and VLP training times are on V100 and exclude pre-training for VLP §4 Quantitative evaluation tab. 1. The paper's own times are on GTX 1080.

Lineage

• Builds on: CLIP, arXiv:2103.00020, used as the frozen image encoder §3.1 [bib].
• Builds on: prefix-tuning, arXiv:2101.00190, which inspired training a prefix for a frozen language model §1 §3.2 [bib].
• Builds on: GPT-2 as the language model §1 §3.1.
• Compared against: VLP (Zhou et al.; the paper's bibliography gives no arXiv id), Oscar, and BUTD, with Oscar using object tags §4 Baselines.
• Superseded by: not stated in this wiki.