Align before Fuse: Vision and Language Representation Learning with Momentum Distillation

In plain words

Many systems that answer questions about pictures first look at the image and the text separately, then mix the two together in one big network. The mixing step has a hard time, because the picture's parts and the words are not yet speaking the same language. ALBEF changes the order: it first trains the picture side and the text side so that a matching image and caption land close together, and only then lets them interact in a joint network.

The web is full of image captions that are loose or wrong, so training on them can teach the model the wrong lesson. ALBEF adds a second, slowly updated copy of itself that acts as a teacher. Instead of insisting that the caption is the only correct word, the teacher says which other words or captions would also fit, and the model learns from that softer answer.

The result is a model that needs no object detector and no boxes, works on smaller images, and answers visual questions, checks whether a sentence matches a picture, and finds the described region of an image with good accuracy. The paper reports that it is also faster at inference than several earlier models.

TL;DR

ALBEF pre-trains a vision transformer and a BERT-based text encoder with a contrastive loss that aligns their outputs, then fuses them with cross-modal attention. Momentum distillation trains each part against a slowly moving teacher, which helps with noisy web captions. It needs no bounding boxes, and it is reported to lead on retrieval, VQA, NLVR2, visual entailment and weakly-supervised grounding. All numbers are author-reported. Use it as a reference design for the align-then-fuse pattern and for teacher-based soft targets.

Problem

Earlier vision-language pre-training feeds region features from a pre-trained object detector and word tokens into one multimodal transformer. The two sets of tokens live in separate spaces, so the fusion encoder must learn their interaction from scratch §1. The detector needs box annotations and high-resolution inputs (for example around 600×1000 at inference) §1. Web image-text data is noisy, and masked-language-modelling on it can overfit to the caption text §1.

The paper assumes large image-text corpora (4M to 14M images here) and fine-tuning on labelled downstream sets §3.4 §5.

Method

ALBEF has three parts: an image encoder, a text encoder, and a multimodal encoder that fuses the two §3.1.

1. Image encoder. A 12-layer ViT-B/16. It is initialised from DeiT weights pre-trained on ImageNet-1k §3.1. Its [CLS] embedding is projected to 256 dimensions and L2-normalised for the contrastive loss §3.2.
2. Text encoder. The first 6 layers of BERT-base. It produces token embeddings and a [CLS] embedding §3.1.
3. Multimodal encoder. The last 6 layers of BERT-base, with cross-attention to the image tokens in every layer §3.1. It is fed the text-encoder output.
4. Image-text contrastive loss (ITC), before fusion. Projected [CLS] features are scored with a dot product and a learnable temperature τ §3.2. Scores are computed against the batch plus a queue of the most recent momentum-encoder features (queue size M) §3.2.

p_m^{\mathrm{i2t}}(I) = \frac{\exp(s(I,T_m)/\tau)}{\sum_{m=1}^{M}\exp(s(I,T_m)/\tau)}, \quad p_m^{\mathrm{t2i}}(T) = \frac{\exp(s(T,I_m)/\tau)}{\sum_{m=1}^{M}\exp(s(T,I_m)/\tau)} \quad \text{[eq. 1]}

\mathcal{L}_{\mathrm{itc}} = \tfrac{1}{2}\,\mathbb{E}_{(I,T)\sim D}\big[\mathrm{H}(\mathbf{y}^{\mathrm{i2t}}(I),\mathbf{p}^{\mathrm{i2t}}(I)) + \mathrm{H}(\mathbf{y}^{\mathrm{t2i}}(T),\mathbf{p}^{\mathrm{t2i}}(T))\big] \quad \text{[eq. 2]}

5. Masked language modelling (MLM) on the fused model. 15% of text tokens are selected. Of those, 80% become [MASK], 10% a random token and 10% stay unchanged, as in BERT §3.2 (footnote). The cross-entropy target is the one-hot vocabulary label eq. 3.

