Unified-IO: A Unified Model for Vision, Language, and Multi-Modal Tasks

Jiasen Lu; Christopher Clark; Rowan Zellers; Roozbeh Mottaghi; Aniruddha Kembhavi

Source-linked AI summary

Unified-IO: A Unified Model for Vision, Language, and Multi-Modal Tasks

Jiasen Lu, Christopher Clark, Rowan Zellers, Roozbeh Mottaghi, Aniruddha Kembhavi

arXiv:2206.08916v2Jun 17, 2022cs.CV

TL;DR

UNIFIED-IO addresses the difficulty of one model handling diverse vision, language, and multimodal representations. It tokenizes these inputs and outputs for a shared transformer trained jointly across many datasets, supporting all seven GRIT tasks and strong results across additional benchmarks without task-specific fine-tuning.

Problem
Heterogeneous inputs and outputs across vision and language tasks make it difficult to build one model without task-specific designs.
Source: §1 · p. 1
Method
UNIFIED-IO homogenizes diverse modalities into discrete token sequences and processes them with a unified encoder-decoder transformer jointly trained across over 90 datasets.
Source: §1 · p. 2, §3.2 · p. 5
Results
UNIFIED-IO supports all 7 GRIT tasks, achieves a 64.3 overall score, and performs strongly across 16 additional benchmarks without benchmark fine-tuning.
Source: §1 · p. 2, §6 · p. 12
Takeaways & Limitations
The unified architecture performs well across heterogeneous tasks, and removing task groups does not significantly affect individual-task performance.
Source: §4.3 · p. 8
Object-detection recall is often poor in cluttered images, while the pre-trained VQ-GAN limits image-generation capabilities and prompts unseen during training are not always understood.
Source: §4.6 · p. 10

Abstract

from arXiv · show

We propose Unified-IO, a model that performs a large variety of AI tasks spanning classical computer vision tasks, including pose estimation, object detection, depth estimation and image generation, vision-and-language tasks such as region captioning and referring expression, to natural language processing tasks such as question answering and paraphrasing. Developing a single unified model for such a large variety of tasks poses unique challenges due to the heterogeneous inputs and outputs pertaining to each task, including RGB images, per-pixel maps, binary masks, bounding boxes, and language. We achieve this unification by homogenizing every supported input and output into a sequence of discrete vocabulary tokens. This common representation across all tasks allows us to train a single transformer-based architecture, jointly on over 90 diverse datasets in the vision and language fields. Unified-IO is the first model capable of performing all 7 tasks on the GRIT benchmark and produces strong results across 16 diverse benchmarks like NYUv2-Depth, ImageNet, VQA2.0, OK-VQA, Swig, VizWizGround, BoolQ, and SciTail, with no task-specific fine-tuning. Code and demos for Unified-IO are available at: https://unified-io.allenai.org.

1 INTRODUCTION

UNIFIED-IO addresses the challenge of heterogeneous vision inputs and outputs by representing diverse tasks with one token-based sequence-to-sequence architecture. Joint training across many datasets yields broad task coverage and strong aggregate benchmark performance.

UNIFIED-IO jointly performs classical vision, image synthesis, vision-and-language, and NLP tasks with one neural model.Source: §1 · p. 1
Vision tasks expose heterogeneous outputs, including bounding boxes, binary masks, text answers, and per-pixel depth maps, complicating single-model design.Source: §1 · p. 1
UNIFIED-IO homogenizes task outputs into discrete token sequences, using VQ-VAE for dense structures, coordinate tokens for sparse structures, and byte-pair encoding for language.Source: §1 · p. 2
UNIFIED-IO jointly trains on over 90 datasets spanning computer vision, vision-and-language, and NLP tasks.Source: §1 · p. 2
64.3 overall score: UNIFIED-IO supports all 7 GRIT tasks and exceeds the second-best model by 32.0.Source: §1 · p. 2

2 VISION, LANGUAGE AND MULTI-MODAL TASKS

UNIFIED-IO organizes 95 datasets into task groups covering language, vision, and multimodal inputs and outputs. The taxonomy spans synthesis, labeling, classification, captioning, vision-language reasoning, and NLP.

UNIFIED-IO gathers 95 datasets from 62 public sources for multi-task learning across language, vision, and multimodal tasks.Source: §2 · p. 3
The task taxonomy distinguishes Text, Image, Sparse, and Dense input or output modalities.Source: §2 · p. 3

Image Synthesis

Image synthesis generates images from text, partial images, inpainting targets, or semantic segmentation maps.Source: §2 · p. 3

Sparse Labelling

Sparse labelling identifies referred regions or keypoints, including detection, localization, pose estimation, and referring expression tasks.Source: §2 · p. 3

Dense Labelling

