Unified-IO: A unified model for vision, language, and multi-modal tasks

In plain words

Most vision systems have a separate output part for each job: one for boxes, one for masks, one for depth, one for answers. Unified-IO drops those separate parts. It turns every input and every output into a sequence of tokens taken from one shared vocabulary. Words become word pieces, positions in an image become special location tokens, and whole images or dense maps become codes from a pre-trained image compressor. One encoder-decoder transformer reads a text prompt plus the input and writes the answer as tokens. Those tokens are then converted back into a box, a mask, a depth map or a sentence. The model is trained on a large mix of tasks at once, so one set of weights answers questions, captions pictures, finds objects and predicts depth. The point is that one model can cover many vision and language tasks without task-specific parts, and the authors report that it stays competitive on about 16 further benchmarks without fine-tuning for any of them.

TL;DR

Unified-IO is one T5-style encoder-decoder that performs the seven GRIT tasks and many vision, vision-language and NLP tasks. Every input and output becomes discrete tokens: text by SentencePiece, points and boxes by 1000 location tokens, images and dense maps by a 16384-entry VQ-GAN codebook. Trained jointly on 95 datasets with no task-specific fine-tuning, the XL model reports 64.3 on GRIT test against 32.0 for GPV-2 tab. 3. The official code is inference-only, so the training recipe rests on the paper alone.

Problem

Vision tasks have very different outputs: boxes, per-pixel labels, depth, text §1. Earlier unified models mostly keep a shared backbone with one head per task or per modality. Language decoders are easy to share, so unified progress has centred on vision-language tasks and often covers only a few tasks §1.

The goal here is one model with little or no task-specific design or parameters. That allows general pre-training, joint training on many tasks, and knowledge sharing across tasks §3. The approach assumes that everything can be written as a token sequence, and it depends on large multi-task data (95 datasets §3.3) and large compute (the XL model used 128-way data parallelism §3.4). Its own limitations include an image tokenizer pre-trained only on ImageNet, weak recall on cluttered object detection and limited image generation §4.6.

Method

Unified-IO is one sequence-to-sequence model. A natural-language prompt tells it which task to do, and inputs and outputs are all discrete tokens §3.1 fig. 2. The network follows T5 with three changes: patch embedding for images, an expanded vocabulary, and position encodings that are 2-D and absolute §3.2.

Text. Text is tokenized with SentencePiece, as in T5 §3.1. Each task gets a natural-language prompt, such as "What is the depth map of the image?" for depth estimation. VQA is the exception: the question itself is the prompt §3.1.

Input images. An image is resized to 384×384, cut into 16×16 patches (576 of them), and each patch is embedded with a linear projection §3.2. The paper does not say how non-square images are fitted to 384×384; the code resizes and pads and returns a patch mask utils.py.

Dense outputs and images as outputs. Depth becomes a normalised grayscale image, and surface normals map x/y/z to r/g/b. Segmentation gives each instance a random colour, and the colour-to-class mapping is written into the text rather than fixed in advance §3.1. Each RGB target is resized to 256×256 and encoded by an ImageNet-pre-trained VQ-GAN with a compression ratio of 16 and a 16384-entry codebook. Those codes join the vocabulary; at inference the VQ-GAN decoder turns the generated codes back into an image §3.1. A 256×256 image gives a 16×16 grid, 256 codes §3.2.

Sparse outputs. The vocabulary gets 1000 location tokens for discretised coordinates §3.1. A point is two tokens, x then y. A box is four tokens: two for one corner and two for the opposite one §3.1. A labelled box is the box followed by a text label. Joints are a sequence of points, each followed by a text visibility label §3.1. The paper names the corners "upper right" and "lower left"; the code uses the opposite order, see Paper vs code.

# Unified-IO forward pass, one example  [§3.1] [§3.2] [fig. 2]
# Input:  prompt text, optional image, target tokens
# Output: logits over the shared vocabulary

def forward(prompt, image, target):
    # Prompt as SentencePiece tokens, at most 256  [§3.2]
    text_tokens = sentencepiece(prompt)
    # Image cut into 16x16 patches, linearly embedded  [§3.2]
    image_patches = patchify(image, size=16)   # (B, 576, 16*16*3)
    image_embeddings = linear(image_patches)   # (B, 576, d)
    # Absolute and 2-D relative position information  [§3.2]
    encoded = encoder(text_tokens, image_embeddings)
    # Text positions first, then image positions  [code: uio/network.py#L911]
    logits = decoder(encoded, target)
    return logits

# Sparse box to location tokens  [§3.1] [code: uio/utils.py#L36-L52]
def box_to_tokens(box, width, height):
    # Scale by the longer image side  [code: uio/utils.py#L36-L52]
    scale = max(width, height)
    # Corners in y, x order, quantised into 1000 bins  [§3.1]
    corners = [box.y1, box.x1, box.y2, box.x2] / scale
    bins = [int(999 * c) for c in corners]     # truncation, as in code
    return [location_token(b) for b in bins]