\mathcal{L}_{\mathrm{mlm}} = \mathbb{E}_{(I,\hat{T})\sim D}\,\mathrm{H}(\mathbf{y}^{\mathrm{msk}}, \mathbf{p}^{\mathrm{msk}}(I,\hat{T})) \quad \text{[eq. 3]}

6. Image-text matching (ITM) with hard negatives. The [CLS] output of the fusion encoder goes to a two-way linear head (match / no match) §3.2. Per image, one negative text is drawn from the batch with probability given by the ITC similarity, so similar texts are more likely; likewise one negative image per text §3.2.

\mathcal{L}_{\mathrm{itm}} = \mathbb{E}_{(I,T)\sim D}\,\mathrm{H}(\mathbf{y}^{\mathrm{itm}}, \mathbf{p}^{\mathrm{itm}}(I,T)) \quad \text{[eq. 4]}

7. Pre-training objective. The three losses are summed with weight 1 each eq. 5.

\mathcal{L} = \mathcal{L}_{\mathrm{itc}} + \mathcal{L}_{\mathrm{mlm}} + \mathcal{L}_{\mathrm{itm}} \quad \text{[eq. 5]}

8. Momentum model. A teacher with the same architecture is kept as an exponential moving average (EMA) of the online encoders, with momentum 0.995 §3.3 §3.5. It is not trained by gradients.
9. Momentum distillation (MoD) for ITC. The teacher's similarities give soft pseudo-targets q. The ITC target is a mix: weight α on the teacher distribution and 1 − α on the one-hot label §3.3 eq. 6.

\mathcal{L}^{\mathrm{mod}}_{\mathrm{itc}} = (1-\alpha)\,\mathcal{L}_{\mathrm{itc}} + \tfrac{\alpha}{2}\,\mathbb{E}_{(I,T)\sim D}\big[\mathrm{KL}(\mathbf{q}^{\mathrm{i2t}}(I)\,\|\,\mathbf{p}^{\mathrm{i2t}}(I)) + \mathrm{KL}(\mathbf{q}^{\mathrm{t2i}}(T)\,\|\,\mathbf{p}^{\mathrm{t2i}}(T))\big] \quad \text{[eq. 6]}

10. Momentum distillation for MLM. Same idea: the MLM loss mixes the one-hot label with the teacher's distribution over the vocabulary at each masked position §3.3 eq. 7.

\mathcal{L}^{\mathrm{mod}}_{\mathrm{mlm}} = (1-\alpha)\,\mathcal{L}_{\mathrm{mlm}} + \alpha\,\mathbb{E}_{(I,\hat{T})\sim D}\,\mathrm{KL}(\mathbf{q}^{\mathrm{msk}}(I,\hat{T})\,\|\,\mathbf{p}^{\mathrm{msk}}(I,\hat{T})) \quad \text{[eq. 7]}

11. Schedule of α. α ramps linearly from 0 to 0.4 during the first epoch of pre-training §3.5.

The pseudocode below follows the steps above. The paper has no algorithm box, so it is composed from §3.1–§3.3 and eq. 1–eq. 7.

# Pre-training step for one batch of B image-text pairs  [§3.1]–[§3.3]
# Input:  images (B, 3, H, W); captions (B, L) as token ids
# Output: scalar loss = L_itc^mod + L_mlm^mod + L_itm  [eq. 5]

