OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework

In plain words

Most systems that handle pictures and text are built for one job: one model writes captions, another answers questions about a picture, a third draws images. OFA tries to fold all of these into one model. It reads a short written instruction together with an input, which can be a picture, some text, or both, and it writes its answer as a sequence of tokens.

Words, picture fragments and box coordinates all share one vocabulary, so the same architecture and training recipe covers captioning, visual question answering, pointing at a region from a description, drawing an image, and plain language tasks. The instruction tells the model which job to do, so new tasks do not need new output layers.

The authors report strong results from a comparatively small pretraining set of public image-text pairs. Whether a single model beats specialised ones depends on the task, and the numbers on this card are the authors' own.

TL;DR

• One encoder-decoder Transformer for vision, language and cross-modal tasks. Every task, pretraining included, is generation from an instruction plus inputs §1 §3.2.
• Text tokens, discrete image codes and location (box) tokens share one vocabulary §3.1.
• Pretrained on 20M public image-text pairs plus unimodal image and text data §1 §3.3. The pretraining set is much smaller than the 1.8B pairs of SimVLM tab. 2 caption.
• Reports the best VQA test-std (82.0), SNLI-VE test (91.2) and COCO CIDEr-optimised captioning (154.9) among the models it lists tab. 2 tab. 3. Results are author-reported.
• Unimodal results are competitive but not always best: ImageNet-1K top-1 is 85.6 against BEiT-384-L's 86.3 tab. 8.

Problem

Pretrained vision-language models often need task-specific parts: heads, adapters or soft prompts. Their pretraining and finetuning objectives also differ, and many rely on an object detector, which fails on open-domain images §1. The paper asks for a model that is task-agnostic (one task format for pretraining and finetuning), modality-agnostic (one input and output format for all modalities) and task-comprehensive (enough task variety to generalise) §1.

The paper argues that existing methods fail on at least one of these: they add learnable components for finetuning, use task-specific formulations, or entangle modality representation with detected objects §1. The headline efficiency comparison is with SimVLM, which uses 1.8B image-text pairs tab. 2 caption.

Method

Everything is a sequence-to-sequence problem. The input is an image (optional), an instruction and optional text. The output is a sequence over one shared vocabulary §3.1 §3.2. The architecture is a Transformer encoder-decoder used for every pretraining, finetuning and zero-shot task §3.1.

Architecture §3.1

1. Image input: the first three ResNet blocks turn the image into patch features, following CoAtNet and SimVLM §3.1 app. A.2. Patch size is 16 app. A.2.
2. Text input: byte-pair encoding (BPE), as in GPT and BART §3.1.
3. Encoder and decoder are both stacks of Transformer layers. The decoder adds cross-attention to the encoder output §3.1.
4. Stability changes: head scaling in self-attention, a post-attention LayerNorm, and a LayerNorm after the first FFN layer, following NormFormer §3.1. The released pretraining scripts enable all three (--scale-attn --scale-fc --scale-heads) pretrain_ofa_large.sh.
5. Positions: two absolute position embeddings (text and image), decoupled from token and patch embeddings as in TUPE. Relative position bias is 1D for text and 2D for images §3.1.
6. Embeddings and output softmax share parameters app. A.2.

Shared vocabulary §3.1

• Text: BPE subwords.
• Images: discrete image codes from an image quantiser (VQGAN-style). A 256×256 image becomes a 16×16 code sequence §3.1.
• Boxes: each corner coordinate is uniformly discretised into integer location tokens. Object labels are plain BPE §3.1. The repository uses 1000 bins pretrain_ofa_large.sh and rounds the normalised coordinate times (bins − 1) unify_dataset.py.

Pseudocode of the input side

# Image side: the first three ResNet blocks turn a picture into patches  [§3.1]
# [app. A.2]
def image_patches(image):
    # Output: (P, d) patch features, where d is the hidden size
    feature_map = resnet_first_three_blocks(image)
    return flatten_to_patches(feature_map)

# Text side: byte-pair encoding into subword ids, then embedding  [§3.1]
def text_tokens(text):
    # Output: ids from the shared BPE vocabulary
    return bpe_encode(text)