Dense labelling produces per-pixel depth, surface-orientation, or semantic-class labels from images.Source: §2 · p. 3

Additional groups cover image classification, image captioning, vision-language reasoning, and text-only NLP tasks.Source: §2 · p. 4

3 UNIFIED-IO

UNIFIED-IO converts heterogeneous modalities into shared token sequences and processes them with a T5-like encoder-decoder transformer. It combines pre-training with large-scale multi-task training without task-specific fine-tuning.

3.1 UNIFIED TASK REPRESENTATIONS

UNIFIED-IO represents text, images, and structured outputs as tokens from a shared finite vocabulary.Source: §3.1 · p. 4
Text uses SentencePiece tokens, while dense labels are converted to RGB images and encoded with a VQ-GAN.Source: §3.1 · pp. 4–5
Bounding boxes and human joints use discretized coordinate tokens, with points encoded by two tokens and boxes by four.Source: §3.1 · p. 5

3.2 UNIFIED ARCHITECTURE

The model is a pure T5-like encoder-decoder transformer with self-attention, decoder cross-attention, feed-forward networks, and residual layers.Source: §3.2 · p. 5

3.3 TRAINING

UNIFIED-IO trains in two stages: unsupervised text, image, and image-text pre-training followed by massive multi-task training without task-specific fine-tuning.Source: §3.3 · p. 6
Multi-task training ensembles 95 datasets and mixes groups and datasets to expose the model to underrepresented tasks.Source: §3.3 · p. 6

4 EXPERIMENTS

Experiments evaluate Unified-IO across GRIT, concept generalization, task-group ablations, and 16 additional benchmarks. The model supports broad heterogeneous task coverage while retaining strong performance, though limitations remain in NLP, detection recall, and image generation.

GRIT Results

Unified-IO supports all seven GRIT tasks and achieves the benchmark’s top overall score of 64.3, versus 32.0 for GPV-2, which supports four tasks.It also outperforms OFA_LARGE on VQA, referring expression, and categorization.Source: §4.1 · p. 7

Same and New Concepts

UNIFIED-IOXL shows little degradation from same to new concepts and sometimes performs better on new concepts, indicating broad training-data concept coverage.The authors note that this training volume can provide supervision nearly as effective as large standard vision datasets such as COCO.Source: §4.2 · p. 8
74.4 versus 40.8 for localization and 64.2 versus 0.3 for segmentation show Unified-IO’s advantage over Mask R-CNN on new concepts.On COCO-like examples, it also leads Mask R-CNN on localization, 65.6 versus 51.9, and segmentation, 53.0 versus 44.9.Source: §4.2 · p. 8

Task-Group Ablations

Removing task groups generally does not significantly affect remaining-task performance, despite reducing the scope learned at fixed model capacity.Removing NLP boosts categorization, removing captioning boosts VQA and some other tasks, and removing image synthesis causes a major change in performance.Source: §4.3 · p. 8

Limitations

Performance is weaker or constrained in several settings: NLP trails SOTA, cluttered-image detection recall is often poor, and the pretrained VQ-GAN limits image-generation ability.Prompt paraphrases and punctuation formatting can also produce substantial performance changes.Source: §4.4 · p. 10, §4.6 · p. 10, §4.5 · p. 10

5 RELATED WORK

Prior unified vision models commonly share a visual backbone but retain task-specific branches or heads. Recent approaches broaden inputs and outputs, yet important gaps remain in supporting structured and dense visual outputs together.

Shared-backbone architectures typically add specialized branches or heads for detection, segmentation, and vision-and-language tasks.Source: §5 · p. 11
Many unified models handle text and image inputs with text outputs, but cannot produce structured or visual outputs.Source: §5 · p. 11
More recent systems represent image locations through detector proposals, output heads, or special coordinate tokens.Source: §5 · p. 11
Gato adds modalities such as game controls and robot torques, while Flamingo accepts interleaved text, images, and videos, but neither produces image outputs.Source: §5 · p. 11
OFA supports image locations and text-to-image synthesis, whereas it does not support dense tasks such as depth estimation, segmentation, and surface-normal estimation.Source: §5 · p. 11

6 CONCLUSION

UNIFIED-IO uses one discrete-token representation to unify diverse vision, language, and multimodal inputs and outputs. The model is jointly trained across more than 90 datasets and performs broadly without benchmark-specific fine-tuning or task-specific modifications.

UNIFIED-IO homogenizes images, continuous maps, binary masks, text, bounding boxes, and keypoints into sequences of discrete tokens.Source: §6 · p. 12
The 2.9B-parameter UNIFIED-IO XL model is jointly trained on 90+ datasets spanning vision, language, and multimodal tasks.Source: §6 · p. 12
UNIFIED-IO is the first model to perform all seven GRIT tasks and obtains strong results across 16 additional benchmarks without benchmark fine-tuning or task-specific modifications.Source: §6 · p. 12