def training_step(images, captions, alpha):
    # Unimodal encoders on the online model  [§3.1]
    image_embeds = vit(images)                        # (B, 1 + N, 768)
    text_embeds = bert_text(captions)                 # (B, L, 768)
    image_feat = normalize(proj_v(image_embeds[:, 0]))  # (B, 256)
    text_feat = normalize(proj_w(text_embeds[:, 0]))    # (B, 256)

    # Momentum (EMA) teacher on the same inputs, without gradients  [§3.3]
    image_feat_m = normalize(proj_v_m(vit_m(images)[:, 0]))  # (B, 256)
    text_feat_m = normalize(proj_w_m(bert_text_m(captions)[:, 0]))

    # Queues hold the most recent momentum features, size M  [§3.2]
    texts_all = concat(text_feat_m, text_queue)       # (B + M, 256)
    images_all = concat(image_feat_m, image_queue)    # (B + M, 256)

    # Contrastive scores, online against momentum  [§3.2] [eq. 1]
    sim_i2t = image_feat @ texts_all.T / tau          # (B, B + M)
    sim_t2i = text_feat @ images_all.T / tau          # (B, B + M)

    # Teacher scores on the same grid  [§3.3]
    sim_i2t_m = image_feat_m @ texts_all.T / tau      # (B, B + M)
    sim_t2i_m = text_feat_m @ images_all.T / tau      # (B, B + M)

    # Soft targets: teacher mass on alpha, one-hot on 1 - alpha  [§3.3] [eq. 6]
    hard = eye_padded(B, B + M)                       # positive is the diagonal
    targets_i2t = alpha * softmax(sim_i2t_m) + (1 - alpha) * hard
    targets_t2i = alpha * softmax(sim_t2i_m) + (1 - alpha) * hard

    # Soft cross-entropy, averaged over the two directions  [§3.2] [eq. 2]
    loss_itc = 0.5 * (
        soft_ce(sim_i2t, targets_i2t) + soft_ce(sim_t2i, targets_t2i)
    )

    # Push the momentum features into the queues  [§3.2]
    enqueue(image_feat_m, text_feat_m)

    # Fused pass on the positive pairs  [§3.1] [§3.2]
    pos_out = fusion(text_embeds, image_embeds, captions)     # (B, 768)

    # Hard negatives: sample from the in-batch ITC scores, no extra encoder
    # pass for mining  [§3.2]
    w_i2t = softmax(sim_i2t[:, :B]) with diagonal set to 0    # (B, B)
    w_t2i = softmax(sim_t2i[:, :B]) with diagonal set to 0    # (B, B)
    neg_text = sample_rows(w_i2t)                     # one per image
    neg_image = sample_rows(w_t2i)                    # one per text

    # Fused pass on 2B negative pairs  [§3.2] [eq. 4]
    neg_out = fusion_pairs(neg_image, neg_text)       # (2B, 768)
    logits = itm_head(concat(pos_out, neg_out))       # (3B, 2)
    labels = [1] * B + [0] * (2 * B)
    loss_itm = cross_entropy(logits, labels)

    # Mask 15% of tokens; teacher distribution over the vocabulary  [§3.2]
    masked, mlm_labels = mask_tokens(captions, p=0.15)
    teacher_logits = fusion_m(masked, image_embeds_m)  # (B, L, V)
    online_logits = fusion(masked, image_embeds)       # (B, L, V)
    loss_mlm = (1 - alpha) * ce(mlm_labels, online_logits)
    loss_mlm += alpha * soft_ce(softmax(teacher_logits), online_logits)

    return loss_itc + loss_mlm + loss_itm

Hyperparameters

Values are from the paper unless the "Where" column says code. Items marked not stated are missing in the paper.

Results

Author-reported. Pre-training: 4M images (ALBEF) or 14M (ALBEF with Conceptual 12M) §3.4. Baselines are the methods named in each table. The paper reports gains from 4M to 14M images on retrieval and the downstream tasks §6.1 §6.2.

Headline comparisons §6.2 §6.3

• Zero-shot Flickr30K text retrieval R@1: 94.1 (ALBEF 14M) vs 88.6 (ALIGN, 1.2B images) tab. 3.
• VQA test-std: 76.04 (ALBEF 14M) vs 73.67 (VILLA); absolute gain 2.37 tab. 4 §6.3.
• NLVR² test-P: 83.14 (ALBEF 14M) vs 79.30 (VILLA); absolute gain 3.84 tab. 4 §6.3.
• SNLI-VE test: 80.91 (ALBEF 14M) vs 79.03 (VILLA); absolute gain 1.88 tab. 4 §6.3 (the paper cautions that the SNLI-VE test set is noisy §5 fn.).
• Inference: more than 10 times faster than VILLA on NLVR² §6.3. The paper attributes this to being detector-free and using lower-resolution images §6.3. Speed figures are not given as absolute numbers in the paper.

