CapPa: Training vision models as captioners

Introduction

Objective

Captioning: We feed images as inputs and try to predict the caption
Masked prediction: We mask part of the caption and predict all the masked tokens simultaneously

Source: Image Captioners Are Scalable Vision Learners Too, Figure 1

Captioning: 25% of the time
Masking: we drop the entire caption 75% of the time

Dataset

We use DataComp 1B and remove exact duplicates of image + caption.
One image can have multiple captions and a caption can be used for different images.
The validation set is made of images that don’t have any other exact duplicate (regardless of the caption).
During training, we don’t apply any augmentation on images.
Images are resized to 256 x 256 (TPUs like multiples of 128 and it limits storage/egress costs)
About 45% of images have an edge smaller than 256 pixels but we still keep them as a form of data augmentation.
We drop the entire caption 75% of the time (with masking objective instead of captioning).

Model architecture

Vision: 328M
Text: 348M (includes 67M for embeddings)

Patches of 16x16 on images of 256x256 pixels, ie 256 patches per image
Learnt positions
MLP uses GeGLU activations (linear + GELU)
No bias
RMSNorm with Normformer positions (start + end of Attention blocks & start + mid of Feed-forward blocks)
24 layers, hidden dimension of 1024, mlp dimension of 3072
Use of vision registers (see "Vision Transformers Need Registers")
we use 8 as it trains faster on TPU’s vs less
we don’t discard any of the registers as we use cross-attention to text vs CLS token in paper ; in downstream tasks we would use a MAP head keeping all tokens + registers

Decoder (causal when captioning) with cross-attention on images
Same as vision except 12 layers (half)
Max length of 64 tokens
Bias only on final layer predicting vocab

Training

Training Configuration

Hardware: TPU v5e-256
Framework: JAX/Flax
Optimizer: distributed shampoo with RSMProp normalized grafting, based on "Evaluation of Distributed Shampoo"
No weight decay, no dropout
Weights in float32, computation in bfloat16 except for attention logits, normalization layers, and loss
Sharding: FSDP, model and data dimension sharded across all devices
Batch size of 8,192 - Note that 16k also fits well with this sharding strategy
Training speed: 0.45 seconds / batch, 1.6B samples per day

Training metrics (updated live)

The validation loss uses only the captioning objective, which explains why it is much lower than the training loss.
To derive top 1/5 scores on ImageNet:
We use the lowest softmax cross-entropy of each possible class against each image as we cannot compute text embeddings separately like with CLIP models
Since it is much slower to perform, we only use 5% of the data
It is expected to be lower than if we fine-tune a CLIP model (SigLiT style) as it is not trained explicitly for that task and the current method may be affected by different token length for each class
You can explore additional metrics (gradients, parameter norms…) in the main workspace or in individual runs:
first stage: warmup + fixed learning rate
second stage: cosine decay

Sample predictions (updated live)

Training Log

June 20th: started first stage with fixed learning rate 1e-4 and 20k warmup steps
June 25th: started second stage with cosine decay

Results

Visual inspection

the model is good at OCR
it has strong cultural knowledge
some details may be hallucinated
there is likely leakage from near duplicates in training data (see last line of below table)

Image	Ground Truth	Prediction
	kristen coates art and home holiday 01.png	christmas coates holiday shop
	my sheltie loves agility design, sable tshirts	sheltie agility t shirt
	the poetry of living off the grid	the road washes out in spring
	articles of confederation are ratified	the constitution
	how to draw spider man full body	how to draw a man side view
	hotel arenal bromelias, view from the street	hostal la merced
	the prince fielder era came to an end for the tigers over the offseason	prince fielder
	south shore fynn 2 drawer nightstand gray oak, 0	south shore fynn collection nightstand, gray oak

CapPa to SigLiT

Composition understanding - SugarCrepe

Registers sttention maps





output, mod_vars = model.apply(
    {"params": params},
    **batch,
    # NOTE: added to capture intermediate values
    capture_intermediates=True, mutable=["intermediates"]
)

attention_weights = mod_vars["intermediates"]["vision"]["encoder"]["layers"]["attention"]["attention_weights"]

How to use the model

@partial(
    jax.jit,
    static_argnames=("num_beams", "do_sample", "temperature", "top_p", "top_k", "max_length", "num_return_sequences"),
)
def generate_caption(pixel_values, *args, **kwargs):
    return model.generate(pixel_values, *args, **kwargs)

Resources

Public codebase, Clip-JAX Github repo
Model weights on Hugging Face
"Image Captioners Are Scalable Vision Learners Too", Michael Tschannen and Manoj Kumar and Andreas Steiner and Xiaohua Zhai and Neil Houlsby and Lucas Beyer
"Sigmoid Loss for Language Image Pre-Training", Xiaohua Zhai and Basil Mustafa and Alexander Kolesnikov and Lucas Beyer
"Vision Transformers Need Registers", Timothée Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski
"SugarCrepe: Fixing Hackable Benchmarks for Vision-Language Compositionality", Cheng-Yu Hsieh and Jieyu Zhang and Zixian Ma and Aniruddha Kembhavi and Ranjay Krishna
Big Vision Github repo which contains a reference implementation of CapPa

Acknowledgements

Google TPU Research Cloud (TRC) program for providing computing resources
Weights & Biases for the infrastructure for experiment tracking and model management
Hugging Face for the support during the project
Lucas Beyer for clarifications/discussions on "Image Captioners Are Scalable Vision Learners Too"
Tim Darcet for clarifications/discussions on "Vision Transformers Need Registers"