A.1 TASKS DETAILS

UNIFIED-IO is jointly trained on a large and diverse collection of vision, language, and vision-language tasks. The task setup includes qualitative comparisons between ground-truth outputs and model predictions.

UNIFIED-IO is jointly trained across a large and diverse set of vision, language, and vision-language tasks.Source: §A.1 · p. 19
The training and inference setup specifies prompts for the supported tasks.Source: §A.1 · p. 19
The paper provides qualitative examples comparing ground-truth outputs with UNIFIED-IO predictions.Source: §A.1 · p. 19

A.1.1 IMAGE SYNTHESIS TASKS

The image-synthesis tasks train UNIFIED-IO to generate complete images from text, fill masked regions using locations and target categories, and synthesize images from semantic segmentation inputs.

Image Synthesis from Text

Text-to-image synthesis generates an image matching a sentence using captioning and image-classification datasets.Source: §A.1.1 · p. 19

Image Inpainting

Image inpainting fills a blanked bounding-box region with a target object using the source image, location, and category as inputs.Source: §A.1.1 · p. 19

Image Synthesis from Segmentation

Image synthesis from segmentation generates an image matching semantic pixel labels represented through a color-coded input image and prompt.Source: §A.1.1 · p. 19

A.1.2 SPARSE LABELLING TASKS

Unified-IO represents sparse labelling tasks as sequences of coordinate tokens, supporting detection, localization, referring expressions, and keypoint estimation within one architecture.

Object Detection

Object detection outputs bounding boxes and class names for all objects in an image from a static prompt and the image.The model randomizes output-object order during training.Source: §A.1.2 · p. 19

Object Localization

Object localization returns all bounding boxes containing instances of a prompted object category.Training examples are constructed from each category present in object-detection images.Source: §A.1.2 · p. 20

Referring Expression Comprehension

Referring expression comprehension localizes an image region described by a natural-language expression.RefCOCO, RefCOCO+, and RefCOCOg provide training data.Source: §A.1.2 · p. 20

Keypoint Estimation

Keypoint estimation returns 17 human-body keypoint locations per person using a multi-step approach to avoid excessively long output sequences.The model first detects keypoints for individual people before producing the complete result.Source: §A.1.2 · p. 20

A.1.3 DENSE LABELLING TASKS

Unified-IO handles dense labelling by generating image-like outputs that encode per-pixel masks, depth, or surface orientation and can be converted into the required maps.

Object Segmentation

Object segmentation generates an RGB image whose uniquely colored regions are converted into binary instance masks for a prompted category.Colors are clustered and tiny connected components are removed during mask cleaning.Source: §A.1.3 · p. 21

Depth Estimation

Depth estimation assigns each pixel a depth value by generating a grayscale image of normalized depth and rescaling it to the training maximum.Training data comes from NYU Depth Dataset V2.Source: §A.1.3 · p. 21

Surface Normal Estimation

Surface normal estimation generates an RGB representation of each pixel’s x/y/z surface orientation, which is converted back into orientations.Training uses FrameNet and BlendedMVS surface-normal datasets.Source: §A.1.3 · p. 21

A.1.4 IMAGE CLASSIFICATION TASKS

The section covers classification, captioning, multimodal question answering, relationship prediction, and language tasks, while illustrating broad multi-task pre-training and depth performance.

Image Classification

Image classification maps an image and static prompt to a class name across six classification datasets.Inference scores candidate labels by log-probability and returns the highest-scoring class.Source: §A.1.4 · p. 21

Object Categorization

Object categorization predicts the class name for an image region specified by an image, bounding box, and prompt.Its labels are derived from Visual Genome, Open Images, and COCO detection annotations.Source: §A.1.4 · p. 22

Captioning

Image and region captioning generate text describing either the whole image or a specified image region.Region captioning uses Visual Genome training data.Source: §A.1.5 · p. 22

Vision & Language

Visual question answering generates answer text from a question prompt, while answer-grounded VQA additionally returns a binary mask for the evidence region.Unified-IO avoids restricting VQA outputs to a fixed list of common answers.Source: §A.1.6 · pp. 22–23

Natural Language Tasks

Question answering, text classification, summarization, and masked language modeling format language inputs as prompts and train the model to generate textual targets.These tasks use datasets including SQuAD, GLUE/SuperGLUE, Gigaword, C4, and Wikipedia.Source: §A.1.7 · pp. 23–24, §A.1.8 · p. 24

Training Distribution

0.385 vs. 0.467 RMSE: Unified-IO outperforms concurrent work on depth estimation despite depth data comprising only 0.43% of sampling.The authors suspect large model capacity and masked image denoising pre-training improve performance.Source: §A.3 · p. 26