# Boxes become location tokens; labels stay BPE  [§3.1]
def box_tokens(x1, y1, x2, y2, num_bins):
    # Coordinates are normalised to [0, 1] before discretisation
    # [code: raw/code/2202.03052/data/pretrain_data/unify_dataset.py#L303-L306]
    bins = [round(c * (num_bins - 1)) for c in (x1, y1, x2, y2)]
    return [location_token(b) for b in bins]

# Target-side image: a quantiser maps the picture to discrete codes  [§3.1]
def image_codes(image):
    # A 256 x 256 image becomes a 16 x 16 sequence of code ids  [§3.1]
    return quantiser_encode(image)

Training objective §3.4

The loss is the negative log-likelihood of the target tokens given the image x, the instruction s and earlier target tokens. The paper writes it inline and does not number it.

# Next-token cross-entropy over the target sequence  [§3.4]
def training_loss(model, source_tokens, target_tokens):
    # Teacher forcing: the decoder reads the target shifted right
    logits = model(source_tokens, shift_right(target_tokens))
    # Sum of -log P(y_i | y_<i, x, s) over all target positions  [§3.4]
    return cross_entropy_sum(logits, target_tokens)

Inference §3.4 app. B

# Beam search over the shared vocabulary  [§3.4] [app. A.3]
def generate(model, source_tokens, beam_size, candidate_trie=None):
    beams = [start_beam()]
    for step in range(max_length):
        # Score every next token for each live beam  [§3.4]
        log_probs = model.next_token_log_probs(beams, source_tokens)
        # Closed label sets: tokens that leave the trie are masked  [app. B]
        if candidate_trie is not None:
            log_probs = mask_outside_trie(log_probs, beams, candidate_trie)
        # Keep the best beam_size continuations  [§3.4]
        beams = top_k_extensions(beams, log_probs, beam_size)
    return best_beam(beams)

Pretraining tasks §3.2

Cross-modal (image-text data):

• Visual grounding: instruction "Which region does the text x^t describe?"; target is the box tokens.
• Grounded captioning: the inverse. Instruction "What does the region describe?" followed by the region's box tokens; target is the caption.
• Image-text matching: instruction "Does the image describe x^t?"; target "Yes" or "No". The negative pairs the image with a substituted caption.
• Image captioning: instruction "What does the image describe?".
• VQA: the question is the input and the answer is the target.

Unimodal:

• Image infilling: the middle of the image is masked. Instruction "What is the image in the middle part?"; target is the image codes §3.2. Code: the masked square spans the central region of a 2×code-size canvas unify_dataset.py.
• Object detection: instruction "What are the objects in the image?"; target is the sequence of box tokens and labels §3.2.
• Text infilling: BART-style span masking §3.2. The repository's masking default is ratio 0.3 with Poisson span lengths unify_task.py.

