PaperlayerView exact revision on arXiv

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.08916v2cs.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.

  • 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.

  • 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.

  • Takeaways & Limitations

    The unified architecture performs well across heterogeneous tasks, and removing task groups does not significantly affect individual-task performance.

  • Takeaways & Limitations

    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.

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.
  • Vision tasks expose heterogeneous outputs, including bounding boxes, binary masks, text answers, and per-pixel depth maps, complicating single-model design.
  • 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.
  • UNIFIED-IO jointly trains on over 90 datasets spanning computer vision, vision-and-language, and NLP tasks.
  • 64.3 overall score: UNIFIED-IO supports all 7 GRIT tasks and exceeds the second-best model by 32.0.

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.
  • The task taxonomy distinguishes Text, Image, Sparse, and Dense input or output modalities.
  • Image Synthesis: Image synthesis generates images from text, partial images, inpainting targets, or semantic segmentation maps.
  • Sparse Labelling: Sparse labelling identifies referred regions or keypoints, including detection, localization, pose estimation, and referring expression tasks.
  • Dense Labelling: Dense labelling produces per-pixel depth, surface-orientation, or semantic-class labels from images.
  • Additional groups cover image classification, image captioning, vision-language reasoning, and text-only NLP tasks.

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.
  • 3.1 UNIFIED TASK REPRESENTATIONS: Text uses SentencePiece tokens, while dense labels are converted to RGB images and encoded with a VQ-GAN.
  • 3.1 UNIFIED TASK REPRESENTATIONS: Bounding boxes and human joints use discretized coordinate tokens, with points encoded by two tokens and boxes by four.
  • 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.
  • 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.
  • 3.3 TRAINING: Multi-task training ensembles 95 datasets and mixes groups and datasets to expose the model to underrepresented tasks.

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.
  • 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.
  • Same and New Concepts: 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.
  • 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.
  • 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.

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.
  • Many unified models handle text and image inputs with text outputs, but cannot produce structured or visual outputs.
  • More recent systems represent image locations through detector proposals, output heads, or special coordinate tokens.
  • Gato adds modalities such as game controls and robot torques, while Flamingo accepts interleaved text, images, and videos, but neither produces image outputs.
  • 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.

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.
  • The 2.9B-parameter UNIFIED-IO XL model is jointly trained on 90+ datasets spanning vision, language, and multimodal tasks.
  • 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.

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.
  • The training and inference setup specifies prompts for the supported tasks.
  • The paper provides qualitative examples comparing ground-truth outputs with UNIFIED-IO predictions.

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.
  • Image Inpainting: Image inpainting fills a blanked bounding-box region with a target object using the source image, location, and category as inputs.
  • Image Synthesis from Segmentation: Image synthesis from segmentation generates an image matching semantic pixel labels represented through a color-coded input image and prompt.

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.
  • 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.
  • Referring Expression Comprehension: Referring expression comprehension localizes an image region described by a natural-language expression.RefCOCO, RefCOCO+, and RefCOCOg provide training data.
  • 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.

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.
  • 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.
  • 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.

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.
  • 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.
  • Captioning: Image and region captioning generate text describing either the whole image or a specified image region.Region captioning uses Visual Genome training data.
  • 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.
  • 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.
  • 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.
Loading 2206.08916v2…