# Dense target to VQ-GAN codes  [§3.1] [§3.2]
def dense_to_tokens(rgb_image):
    # Depth as normalised grayscale; normals as r/g/b  [§3.1]
    # Resize to 256x256, 16x downsampling  [§3.1]
    codes = vqgan_encode(rgb_image)            # 256 codes, 16384 entries
    return codes

Architecture. Encoder and decoder are stacks of residual transformer layers: self-attention, cross-attention in the decoder, and feed-forward blocks, with layer norm before each sub-block, as in T5 §3.2. The text limits are 256 input tokens and 128 output tokens. Image inputs are 576 patches; image outputs are 256 codes §3.2. The relative positions are extended from 1-D to 2-D with a fixed number of learned embeddings, and an absolute position embedding is added to the token embedding §3.2.

Pre-training §3.3 uses two denoising objectives, given equal sampling. Text span denoising corrupts 15% of tokens, and each run of consecutive corrupted tokens is replaced by one unique mask token. Masked image denoising masks 75% of the image patches and asks for the whole image back, following BEiT and MAE. When the other modality is present, the model can use it. The text data is C4 and Wikipedia, and the image-plus-text data is ImageNet-21K class names and YFCC15M captions, with the prefix "An image of" §3.3 app. A.2.

Multi-task training §3.3 mixes 95 datasets from 62 public sources, with examples mixed inside each batch §3.3 tab. 1. The eight task groups are sampled at 1/8 each, except image synthesis (3/16) and dense labelling (1/16) §3.3. Within a group, datasets are sampled in proportion to the square root of their size §3.3.

# Multi-task batch sampling  [§3.3] [app. A.3]
def sample_dataset(groups):
    # Six groups get 1/8 of each batch; image synthesis 3/16,
    # dense labelling 1/16  [§3.3]
    group = draw(groups, weights=group_weights)
    # Inside a group, weight grows with the square root of size  [§3.3]
    dataset = draw(group.datasets, weights=sqrt(group.sizes))
    return dataset

Inference runner.py. The released runner decodes text with beam search by default and images by sampling. Text output is capped at 128 tokens. Keypoints are a two-stage procedure: first find the person with the localisation prompt, then find keypoints for each region §4.1.

Hyperparameters

Values are as stated in the paper unless a row says it comes from code. The training code is not in the released repository, so the optimiser and schedule rows are paper-only.

Model sizes tab. 2:

VQ-GAN parameters are excluded from these counts tab. 2.

Training data, grouped tab. 1. The table has one row per task; only group rows are shown here.

Results

Author-reported. Setup: GRIT is an evaluation-only benchmark with seven tasks; the main comparison is on the GRIT test split against Mask R-CNN, OFA, GPV-1 and GPV-2 §4.1 tab. 3. Each cell is ablation split / test split; dashes mean the entry does not report that cell. The Small, Base and Large models only report the ablation split.

Same-concept vs new-concept split on the GRIT ablation set tab. 4. Mask R-CNN and GPV-2 are the reference points the text uses; "restricted" is a column of the source table and is not defined in the active text.

Ablation: each task group is held out of multi-task training. LARGE, 250k steps, GRIT ablation split and MNLI tab. 5 §4.3.

Sixteen further benchmarks, no fine-tuning tab. 6 §4.4. Each column reports one benchmark; "Unified SOTA" is the best prior result from another multi-task, unified model, and "Single / fine-tuned SOTA" is the best specialised or fine-tuned result. Table 6 has 17 numeric columns; see Pitfalls.

Text claims, checked against the tables. Depth 0.385 RMSE is behind the specialised SOTA but ahead of UViM §4.4 tab. 6. A-OKVQA is a state-of-the-art result §4.4 tab. 6. OK-VQA is below Flamingo §4.4 tab. 6. Prompt paraphrases change performance a lot in some cases, and spacing around punctuation matters §4.5 tab. 7.

Paper vs code

The official repository is allenai/unified-io-inference (Apache-2.0, commit ee0bd9d848aafc855c0f3dbe24ac8714baa87496). Its README cites this paper, and its owner is the Allen Institute for AI, where the authors work README.md. The repository is inference only: it contains the network, configs and runner, but no training loop, data pipeline, optimiser or loss. Training values therefore come from the paper alone. Weights are linked in the README from a public bucket, but the repository states no licence for them, and this card does not use them README.md.

Differences found by reading the code against the paper:

• Attention width of Small. The paper gives Small a model dim of 512 and 6 heads tab. 2. The config fixes head_dim at 64 for every size network.py, and Small sets 6 heads configs.py. So Small's attention width is 6 × 64 = 384, not 512. This is inferred from code, not stated.
• Relative position buckets. The paper says only that the 1-D relative embedding extends to 2-D with learned embeddings §3.2. The code uses 32 buckets with max distance 128 for text, and 8 buckets with max distance 20 for image positions network.py. Image position ids are offset by the maximum text length of 256 network.py.
• Logits. When logits are tied to the embedding, the code divides by √(model dim) network.py. The paper does not mention this.
• Output masking. Text positions may emit only text or location tokens, and image positions only image codes network.py. The paper implies but does not state this rule.
• MLP activation. The code uses a GELU/linear pair, a gated MLP network.py. The paper does not name the activation.
• XL segment embedding. XL sets num_seg_emb=8 configs.py; the paper does not mention segment embeddings.
• Location token layout. The paper gives the vocabulary as 32152 + 1000 + 16384 §3.4. The config stores text plus location as 33152 network.py. A helper file sets vocab_size = 33100 and derives location ids from that value utils.py. These three numbers cannot be reconciled from this repository alone. Use region_to_tokens for the exact mapping utils.py.
• Box corner order. The paper says a box is encoded as the upper-right corner then the lower-left corner §3.1. The code takes the box as [x1, y1, x2, y2], reorders it to y, x, scales by the longer image side and quantises to 999 steps utils.py. So the order is top-left then bottom-right, with y before x. Follow the code, since it is what the released model was run with.
• Input resizing. The code resizes and pads a non-square image to 384×384 and returns a patch mask utils.py. The paper does not say how shapes are handled.
• Decoding. Text is decoded by beam search by default and images by sampling runner.py. Sampling temperature defaults to 1.0 decoding.py. The paper does not describe decoding.
• Text output length. The paper caps outputs at 128 text tokens §3.2. The code allows 256 decoder text positions network.py but the runner asserts at most 128 runner.py. The two agree in practice.
• Image sizes. Image input is 384×384 §3.2 and the code agrees utils.py. Image output is 256×256 §3.2 and the code agrees.

Pitfalls

• Learning-rate decay is underspecified. The rate is 0.01 for the first 10,000 steps, then falls with the inverse square root of the step count §3.4. The text does not say whether that is a multiplier or an absolute schedule, or whether there is warm-up. The training code is not released, so this cannot be checked. Pick one reading and note it in your code.
• The loss function is never written out. The paper describes the targets as discrete tokens for both text and images §3.3, but gives no loss equation. The inference code computes logits with masks network.py and does not contain a training loss.
• Two statements about task mixing. Section 3.3 says datasets within a group are sampled in proportion to the square root of their size §3.3. The appendix says each task's mixing rate is raised to the power 1/T with T = 2 and then renormalised app. A.3. Both produce a square root, but the appendix works per task and the text works per dataset. It is not stated which one training used.
• Normal estimation compares different splits. The text says Unified-IO gets "strong results" against NLL-AngMF, 44.3 against 49.6 §4.1. The 44.3 is Unified-IO's test score and 49.6 is NLL-AngMF's ablation score. On matched splits Unified-IO is behind either way: 45.0 against 49.6 on ablation, and 44.3 against 50.5 on test tab. 3. Quote the matched numbers, not the text's comparison.
• Table 6 has one missing cell and an unexplained column. The LARGE row has 16 values for 17 columns. Swig is the missing one (the row goes from VizWiz-G straight to SNLI-VE), so it is shown as "not reported" above. The table has two COCO columns, but the caption names only "COCO Captions". The abstract says 16 benchmarks, and the caption lists 15.
• Dataset counts vary. The introduction and abstract say "over 90" datasets §1, and Section 2 and Table 1 say 95 datasets from 62 sources §2 tab. 1. Use 95 for the training mix.
• Ablations are not the full model. The ablation runs use LARGE for 250k steps and the paper warns that conclusions may change at XL size §4.3. Do not read the ablation deltas as the full model's behaviour.
• Image tokens cap image quality. The VQ-GAN is ImageNet-pre-trained and its codebook is fixed for all tasks. The authors say this limits image generation and depth and normal quality, and that a stronger tokenizer would help §4.1 §4.6.
• Recall and prompts are weak spots. Object detection recall is poor in cluttered images §4.6. Paraphrased prompts can drop performance sharply §4.5 tab. 7.
• Name clash. The allenai organisation also publishes unified-io-2, a different paper (Unified-IO 2) with its own code. Do not mix its checkpoints or code with this card, which covers the first version.
• Model size wording. The paper says XL has "2.9 billion" parameters §3.2 and Table 2 gives 2925M tab. 2. Use 2925M when you need an exact count.

Lineage

• Builds on: T5, the text-to-text encoder-decoder whose layout Unified-IO copies §3.2; VQ-GAN, the image tokenizer §3.1; Pix2Seq, which encodes boxes as discrete tokens §3.1; BEiT and MAE for masked image denoising §3.3.
• Compared against: Mask R-CNN, GPV-1, GPV-2 and OFA on GRIT tab. 3; CLIP for categorisation tab. 3; UViM on depth §4.4; Flamingo, KAT and CoCa on VQA and captioning tab. 6.
• Superseded by: not stated in this wiki.