Masking for text infilling (cites are for the repository's implementation)

# Text infilling: mask word spans, then reconstruct the text  [§3.2]
def mask_text(subword_ids, mask_ratio=0.3):
    # Mask about mask_ratio of the words; span lengths follow a Poisson law
    # [code: raw/code/2202.03052/tasks/pretrain_tasks/unify_task.py#L66-L80]
    masked_ids = span_mask(subword_ids, mask_ratio)
    # Instruction used by the repository; the paper gives none for this task
    # [code: raw/code/2202.03052/data/pretrain_data/unify_dataset.py#L376]
    source = instruction_tokens(masked_ids)
    return source, subword_ids

Pretraining data §3.3 app. A.1 tab. 11

Images appearing in downstream validation and test sets were excluded §3.3.

Model sizes §3.5 tab. 1

The paper scales five models from 33M to 930M parameters (§3.5 says 940M; the table says 930M). Base and Large are sized like BART-Base and BART-Large §3.5.

Hyperparameters

Model configuration tab. 1:

Pretraining app. A.2:

Finetuning app. A.3:

Downstream instructions tab. 12: captioning "What does the image describe?"; VQA the question itself; visual entailment "Can image and text1 ... imply text2 ...?"; referring expressions "Which region does the text ... describe?"; image generation "What is the complete image? caption: ..."; GLUE tasks as yes/no questions.

Results

Author-reported. Setup: pretraining on the public data in tab. 11; models from Tiny to Huge tab. 1. Baselines are the strongest listed prior work in each table. Hardware is not stated in the LaTeX source.

GLUE tab. 6, OFA per task: SST-2 96.6, RTE 91.0, MRPC 91.7, QQP 92.5, MNLI 90.2, QNLI 94.8. OFA is not best on every task: XLNet has 97.0 on SST-2 and DeBERTa 91.9 on MRPC tab. 6.

Zero-shot tab. 9: OFA-Base against Uni-Perceiver. SST-2 71.6 vs 70.6; RTE 56.7 vs 55.6; MRPC F1 79.5 vs 76.1; QQP F1 54.0 vs 53.6; QNLI 51.4 vs 51.0; MNLI 37.3 vs 49.6 (worse); SNLI-VE dev/test 49.71 / 49.18 (Uni-Perceiver not listed).

Ablations, pretrained 250K steps, Base size tab. 10: full model CIDEr 135.6, VQA 76.0, ImageNet 82.2, image generation FID 20.8 / CLIPSIM 31.6 / IS 21.5. Removing text infilling: 134.8, 75.6, 83.2, 20.3 / 31.7 / 21.8. Removing image infilling: 136.3, 76.3, 81.8, 23.2 / 31.0 / 20.0. Removing detection: 133.3, 75.4, 81.4, 20.9 / 31.5 / 21.6. Removing grounding: 134.2, 75.5, 82.0, 21.2 / 31.5 / 21.5.

Trie-based search on Base, 250K steps tab. 13: VQA 76.03 vs 75.86 without Trie; SNLI-VE 89.2 vs 89.0; ImageNet 82.2 vs 81.9; MRPC F1 90.6 vs 90.1; QQP F1 88.4 vs 88.2.

Paper vs code

The official repository is linked from the abstract: https://github.com/OFA-Sys/OFA, Apache-2.0, commit a36b91ce86ff105ac8d9e513aa88f42b85e33479. The following differences were checked against the repository's scripts and code. The pretraining scripts are single-GPU examples.

• Peak pretraining LR. Paper: 2e-4 app. A.2. Code: 1e-4 in both pretraining scripts pretrain_ofa_large.sh pretrain_ofa_base.sh.
• Random patch sampling. Paper: random sampling of patches "does not bring additional benefits"; the LaTeX sentence about sampling 196 patches is commented out app. A.2. Code: the pretraining scripts set sample_patch_num=196 pretrain_ofa_large.sh, the criterion passes it on when positive label_smoothed_cross_entropy.py, and the model draws a random subset of patches per example during training unify_transformer.py. Which setting the released checkpoints used is not determined here.
• Pretraining resolution for Large. Paper: trained at 384 and continued at 480 for Large and Huge app. A.2. Code: the large pretraining script uses 384 and no 480 continuation script is in the repository pretrain_ofa_large.sh.
• Text length. Paper: encoder and decoder maximum 256 app. A.2. Code: pretraining scripts set source 80 and target 30 pretrain_ofa_large.sh. The text-only branch encodes up to 512 tokens and keeps the last 256 unify_dataset.py.
• Image-text matching. Paper: a separate cross-modal task with randomly substituted negative captions §3.2. Code: each caption or QA example also yields a matching example, positive or negative with probability 0.5 each unify_dataset.py. For negatives, when object tags exist, the code substitutes an object name with probability 0.4; otherwise it draws a random caption unify_dataset.py.
• Captioning cross-entropy stage. Paper: 2 epochs, batch 128 app. A.3. Code: the base-size caption script runs 5 epochs train_caption_stage1_base.sh train_caption_stage1_base.sh. The large-size caption script was not checked.
• Gigaword decoding. Paper: beam 6, length penalty 0.7, maximum output 32 app. A.3. Code: the same plus no_repeat_ngram_size=3, which the paper does not state train_gigaword.sh.
• Text infilling instruction. Paper: no instruction given §3.2. Code: "what is the complete text of ..." unify_dataset.py.
• Matches. Architecture flags (head scaling, LayerNorms) pretrain_ofa_large.sh. Dropout 0.1, stochastic depth 0.1 and weight decay 0.01 pretrain_ofa_large.sh pretrain_ofa_large.sh. Warmup ratio 0.01 with linear decay pretrain_ofa_large.sh pretrain_ofa_large.sh. Model shapes ofa.py. VQA (40K updates, lr 5e-5, label smoothing 0.1, EMA 0.9999) train_vqa_distributed.sh train_vqa_distributed.sh train_vqa_distributed.sh train_vqa_distributed.sh. RefCOCO (10 epochs, lr 3e-5, input 512) train_refcoco_el.sh train_refcoco_el.sh train_refcoco_el.sh. SNLI-VE (6 epochs, lr 2e-5) train_snli_ve.sh train_snli_ve.sh. ImageNet (160,000 updates, lr 5e-5, input 480) train_imagenet_distributed.sh train_imagenet_distributed.sh train_imagenet_distributed.sh. Gigaword (6 epochs, lr 1e-4) train_gigaword.sh train_gigaword.sh.
• Not checkable. The VQA candidate-answer file (trainval_ans2label.pkl) is referenced by the script but not in the repository train_vqa_distributed.sh. The pretraining TSV files are read from ../../dataset/, which is not in the repository pretrain_ofa_large.sh. Global batch sizes are not determinable from the single-GPU scripts.

Pitfalls

• Batch sizes are not reproducible from the scripts. The scripts use one GPU with batch 4 pretrain_ofa_large.sh pretrain_ofa_large.sh. The paper's per-step batch composition is not reproduced by these scripts, and the per-device split is not stated.
• Data is not in the repository. The pretraining TSV files, the negative-caption lists and the VQA answer vocabulary must be rebuilt from public sources. The paper describes filtering out downstream validation and test images §3.3; the repository files read here contain no such filter.
• Headline data size. The abstract says 20M image-text pairs. Table 11 lists 15.25M captioning/ITM samples, 2.92M VQA samples and 3.20M grounding samples tab. 11. The text does not explain how the two figures relate.
• Model size. §3.5 says 940M parameters; Table 1 says 930M for Huge §3.5 tab. 1.
• Trie-constrained decoding gives small gains in the authors' ablation: 0.17 to 0.5 points on each task tab. 13. The paper's stated reason for it is that the model may otherwise generate labels outside the closed set §3.4.
• Image generation sampling differs from the baselines. OFA samples 24 images and keeps the best by CLIP app. A.3. The paper says DALLE uses 512 samples and CogView and NÜWA use 60 §4.1. FID and IS are not computed under the same sampling budget.
• Zero-shot is brittle. The paper says results are highly sensitive to instruction wording and that sentence-pair tasks score below 60% §4.3. On MNLI, OFA-Base scores 37.3 against Uni-Perceiver's 49.6 tab. 9.
• No variance reported. The ablation tables show one number per setting, with differences of 0.2 to 0.5 points in several rows tab. 10 tab. 13. The tables give no spread across runs.
• Table and section numbers were counted from the LaTeX source of arXiv v2 (updated 2022-06-01), not from a rendered PDF. Check against the PDF before citing a number to a reader.
• No images on this card. The paper's figures show generation and grounding examples. They are described here only by section (§4.1, §4.3, app. C).

Lineage

• Builds on: BART (architecture and BPE tokenizer) §3.1 app. A.2; VQGAN-style image quantisation §3.1; patch features from ResNet blocks, following CoAtNet and SimVLM §3.1; discrete coordinate tokens following Pix2Seq §3.1; instruction-based task formatting from FLAN and T0 §3.2; NormFormer stability changes §3.1; TUPE position decoupling §3.1; BEiT-style masked image modelling §3.2; E2E-VLP for detection-style pretraining §3.2. The paper calls this work the latest in the authors' M6 series §1.
• Compared against: SimVLM, VLMo, UNICORN, VinVL, UNIMO, LEMON, ALBEF, METER, Florence (classification and captioning) tab. 2 tab. 3; MDETR and UNICORN (grounding) tab. 4; DALLE, CogView, GLIDE and NÜWA (image generation) tab. 5; Uni-Perceiver (zero-shot) tab. 9. For unimodal tasks, RoBERTa, XLNet, ELECTRA, DeBERTa, BEiT and MAE tab. 6 tab. 8.
• Superseded by: not stated in content already on this wiki.