Table 1. Effect of the proposed components (4M images unless noted) tab. 1

Caption: TR and IR are the average of R@1, R@5 and R@10 on Flickr30K; the others are accuracy or score on the stated test split. Source: tab. 1.

Table 2. Fine-tuned image-text retrieval, Flickr30K and COCO tab. 2

Caption: as printed in tab. 2. Recall at 1, 5 and 10. "–" means not reported.

Table 3. Zero-shot image-text retrieval, Flickr30K tab. 3

Caption: as printed in tab. 3. The zero-shot model is fine-tuned on COCO and evaluated on Flickr30K §5.

Table 4. VQA, NLVR², and SNLI-VE tab. 4

Caption: as printed in tab. 4. Dashes mean not reported.

Table 5. Weakly-supervised visual grounding, RefCOCO+ tab. 5

Caption: as printed in tab. 5. Metric is accuracy on the RefCOCO+ split. ALBEF_itc uses the last-layer self-attention; ALBEF_itm uses the 3rd cross-attention layer §6.4 app. A.

Table 6. Retrieval ablation: top-k and hard negatives (Flickr30K, 4M) tab. 6

Caption: as printed in tab. 6. Average recall on the test set. s_itc means ranking by the ITC score alone.

Table 7. NLVR² ablation: text-assignment pre-training and parameter sharing tab. 7

Caption: as printed in tab. 7. "share all" shares all parameters of the two blocks, "share CA" shares only the cross-attention layers, "no share" shares none.

Other results: the paper's own Grad-CAM heat maps for grounding fig. 4, VQA fig. 5 and per-word COCO maps fig. 6; figures are described, not reproduced, here.

Table 8. Pre-training data statistics app. E tab. 8

Paper vs code

Code: official repository salesforce/ALBEF at b9727e43c3040491774d1b22cc27718aa7772fac, linked from the abstract, BSD-3-Clause. I read the pre-training model, training loop, schedulers, dataset transforms and the fine-tuning configs.

Agrees with the paper: loss sum with weight 1 Pretrain.py; EMA momentum 0.995 and queue 65,536 Pretrain.yaml; MLM 15% with 80/10/10 model_pretrain.py; ITM with one in-batch hard negative per image and per text model_pretrain.py; multi-positive ITC with weight 1/#positives in fine-tuning model_retrieval.py; α = 0.4 downstream VQA.yaml; AdamW, weight decay 0.02, 30 epochs Pretrain.yaml.

Differs or is not in the paper:

1. Warm-up length. Paper: warm up to 1e-4 in the first 1000 iterations §3.5. Code: warmup_epochs: 20 with a hard-coded 100-iteration step, so the warm-up runs over 2000 iterations and starts at 1e-5 Pretrain.yaml Pretrain.py Pretrain.py. This is my reading of the scheduler code; check it on a run.
2. Temperature. Paper: learnable, initial value not stated §3.2. Code: learnable, starts at 0.07 and is clamped to [0.001, 0.5] at every forward pass model_pretrain.py model_pretrain.py. The clamp is not in the paper.
3. Colour augmentation. Paper footnote: colour changes were removed from RandAugment because captions mention colours §3.5. Code: the list still contains AutoContrast, Equalize, Brightness and Sharpness [code: dataset/init.py#L21-L22]. The paper's claim and the released code disagree; pick one and record it.
4. Crop scale. Paper: random crops at 256 px, no scale given §3.5. Code: RandomResizedCrop with scale (0.2, 1.0) for pre-training and (0.5, 1.0) for fine-tuning [code: dataset/init.py#L19] [code: dataset/init.py#L27].
5. Hard-negative pool. Paper: negatives are sampled "from the same batch" §3.2. Code: negatives come from the local per-GPU batch (64 in the config), not the global batch of 512 model_pretrain.py Pretrain.yaml.
6. ITC loss form. Code writes the mixed target as α·CE(q, p) + (1 − α)·CE(y, p) model_pretrain.py model_pretrain.py. This equals Eq. 6 up to a constant (the entropy of q), so gradients match but logged loss values are not comparable with the paper's.
7. COCO retrieval epochs. Paper: 10 epochs for both COCO and Flickr30K app. A. Code: 5 epochs for COCO Retrieval_coco.yaml; Flickr30K config has 10 Retrieval_flickr.yaml.
8. Fine-tuning warm-up. Paper: the linear α ramp is described for pre-training §3.5. Code: the same ramp runs during the first fine-tuning epoch for retrieval, VQA, NLVR² and grounding, and is off for visual entailment Retrieval.py VE.yaml. Paper silent on this.
9. Downstream warm-up length. Paper: not stated. Code: 1 epoch for most tasks and 4 for VQA VQA.yaml.
10. Top-k for retrieval inference. Paper main text: not stated. Code: 128 for Flickr30K and 256 for COCO Retrieval_flickr.yaml Retrieval_coco.yaml.
11. Batch sizes. Paper gives total batch sizes app. A; code gives per-GPU sizes (32 for retrieval, VQA, VE; 16 for NLVR²). The totals are consistent with 8 GPUs, but the GPU count for fine-tuning is not stated in the paper.

Pitfalls

• Queue assertion. The queue size must be divisible by the gathered batch size or the enqueue step fails model_pretrain.py. The queue is random normalised noise at start model_pretrain.py; it is not stated in the paper, and early steps contrast against that noise.
• Distributed setup is required. The enqueue step gathers features with torch.distributed even for a single device model_pretrain.py. A one-GPU run needs a process group.
• Memory. The EMA copy doubles weight memory: the momentum ViT, BERT and projection heads run in a second forward pass each step model_pretrain.py. The paper's "zero computational overhead" refers to the mining step only §3.2. ITM still runs the fusion encoder on three pairs per image (one positive, two negatives) model_pretrain.py; MLM runs a second fusion pass through the teacher model_pretrain.py.
• Masking. Only padding and [CLS] are protected from masking, so [SEP] can be masked model_pretrain.py. Not stated in the paper.
• Unstated caption handling. Captions are cleaned, lower-cased, truncated to 30 words utils.py and tokenised to at most 25 tokens Pretrain.py. The paper states neither; the 25-token cap can cut longer captions.
• Warm-up and temperature are the two numerically sensitive knobs. Warm-up length changes how long the model trains at low learning rate; the temperature clamp bounds the logit scale. Both are above under Paper vs code.
• Teacher quality is not measured. The paper shows top-5 teacher candidates as examples fig. 2 and does not quantify teacher error. α = 0.4 was chosen "for simplicity", with 0.3 to 0.5 reported as similar §3.3 fn.
• Data availability. Some Conceptual Captions and SBU image URLs are invalid §3.4 fn., so a full reproduction of the 4M or 14M set is not possible from the paper alone.
• Test leakage in grounding. The authors admit RefCOCO+ val and test images may appear in the 14M pre-training data, and they leave decontamination to future work app. A.
• Noisy SNLI-VE test. The paper cautions that its test data has been reported as noisy §5 fn.
• Speed claim. The speed comparison is qualitative (more than 10 times faster than VILLA on NLVR² §6.3); hardware and batch settings for the timing are not given in the text.

Lineage

• Builds on: ViT-B/16 image encoder with DeiT ImageNet-1k weights §3.1; BERT-base text and fusion encoders §3.1; MoCo-style queues and momentum encoder §3.2; InfoNCE and its mutual-information view §4. These works are named in the paper and no cards exist for them in this repository.
• Replaces: detector-based vision-language pre-training (LXMERT, UNITER, OSCAR) that feeds region features into a fusion encoder §1.
• Compared against: UNITER, OSCAR, VILLA, LXMERT, VisualBERT, VL-BERT, 12-in-1, ViLT, VL-BART/T5, CLIP, ALIGN, ARN and CCL tab. 2–tab. 5.
• Superseded by: not stated in the paper's own content.