Given any two conjugate numbers $p,q\in[1,\infty]$ s.t. $1/p+1/q=1$, we have

$\|AB\|_{p,\infty}\leq\|A\|_{p,q}\|B\|_{p,\infty}$

for any matrix $A\in\mathbb{R}^{m\times n}$ and matrix
$B\in\mathbb{R}^{n\times r}$.

######  Lemma F.7 (Lemma M.9 in Zhang et al. (2022)).

Given any two vectors $x,y\in\mathcal{R}^{d}$, we have

$\|\texttt{softmax}(x)-\texttt{softmax}(y)\|_{1}\leq 2\|x-y\|_{\infty}.$

######  Lemma F.8 (Property of Kullback-Leibler divergence, Proposition 7.2 in
Polyanskiy and Wu (2024)).

Given any two probability distributions $\mu_{1}$ and $\mu_{2}$ over
$(\Omega,\mathcal{F})$ and any two distributions $\nu_{1}$ and $\nu_{1}$ over
$(\Omega^{\prime},\mathcal{F}^{\prime})$, if $\mu_{1}\ll\mu_{2}$ and
$\nu_{1}\ll\nu_{2}$, then we have

$D_{\mathrm{kl}}(\mu_{1}\otimes\nu_{1}\|\mu_{2}\otimes\nu_{2})=D_{\mathrm{kl}}(\mu_{1}\|\mu_{2})+D_{\mathrm{kl}}(\nu_{1}\|\nu_{2}).$

##  Appendix G Experiment Details

###  G.1 Training data generation

We first describe a basic setup of all our experiments. For some experiments,
we change some part(s) in below to design the corresponding “flipped”
experiment or to examine the OOD ability of the trained transformer. In
particular, the $i$-th thread the training data

$\left(x_{1}^{(i)},y_{1}^{(i)},x_{2}^{(i)},y_{2}^{(i)},...,x_{T}^{(i)},y_{T}^{(i)}\right)$

is generated by the following distributions:

  * •

$\mathcal{P}_{X}$: the feature vector
$x_{t}^{(i)}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,I_{d})$ where $I_{d}$
is $d$-dimensional identity matrix.

  * •

$\mathcal{P}_{\epsilon}$: the noise
$\epsilon_{t}^{(i)}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,1)$.

  * •

$\mathcal{P}_{\sigma}$: the noise intensity $\sigma_{i}$ is sampled i.i.d.
from

$\tau_{i}\sim\mathrm{Gamma}(\underline{\tau},\bar{\tau}),\quad\sigma_{i}=\dfrac{1}{\sqrt{\tau_{i}}}$

where the parameters $\underline{\tau}=\bar{\tau}=20$ for the basic setup of
the experiment. We change these two parameters for some OOD experiments.

  * •

$\mathcal{P}_{\mathcal{F}}$: The function $f_{i}(x)\coloneqq w_{i}^{\top}x$
where $w_{i}$ is generated from

$w_{i}|\sigma_{i}\sim\mathcal{N}(\bar{w},\sigma_{i}^{2}\cdot I_{d})$

where $I_{d}$ is the $d$-dimension identity matrix and $\bar{w}$ is set to be
an all-one vector of dimension $d$. The covariance matrix of $w_{i}$ is
related with the noise intensity $\sigma_{i}$ to control the signal-to-noise
ratio.

Finally, the target variable is calculated by

$y_{t}^{(i)}=w_{i}^{\top}x_{t}^{(i)}+\sigma_{i}\epsilon_{t}^{(i)}.$

Throughout the paper, we consider the dimension $d=8.$

###  G.2 Number of Tasks $N$ and Training Procedure

In the previous Section G.1, we define how we generate the training data. As
in the previous work, we introduce the notion of task where each realization
of $(w_{i},\sigma_{i})$ is referred to as one task. The rationale is that each
configuration of $(w_{i},\sigma_{i})$ corresponds to one pattern of the
sequence $(x_{t},y_{t})$’s. While the distribution of $(w_{i},\sigma_{i})$
corresponds to infinitely many possible task configurations, we use a finite
pool of tasks for training the Transformer. Specifically, we generate

$\mathcal{T}\coloneqq\\{(w_{i},\sigma_{i})\\}_{i=1}^{N}$

from the distributions discussed above. Throughout the paper, we use $N$ to
refer to the total number of tasks or the pool size.

Training the Transformer for our setting is slightly different from the
classic ML model’s training. We do not use a fixed set of training data.
Rather, we generate a new batch of training data freshly for each batch.

  * •

The batch size $b=64.$ For each batch, we first sample with replacement $b$
tasks from the task pool $\mathcal{T}.$ And based on each sampled
$(w_{i},\sigma_{i}),$ we generate a training sequence
$\left(x_{1}^{(i)},y_{1}^{(i)},x_{2}^{(i)},y_{2}^{(i)},...,x_{T}^{(i)},y_{T}^{(i)}\right)$
following the setting in the previous Section G.1.

  * •

All the numerical experiments in our paper run for 200,000 batches.

The validation and testing sets are also randomly generated instead of fixed
beforehand. But unlike the training phase which draws the task configuration
from the task pool $\mathcal{T}$, the validation and test phase samples
$(w_{i},\sigma_{i})$ directly from the original distribution described in the
previous Section G.1. This is aimed to validate or test whether the trained
model has learned the ability to solve _a family of_ problems, or it only just
memorizes a fixed pool of tasks $\mathcal{T}$.

###  G.3 Derivation of Bayes-optimal Predictor

In Proposition 3.1, we state the Bayes-optimal predictor in the form of a
posterior expectation. Now we calculate the Bayes-optimal predictor explicitly
under the generation mechanism specified in Section G.1. Conditional on
history $H_{t}=(x_{1},y_{1},\ldots,x_{t})$, the posterior distribution of
$(w,\sigma)$ that governs the generation of $H_{t}$ can be calculated based on
the Bayesian posterior as

$\mathbb{P}(\tau|H_{t})=\mathrm{Gamma}(\tau;\underline{\tau}_{t},\bar{\tau}_{t}),\quad\sigma=\dfrac{1}{\sqrt{\tau}},$
$\mathbb{P}(w|\sigma,H_{t})=\mathcal{N}(w_{t},\sigma^{2}\cdot\Sigma_{t}),$

where

$\displaystyle\Sigma_{t}=\left(I_{d}+\sum_{s=1}^{t-1}x_{s}x_{s}^{\top}\right)^{-1},\quad$
$\displaystyle
w_{t}=\Sigma_{t}\left(\bar{w}+\sum_{s=1}^{t-1}x_{s}y_{s}\right)$
$\displaystyle\underline{\tau}_{t}=\underline{\tau}+\dfrac{t}{2},\quad$
$\displaystyle\bar{{\tau}}_{t}=\bar{\tau}+\dfrac{1}{2}\sum_{s=1}^{t-1}\left(y_{s}^{2}+\bar{w}^{\top}\bar{w}-w_{s}^{\top}\Sigma_{t}^{-1}w_{s}\right).$

Accordingly, the Bayes-optimal predictor becomes

$y_{t}^{*}(H_{t})=\mathbb{E}[y_{t}|H_{t}]=w_{t}^{\top}x_{t},$

and

$\displaystyle\sigma_{t}^{*2}(H_{t})$
$\displaystyle=\mathbb{E}[(y_{t}-y_{t}^{*}(H_{t}))^{2}|H_{t}]=\mathbb{E}[(f(x_{t})-y_{t}^{*}(H_{t}))^{2}|H_{t}]+\mathbb{E}[\sigma^{2}|H_{t}]$
$\displaystyle=\dfrac{\bar{\tau}_{t}}{\underline{\tau}_{t}-1}\cdot(\mathrm{tr}\left(x_{t}x_{t}^{\top}\Sigma_{t}\right)+1).$

These formulas are used to generate the Bayes-optimal curves in the figures.


Audio Generation | 164m | 7.50 | 9 | ✓ | - | - | ✓ | ✓ | ✓ | ✓ | - | - | - | - | ✓
Audio from Text | 19m | 5.62 | 8 | ✓ | - | - | - | - | - | ✓ | - | - | - | - | ✓
Audio from Video | 145m | 1.88 | 1 | ✓ | - | - | ✓ | ✓ | ✓ | ✓ | - | - | - | - | ✓
Image Understanding | 53m | 17.8 | 73 | ✓ | ✓ | - | ✓ | - | ✓ | - | ✓ | - | - | - | -
VQA | 5.8m | 6.23 | 31 | ✓ | - | - | ✓ | - | - | - | ✓ | - | - | - | -
Image Captioning | 32m | 4.25 | 14 | ✓ | - | - | ✓ | - | - | - | ✓ | - | - | - | -
Region Classification | 6.1m | 2.41 | 4 | ✓ | ✓ | - | ✓ | - | - | - | ✓ | - | - | - | -
Image Tagging | 3.8m | 2.38 | 8 | ✓ | - | - | ✓ | - | - | - | ✓ | - | - | - | -
Relationship Prediction | 0.8m | 1.41 | 6 | ✓ | ✓ | - | ✓ | - | - | - | ✓ | - | - | - | -
Region Captioning | 3.5m | 0.60 | 1 | ✓ | ✓ | - | ✓ | - | - | - | ✓ | - | - | - | -
Image Instruction Following | 0.4m | 0.37 | 6 | ✓ | - | - | ✓ | - | - | - | ✓ | - | - | - | -
Image Pair QA | 0.1m | 0.17 | 3 | ✓ | - | - | ✓ | - | ✓ | - | ✓ | - | - | - | -
Image Sparse Labelling | 13m | 7.25 | 26 | ✓ | ✓ | - | ✓ | - | ✓ | - | - | ✓ | - | ✓ | -
Object Detection | 5.3m | 3.08 | 9 | ✓ | - | - | ✓ | - | - | - | - | ✓ | - | - | -
Object Localization | 6.0m | 1.31 | 3 | ✓ | - | - | ✓ | - | - | - | - | ✓ | - | - | -
Referring Expression | 0.2m | 1.08 | 7 | ✓ | - | - | ✓ | - | - | - | - | ✓ | - | - | -
3D | 1.0m | 1.00 | 2 | ✓ | - | - | ✓ | - | ✓ | - | - | ✓ | - | ✓ | -
Text Detection | 37k | 0.41 | 3 | ✓ | - | - | ✓ | - | - | - | - | ✓ | - | - | -
Keypoint Detection | 0.3m | 0.38 | 2 | ✓ | ✓ | - | ✓ | - | - | - | - | ✓ | - | - | -
Image Dense Labelling | 6.9m | 4.06 | 19 | ✓ | ✓ | - | ✓ | - | ✓ | - | - | - | ✓ | - | -
Semantic Segmentation | 2.4m | 1.23 | 4 | ✓ | - | - | ✓ | - | - | - | - | - | ✓ | - | -
Localized Segmentation | 3.2m | 1.17 | 3 | ✓ | ✓ | - | ✓ | - | - | - | - | - | ✓ | - | -
Surface Normal Estimation | 1.1m | 1.03 | 6 | ✓ | - | - | ✓ | - | - | - | - | - | ✓ | - | -
Referring Expression Segmentation | 0.1m | 0.47 | 3 | ✓ | - | - | ✓ | - | - | - | - | - | ✓ | - | -
Depth Estimation | 47k | 0.11 | 1 | ✓ | - | - | ✓ | - | - | - | - | - | ✓ | - | -
Optical Flow | 24k | 0.06 | 2 | ✓ | - | - | ✓ | - | ✓ | - | - | - | ✓ | - | -
Video Understanding | 13m | 10.6 | 24 | ✓ | - | - | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | - | -
Video Captioning | 9.1m | 3.75 | 3 | ✓ | - | - | ✓ | - | ✓ | ✓ | ✓ | - | - | - | -
Video Tagging | 1.1m | 3.75 | 6 | ✓ | - | - | ✓ | - | ✓ | ✓ | ✓ | - | - | - | -
Video Question Answering | 2.5m | 2.84 | 9 | ✓ | - | - | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | - | -
Video Instruction Following | 0.2m | 0.21 | 6 | ✓ | - | - | ✓ | - | ✓ | ✓ | ✓ | - | - | - | -
Video Sparse Labelling | 0.4m | 3.42 | 5 | ✓ | ✓ | - | ✓ | - | ✓ | ✓ | - | ✓ | - | - | -
Video Tracking | 0.2m | 2.50 | 3 | ✓ | ✓ | - | ✓ | - | ✓ | ✓ | - | ✓ | - | - | -
Video Action Localization | 0.2m | 0.61 | 1 | ✓ | - | - | ✓ | - | ✓ | ✓ | - | ✓ | - | - | -
Video Sound Localization | 2.5k | 0.31 | 1 | ✓ | - | - | ✓ | - | ✓ | ✓ | - | ✓ | - | - | -
Audio Understanding | 2.2m | 2.50 | 10 | ✓ | - | - | ✓ | ✓ | ✓ | - | ✓ | - | - | - | -
Audio Tagging | 2.1m | 1.25 | 5 | ✓ | - | - | ✓ | ✓ | ✓ | - | ✓ | - | - | - | -
Audio Captioning | 75k | 1.25 | 5 | ✓ | - | - | - | ✓ | - | - | ✓ | - | - | - | -
Natural Language | 11m | 25.0 | 17 | ✓ | - | - | - | - | - | - | ✓ | - | - | - | -
Text Instruction Following | 11m | 12.5 | 10 | ✓ | - | - | - | - | - | - | ✓ | - | - | - | -
Language Modeling | - | 12.5 | 7 | ✓ | - | - | - | - | - | - | ✓ | - | - | - | -
Embodied AI | 7.2m | 4.33 | 23 | ✓ | - | - | ✓ | - | ✓ | - | ✓ | ✓ | - | ✓ | -
Action Prediction | 4.3m | 3.37 | 12 | ✓ | - | - | ✓ | - | ✓ | - | ✓ | - | - | - | -
Next Frame/State Prediction | 1.3m | 0.33 | 2 | ✓ | - | - | ✓ | - | ✓ | - | ✓ | - | - | ✓ | -
Goal Generation | 0.7m | 0.33 | 3 | ✓ | - | - | ✓ | - | ✓ | - | - | - | - | ✓ | -
Embodied QA | 1.0m | 0.30 | 6 | ✓ | - | - | ✓ | - | ✓ | - | ✓ | ✓ | - | - | -
All Tasks | 775m | 100 | 227 | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓
Table 11: Instruction tuning training mixture. Due to the number of datasets
used, we group them by task and only show statistics for each group. The rate
shows the sampling rate, size shows the number of examples of iterating
through the data once, and datasets show the number of individual data sources
used for the tasks.

In this section, we provide additional details about the instruction tuning
data and individual tasks Unified-IO 2 supports. An overview of the
instruction tuning data is shown in Table 11. We show a visualization
including individual datasets in Figure 10. We sample broad categories of
tasks evenly and then generally sample individual datasets in proportion to
the square root of their size, although with some minor hand-engineered
adjustments to downweight noisy datasets or upweight very rare tasks.

###  D.1 Natural Language

For natural language data we use the mixture from FlanV2 [122], which in turn
includes data from Muffin [193], T0-SF [156], NIV2 [191], and CoT annotations,
as well data from Alpaca [142], Dolly [33], Open Assistant [99], and MDPP [8].
In addition, we continue pre-training on our unsupervised NLP mixture from our
fine-tuning stage to ensure the model does not forget information learned from
unsupervised data during the extensive instruction-tuning stage.

###  D.2 Image Generation

For text-to-image generation, we use the same image/text pairs we used during
pre-training, as well as localized narratives from Open Images [103] and
captions from COCO [115] and Visual Genome (VG) [102]. Our prompts for these
tasks specify that the image might be noisy or approximate for unsupervised
corpora (_e.g_. “Generate an image that roughly matches this text: {caption}”)
and give hints as to the style for supervised corpora (_e.g_. “What do you see
in this image? Plainly describe the individual element you observe.” for
localized narratives) to help disambiguate the stylistic differences between
the datasets. We use simple prompts (_e.g_. “Caption this image.”) for the
COCO captions.

We additionally train the model to generate images through view synthesis
[194, 40] as was done during pre-training. We also integrate data for image
editing [18, 217] and image editing based on various dense control signals
such as depth maps, edges, segmentation, etc. Following [145], and the
segmentation-based image generation from Unified-IO using data from COCO and
LVIS [62]. Finally, we train on inpainting by masking a region of an input
image that contains an object and training the model to generate the complete
image given the object name and location. We derive data for this task from
the object annotation data in COCO, Open Images, and VG.

During inference, we use top-p sampling, also known as nucleus sampling [75],
for generating images with the temperature $t=1.0$ and $p=0.95$. We also
enable classifier-free guidance [74] by replacing the prompt with the un-
informative prompt “An image of a random picture.” 10% of the time during
training. That prompt is then used as the classifier-free prompt with a
guidance scale of $\alpha=10.0$ during inference.

###  D.3 Audio Generation

Datasets for audio generation from text include AudioCaps [93], Clotho [47],
MACS [131], MusicCaps [2], and LJSpeech [85]. During training, we divided the
audio into 4-second-long segments and then generated one segment of the target
audio, giving both the text and any previous segments as input. We also train
on the next-frame prediction task, which aims to generate the audio for the
next frame in a video from YT-Temporal-1B [215].

Our prompts for these tasks specify the characteristics of target audio;
_e.g_., “Generate the sound/music based on the description: {caption}” for
natural sound and music, respectively, and “Speak: {passage}” for speech. We
use the same sampling method as the image generation, the top-p sampling with
the temperature $t=1.0$ and $p=0.95$. We do not use the classifier-free
guidance because it can lead to poor performance. When generating audio longer
than 4.08 seconds during inference, we generate an initial segment that is
4.08 seconds long and then extend it by generating additional segments using
previous audio segments as the audio history input.

###  D.4 Image Understanding

These tasks require generating text in response to a query about an image or a
pair of images. We use the data from M3IT [112] and MIMIC-IT [108, 107], as
well as a variety of other additional sources. For VQA, we add GQA [83],
TallyQA [1], OK-VQA [130], A-OKVQA [160], OCR-based VQA datasets [133, 165],
Visual Genome, ScienceQA [124], VCR [213] and VizWiz [67]. For image tagging
we add Caltech Birds [195], iNaturalist [180], Sun397 [197], and Places365
[224]. For region classification, we add examples derived from object
annotation from Open Images, VG, and COCO. We categorize datasets with open-
ended responses such as LLaVa [119], Visual Storytelling [82], and Visual
Dialog [36] as visual instruction following, and we categorize NLVR [171] and
the “spot the differences” tasks from MIMIC-IT as image pair QA. For image
pair QA tasks, we encode the second image in the image history modality.

We also add a grounded relationship prediction task using data from Visual
Genome and VSR [116] as well as image captioning using the same supervised
sources we use for image generation.

We again put stylistic hints in the prompts for these tasks. For example, in
VQA and captioning datasets, we specify to return a short answer (_e.g_.
“Answer this question very succinctly: {question}”), which we find is critical
to allow the model to produce longer, more natural responses when asked user
questions. Likewise, we roughly specify the kind of class to output for image
tagging, _e.g_., “”What is the scientific name of this animal?” for the
iNaturalist dataset.

###  D.5 Image Sparse Labelling

These tasks require outputting sparse coordinates based on an input image. We
use Open Images, Visual Genome, and COCO for object detection and
localization, which requires detecting all objects belonging to a specific
class and three COCO referring expression datasets [91, 209, 136] for
referring expressions.

In addition, we train on the OmniLabel [16] 3D detection dataset by generating
the projected 3D center, virtual depth, log-normalized box size, and rotation
of each 3D box, again by normalizing these values between 0 and 1 and then
encoding them using the location tokens. We also added the camera pose
prediction tasks using Objaverse objects that were used during pre-training.

We include 3 text detection datasets from COCO-Text [183], including finding
the bounding box of an input text string for multiple text strings or finding
and listing all text along with their bounding boxes in an image.

Lastly, we do keypoint detection using COCO pose data. For keypoint detection,
we input a bounding box around a person in the image and train the model to
return a list of keypoints for that person. During inference, we first
localize all people in the image and then use each returned bounding box as a
keypoint query to find that person’s keypoints. During training, the model
predicts “MISSING” for keypoints that are not visible (_e.g_. “right elbow:
MISSING”). During inference, we use a masking function over the model’s logit
to force it to guess a valid point for each keypoint since the keypoint metric
does not award points for correctly identifying a keypoint as being not
visible.

###  D.6 Image Dense Labelling

We do several image labeling tasks, including surface normal estimation on
FramNet [78], BlendedMVS [204] and Taskonomy [211], depth on NYU Depth [138],
and optical flow on Flying Chairs [44] and MPI Sintel [21].

We additionally train on several segmentation tasks: semantic segmentation
(segmenting a particular class), localization segmentation (segmenting an
object in an input bounding box), and referring expression segmentation
(segmenting an object matching a referring expression). Data comes from Open
Images, COCO, LVIS, and referring expressions from the COCO refexp datasets
[91, 209, 136]. To do instance segmentation, as needed for GRIT, we first do
localization on the target class and then perform localized segmentation on
each returned bounding box.

During inference, we do temperature sampling with a top-p of 0.95 as before,
but without classifier-free guidance. For segmentation, we find it beneficial
to increase the value of p to 0.97.

###  D.7 Video Understanding

These tasks require generating text in response to a query about a video. For
video captioning, we add VATEX [190] and MSR-VTT [199]. For action
classification (video tagging), we add UCF101 [168], Kinetics-710 [111],
Something-Something v2 [58] and EPIC-KITCHENS-100 [35]. We also use examples
from EPIC-KITCHENS-100 for action anticipation. For video question answering,
we add MSRVTT-QA [198], MSVD-QA [198], STAR [196] and M4-ViteVQA [221].
Lastly, we use examples from M3IT and MIMIC-IT for the video instruction
following.

To cover the visual content of the entire video with a small number of frames
(5), we use the segment-based sampling following [185]; we first divide the
video into five segments of equal duration and then randomly sample one frame
from each of the segments during training, and the middle frame at inference.
We use the first four frames as the image history input and the final frame as
the image input for action classification and video captioning. We empirically
found that using the third frame as the image input while using the other
frames as the image history input performs better for video question
answering.

We use similar prompts to those for image understanding tasks, _e.g_., “Write
a short description of this video.”, “The question {question} can be answered
using the video. A short answer is” and “What are they doing in this video?
Short answer:” in video captioning, video question answering, and video
tagging, respectively, for ensuring a short answer.

###  D.8 Video Sparse Labelling

We do single object tracking and spatial-temporal action localization on video
data. We train on YouTube-BB [151], LaSOT [50] and GOT-10k [79] by inputting
bounding boxes around a target object in each of previous frames and having
the model return the next location as a bounding box (“Anticipate the object’s
next location from all previous images and the location of the object in those
frames: {locations}.”). We also train the model on AVA [61] by inputting a
video snippet consisting of five frames and requiring the model to detect all
actions of humans appearing in the middle (third) frame of the video snippet
(“Given the temporal context from the video, detect all of the humans
performing actions in the image.”). Note that we provide the video snippet,
not a single video frame, because some of the actions require temporal context
to answer (_e.g_., stand and sit) correctly. We use the final/middle frame of
five consecutive frames in the video as the image input and the other frames
as the image history input for single object tracking and action localization,
respectively.

###  D.9 Audio Understanding

We train the model on audio tagging and audio captioning tasks. For audio
tagging, we add AudioSet [54], VGG-Sound [24], and MACS. For audio captioning,
we use the same datasets as text-to-audio generation, that is, AudioCaps,
Clotho, MACS, MusicCaps, and LJSpeech. For audio-visual action classification,
we train on Kinetics-Sounds [7] and VGG-Sound.

We again use stylistic hints in the prompts for these tasks. For example, we
specify the characteristics of target audio (_e.g_., “Describe the music.” and
“Transcribe the audio to text.” for MusicCaps and LJSpeech, respectively),
enforce a short answer (_e.g_., “What is this in the audio? Short answer:” and
“Give a short description of this audio.”), and specify the kind of class to
output for audio tagging, _e.g_., “This audio depicts a scene of a” for MACS.
We use the same prompts as video tagging for audio-visual action
classification.

We use the same sampling strategy as the video understanding; we sample five
audio segments with uniform intervals from the whole audio and use the
middle/final audio segment as the audio input while using the other segments
as the audio history input for audio classification and audio captioning,
respectively.

###  D.10 Embodied AI

Figure 11: Examples of input and target representations for embodied and robot
tasks.

While many robot manipulation tasks can be formulated by multimodal prompts
that interleave language and images or video frames, we use VIMA-Bench [87] to
evaluate the robot manipulation skills. We use the image input as the initial
observation of the environment and the image history for the images or videos
in the prompt. The text inputs, or the language instructions, also include
special tokens to explicitly express the interleaved multimodal prompt. The
action space consists of primitive actions of “pick and place” for tasks with
a suction cup as the end effector or “push” for tasks with a spatula. Both
primitive actions contain two poses and one rotation $\in\mathbb{R}^{3}$,
specifying the start and target states of the end effector.

With the action representation described in B.1, we seamlessly add large-scale
manipulation datasets Language Table [127], BridgeData V2 [184], and
FrankaKitchen [63] with the continuous control in both simulated and real-
world environments. The model directly predicts the next action as the text
target based on the current observation as image input, previous frames as
image history, and language instruction and previous actions as text inputs.

Due to the non-causality of the model and limited sequence length for the
image history, we only added the PointNav task from Habitat [157] Gibson
scenes for the navigation. The model is required to predict the next action,
with random augmentation for predicting the next position and rotation state,
based on the point goal (positions $\in\mathbb{R}^{2}$), visual observations,
and previous actions and states, if any.

###  D.11 Task Augmentation

In addition to these sources, we derive several additional tasks that use the
same supervised annotations as other tasks but require performing slightly
different functions. We call this task augmentation. The new tasks include
prompts that specify the desired output. These tasks serve to add diversity to
our instruction following data. We review the task augmentation we construct
below.

Segmentation. We build several augmentations of the segmentation tasks,
including (1) segmenting pixels belonging to one of a set of 2-4 categories,
possibly including categories that do not exist in the image, (2) segmenting
pixels belonging to a class and are within an input bounding box and (3) build
a map of pixels that do not belong to a set 1-4 classes. Prompts are designed
for these that state the requirement, _e.g_., “Show pixels that are part of
chair, paper and in <extra_id_289> <extra_id_871> <extra_id_781>
<extra_id_1156>”.

Prompt | Model Response
---|---
A video of a man (woman) saying Unified-IO 2 is a model that works with vision, language, audio, and action. |  ()
A video of a man playing guitar.  |

Table 12: Audio generation samples from the pre-trained model.

Figure 12: Future frame prediction samples from the pre-trained model. Given
the initial input image and action, the model can generate a plausible frame
that reflects the result of the action.

Detection and Referring Expression. For detection, localization, and referring
expressions, we also train the model to output various properties of the
output bounding boxes instead of the boxes themselves. Properties include the
width, height, area, left/right/top/bottom half, center coordinates, distance
from the left/right/top/bottom edge of the image, or the coordinates of
different corners of the bounding box. We also change the format of the output
bounding box (_e.g_., $[x_{1},y_{1},w,h]$ instead of
$[y_{1},x_{1},y_{2},x_{2}]$ format), and change whether the model labels the
boxes with the object category or not.

For detection, we train the model to detect any object belonging to a set of
1-4 classes. For referring expressions, we train the model to locate multiple
referring expressions from a single query. In this case, we sometimes train
the model to predict a property of both referenced boxes instead of outputting
the directly, for example, which box is the smallest, which is the largest,
the area of intersection, a box containing both boxes, etc.

Relationship Prediction. We train the model to list all relationships between
a particular object in the image and any other object. A bounding box and
category specify the target object. Similarly, we train the model to predict
all relationships between any instance of a particular class of objects and
any other object in the image.

Captioning. For captioning, we train the model to generate a caption that is
longer or shorter than a given character or word length or contains a
particular word or set of words. We also randomly require the caption to start
with a particular prefix. Again, these requirements are specified in the
prompt, for example, “Generate a caption longer than five words for this
image. Start your output with the text ‘My caption is:’”.

Surface Normal Estimation. For surface normal estimation, we train the model
to generate RGB images that encode the pixel orientation differently. This
includes changing which RGB channels correspond to the x, y, and z
orientations and only including a subset of those orientations. We also
include tasks that require specifying the x, y, and z orientation at a
particular point specified in the prompt using location tokens. Finally, we
include tasks requiring segmentation masks over pixels with particular
orientations, _e.g_., “Build a binary mask over surfaces with an upward
orientation”.

Figure 13: Image generation samples from the pre-trained model. Prompts from
left to right: 1) An image of an astronaut riding a horse in the forest. There
is a river in front of them with water lilies, Fantasy, HD. 2) A image of a
cat wearing sunglasses, HD. 3) A image of a black german shepherd wearing a
red beret, HD. 4) An image of a stop sign in a Fantasy style with the text
“1991”.

Embodied AI. We further augment the embodiment datasets with the video QA and
goal image generation tasks. The QA augmentation aims for the robot’s planning
and affordance. For example, given a robot video trajectory, the model is
supposed to predict the plan (caption), or whether a given action is
reasonable from the language instruction. Applying image editing in embodied
space, we further let the model generate the goal or subgoal images based on
the initial visual observation in the image input and the language prompt in
the text input. While recent works show that embodiment QA with VLM [46, 162]
and (sub-)goal generation with diffusion model [134] are effective in the
decision-making downstream tasks, our model combines the both augmentation
strategies.

##  Appendix E Experiment Details

|  | Categorization | Localization | VQA | Refexp | Segmentation | Keypoint | Normal | All
---|---|---|---|---|---|---|---|---|---
|  | ablation | test | ablation | test | ablation | test | ablation | test | ablation | test | ablation | test | ablation | test | ablation | test
0 | NLL-AngMF [11] | - | - | - | - | - | - | - | - | - | - | - | - | 49.6 | 50.5 | 7.2 | 7.1
1 | Mask R-CNN [70] | - | - | 44.7 | 45.1 | - | - | - | - | 26.2 | 26.2 | 70.8 | 70.6 | - | - | 20.2 | 20.3
2 | GPV-1 [65] | 33.2 | 33.2 | 42.8 | 42.7 | 50.6 | 49.8 | 25.8 | 26.8 | - | - | - | - | - | - | 21.8 | 21.8
3 | CLIP [146] | 48.1 | - | - | - | - | - | - | - | - | - | - | - | - | - | 6.9 | -
4 | OFA${}_{\texttt{LARGE}}$ [186] | 22.6 | - | - | - | 72.4 | - | 61.7 | - | - | - | - | - | - | - | 22.4 | -
5 | GPV-2 [89] | 54.7 | 55.1 | 53.6 | 53.6 | 63.5 | 63.2 | 51.5 | 52.1 | - | - | - | - | - | - | 31.9 | 32.0
5 | DINO + SAM [94, 139] | - | - | 66.0 | 66.0 | - | - | - | - | 60.2 | 60.1 | - | - | - | - | 18.0 | 18.0
6 |  Unified-IO${}_{\texttt{SMALL}}$ | 42.6 | - | 50.4 | - | 52.9 | - | 51.1 | - | 40.7 | - | 46.5 | - | 33.5 | - | 45.4 | -
7 |  Unified-IO${}_{\texttt{BASE}}$ | 53.1 | - | 59.7 | - | 63.0 | - | 68.3 | - | 49.3 | - | 60.2 | - | 37.5 | - | 55.9 | -
8 |  Unified-IO${}_{\texttt{LARGE}}$ | 57.0 | - | 64.2 | - | 67.4 | - | 74.1 | - | 54.0 | - | 67.6 | - | 40.2 | - | 60.7 | -
9 |  Unified-IO${}_{\texttt{XL}}$ | 61.7 | 60.8 | 67.0 | 67.1 | 74.5 | 74.5 | 78.6 | 78.9 | 56.3 | 56.5 | 68.1 | 67.7 | 45.0 | 44.3 | 64.5 | 64.3
9 | UIO-2${}_{\texttt{L}}$ | 70.1 | - | 66.1 | - | 67.6 | - | 66.6 | - | 53.8 | - | 56.8 | - | 44.5 | - | 60.8 | -
10 | UIO-2${}_{\texttt{XL}}$ | 74.2 | - | 69.1 | - | 69.0 | - | 71.9 | - | 57.3 | - | 68.2 | - | 46.7 | - | 65.2 | -
11 | UIO-2${}_{\texttt{XXL}}$ | 74.9 | 75.2 | 70.3 | 70.2 | 71.3 | 71.1 | 75.5 | 75.5 | 58.2 | 58.8 | 72.8 | 73.2 | 45.2 | 44.7 | 66.9 | 67.0

Table 13: GRIT results and additional baselines from the GRIT leaderboard.

###  E.1 Pre-training Visualization

In the main paper, we evaluate the effectiveness of our pre-training by
evaluating Unified-IO 2 quantitively on a variety of benchmarks. Here, we
qualitatively show the visualizations from the pre-trained
UIO-2${}_{\texttt{XXL}}$model. Table 12 shows audio generation from text (top)
and text + video (bottom). We can see the pre-trained model learns text-to-
speech synthesis through video pre-training, and the model can also synthesize
music that matches the video input. Figure 12 shows the future frame
prediction samples given the initial input image and action sequence. Figure
13 shows the image generation samples given prompts. The model has a good
understanding of different objects. However, it struggles to generate the
correct text from the given caption.

###  E.2 NLP Results

| HellaSwag | MMLU | Arc Easy | Arc Cha. | BoolQ
---|---|---|---|---|---
UIO-2${}_{\texttt{L}}$ | 39.4 | 28.4 | 41.8 | 26.2 | 66.6
UIO-2${}_{\texttt{XL}}$ | 49.9 | 29.7 | 49.5 | 31.3 | 72.8
UIO-2${}_{\texttt{XXL}}$ | 52.7 | 30.4 | 55.3 | 33.5 | 77.3
Open LLaMA 3B | 52.0 | 23.9 | 69.3 | 33.8 | 67.0
LLaMA 7B | 57.1 | 42.6 | 76.4 | 43.5 | 77.7
LLaMA 7B Chat | 57.7 | 47.6 | 74.4 | 44.0 | 80.7

Table 14: Results on NLP tasks.

We present results on a set of NLP tasks to evaluate the model’s language
understanding abilities. We evaluate using the EleutherAI LM-Eval harness
[51], tasks are evaluated zero-shot using the default prompts without any
adjustments aside from adding the [Text] [S] prefix used for all text
generation tasks. We evaluate on HellaSwag [214] and a selection of other
question answering benchmarks: MMLU [72], ARC [31], and BoolQ [30]. Results
are shown in Table 14. Baselines were evaluated in the same setting, _i.e_.,
zero-shot, with the default prompts, and using LM-Eval. Unified-IO 2 is
generally ahead of Open LLaMA 3B but behind LLaMA.

###  E.3 GRIT Details

We present GRIT results in more detail in Table 13. Notably, Unified-IO 2 is
the first unified model to pass the Masked R-CNN baseline for localization and
goes a long way toward closing the gap between SAM and unified models on
segmentation.

For GRIT VQA, looking at the scores from GRIT on different VQA subsets, we
find that Unified-IO 2 does better on the same-source subset (84.6 vs 58.5)
but worse on the new-source subset (57.7 vs 67.2). Same-source questions come
from VQA 2.0, and new-source questions come from VG, so the difference can be
attributed to the kinds of questions being asked. Qualitatively, it is hard to
understand why the scores differ on these subsets since the GRIT ablation
questions lack ground truth annotations. However, we notice the models often
produce different answers when faced with ambiguous questions (_e.g_. “What
color is black on the horse”, “hair” for Unified-IO _vs_. “mane” for Unified-
IO 2), so one possibility is that Unified-IO 2 does not match the VG answer
style as well as Unified-IO, which would likely be due to differences in the
kind of VQA training data the models were trained on.

For GRIT localization, we find the model can struggle with images with many
instances of the target class, particularly when using beam search. We
hypothesize that this is because the probability mass can get split between
many similar location tokens, resulting in EOS becoming the most probable
token even if its probability is low. As a solution, during inference, we only
output EOS if the EOS token itself has a probability of over 0.5, which we
find significantly improves the performance on crowded images. In rare cases,
we observe this leads to the model generating bounding boxes for the same
instance multiple times. As a solution, we apply Non-maximum suppression with
a higher threshold of 0.8 to remove these duplicates. We apply this inference
trick for localization and when doing the initial localization step for the
keypoint and segmentation tasks.

###  E.4 Multimodal Benchmark Details

Splits | Metrics |  UIO-2${}_{\texttt{XXL}}$ |  UIO-2${}_{\texttt{XL}}$ |  UIO-2${}_{\texttt{L}}$ |  [206] |  [207] |  [26]
---|---|---|---|---|---|---|---
Random | Accuracy ($\uparrow$) | 90.90 | 88.27 | 84.03 | 88.28 | 90.24 | 86.90
Precision ($\uparrow$) | 94.30 | 97.44 | 77.73 | 94.34 | 97.72 | 94.40
Recall ($\uparrow$) | 87.07 | 78.60 | 95.40 | 82.20 | 83.00 | 79.27
F1-Score ($\uparrow$) | 90.54 | 87.01 | 85.66 | 87.85 | 89.76 | 86.19
% Yes | 46.17 | 40.33 | 61.37 | 44.91 | 43.78 | 43.26
Popular | Accuracy ($\uparrow$) | 88.17 | 87.47 | 77.27 | 86.20 | 84.90 | 83.97
Precision ($\uparrow$) | 89.13 | 95.69 | 70.03 | 89.46 | 88.24 | 87.55
Recall ($\uparrow$) | 86.93 | 78.47 | 95.33 | 82.06 | 80.53 | 79.20
F1-Score ($\uparrow$) | 88.02 | 86.23 | 80.75 | 85.60 | 84.21 | 83.16
% Yes | 48.77 | 41.00 | 68.07 | 45.86 | 45.63 | 45.23
Adversarial | Accuracy ($\uparrow$) | 84.17 | 85.77 | 72.00 | 84.12 | 82.36 | 83.10
Precision ($\uparrow$) | 82.17 | 92.01 | 65.00 | 85.54 | 83.60 | 85.60
Recall ($\uparrow$) | 87.27 | 78.33 | 95.33 | 82.13 | 80.53 | 79.60
F1-Score ($\uparrow$) | 84.64 | 84.62 | 77.30 | 83.80 | 82.00 | 82.49
% Yes | 53.10 | 42.57 | 73.30 | 48.00 | 48.18 | 46.50

Table 15: Object hallucination benchmark POPE results, in comparison with
mPLUG-Owl2 [206], Ferret [207], and Shikra [26].

| Img | Video | All
---|---|---|---
InstructBLIP [34] | 58.8 | 38.1 | 53.4
VideoChat-7B [110] | 39.0 | 33.9 | 37.6
Otter-7B [108] | 42.9 | 30.6 | 39.7
Qwen-VL-7B [12] | 62.3 | 39.1 | 56.3
Qwen-VL-chat-7B [12] | 65.4 | 37.8 | 58.2
mPLUG-Owl2-7B [206] | 64.1 | 39.8 | 57.8
LLaVA-1.5-7B [118] | - | - | 58.6
LLaVA-1.5-13B [118] | 68.2 | 42.7 | 61.6
Unified-IO 2${}_{\texttt{L}}$ | 56.0 | 37.5 | 51.1
Unified-IO 2${}_{\texttt{XL}}$ | 64.1 | 45.6 | 60.2
Unified-IO 2${}_{\texttt{XXL}}$ | 65.7 | 46.8 | 61.8

Table 16: Results on SEED-Bench [106]. Our XXL model outperforms all 7B vision
language models and is even slightly better than the LLaVA-1.5 13B model.

We now provide the breakdown results for the evaluation-only multimodal
benchmarks, POPE [113] and SEED-Bench [106]. POPE is the object hallucination
benchmark, requiring ‘yes’ or ‘no’ answers. As shown in Table 15, our largest
model achieves the highest F1 score in all 3 dimensions. Interestingly,
smaller models favored ‘no’ responses, possibly due to a bias from negative
examples encountered during the instruction tuning phase. SEED-Bench offers
19k multiple-choice questions with human annotations for evaluating multimodal
models across 12 dimensions, including spatial (Image) and temporal (Video)
understanding. As shown in Table 16, our XXL model outperforms all other 7B
vision/video language models, and is even slightly better than the LLaVA-1.5
13B model. Notably, our XL (3B) model has already outperformed all other
counterparts in the temporal understanding split. While recent video language
models [128, 110, 108] have shown proficiency in conventional video tasks like
video tagging and captioning, their performance in SEED-Bench’s temporal
understanding is even worse than that of vision language models, which might
be attributed to their limited instruction-following capabilities.

###  E.5 Image Generation Details

Figure 14: Samples generated by UIO-2${}_{\texttt{XXL}}$ for the MS COCO
captions [115]. While the classifier-free guidance [74] significantly boosts
image quality and fidelity, it achieves poor FID.

Figure 15: Samples generated for the TIFA bnechmark captions [76]. Some of the
images generated by baselines (_e.g_., rows 1-2, 6, 9) have high quality but
do not fully follow the input text while Unified-IO 2 generates faithful
images.

Prompt | Model Response
---|---
What is the sound of a duck quacking? |
A bell is ringing loudly and quickly. |
Speak out the text “Police protection was better and more effective;” for me. |
Speak: Printing, then, for our purpose, may be considered as the art of making books by means of movable types. |
Read “Many animals of even complex structure which live parasitically within others are wholly devoid of an alimentary cavity.” |
Generate the music based on the description “Slow tempo, bass-and-drums-led reggae song.” |
Based on the given description “Industrial techno sounds, repetitive, hypnotic rhythms”, produce a corresponding piece of music. |

Table 17: Audio generation examples. Unified-IO 2 can generate not only
environmental sound (rows 1-2), but also speech (rows 3-5) and music (rows
6-7). Note that some of the outputs longer than 4.08 seconds have
discontinuity in sound, or changes in tone, speed or melody (rows 4-5, 7).
Since our model can output 4.08-second audio at a time, we complete the audio
clip by using any previously generated clips as additional input. Click for
audio samples.

Figure 15 shows generated images for the TIFA benchmark captions [76] using
several baselines as well as UIO-2${}_{\texttt{XXL}}$. We use the official
implementation code (Emu [172] and CoDi [174]) or the images shared in the
official GitHub repository of TIFA555https://github.com/Yushi-
Hu/tifa/tree/main/human_annotations (Stable Diffusion v1.5 [154] and
miniDALL-E [37]) for baselines. All the baselines except miniDALL-E use the
Stable Diffusion decoder trained on large-scale, high-quality image datasets,
generating images of high fidelity. However, they often generate images that
do not fully follow the input captions while Unified-IO 2 generates faithful
images.

For text-to-image generation on MS COCO [115], we follow the standard
convention [226]; we evaluate on a subset of 30K captions sampled from the
validation set.666We use the evaluation code at
https://github.com/MinfengZhu/DM-GAN Following [43], we generate 8 images for
each caption and select the best one using CLIP text-image similarity [146].
Despite classifier-free guidance [74] resulting in generated images of
qualitatively higher quality, the computed FID score [73] is significantly
worse compared to what would have been achieved without employing it (33.77 vs
13.39); see Figure 14.

###  E.6 Audio Generation Details

Input Image | Prompt | Model Response
---|---|---
| What is the sound of this instrument? |
| What is the sound of this instrument? |
| Generate music about this scene. |
| Generate music about this scene. |
|  Locate the bounding boxes of the sound sources in the given image. |
|  Identify the locations of the sound sources in the given image. |
|  Identify the locations of the instruments producing the given sound. |
|  Identify the locations of the instruments producing the given sound. |

Table 18: Audio-visual qualitative examples showcasing the ability of Unified-
IO 2 to reason across modalities. Unified-IO 2 can generate the sound of the
instrument in the input image (rows 1-2), and generate the music that matches
the mood of the input image (rows 3-4). The last four examples (rows 5-8) show
the results of visual sound localization. Note that Unified-IO 2 can
accurately identify the instruments that make and do not make sounds (rows
7-8). Click for audio samples.

For text-to-audio generation, we evaluate on the AudioCaps [93] test set. Note
that we cannot do an apples-to-apples comparison with other methods because
AudioCaps consists of 10-second audio clips while our model can generate
4.08-second audio at a time. Instead, we evaluate the dataset in the following
setup: we first sample four 2-second audio segments, convert them to log-mel-
spectrograms with zero-padding, and generate the following audio with the
prompt “Generate the following sound based on what you heard and the
description: {caption}”. We convert the model output, that is, a log-mel-
scaled spectrogram, into a waveform using the pretrained HiFi-GAN, and compare
the ground-truth audio and generated audio using computational metrics
including Fréchet Audio Distance [92], Inception Score [155] and
Kullback–Leibler divergence. We use the same evaluation code as
AudioLDM777https://github.com/haoheliu/audioldm_eval [117]. We show the audio
generation examples in Table 17 and audio-visual qualitative examples in Table
18.

###  E.7 Video and Audio Understanding Details

Figure 16: Single object tracking examples from LaSOT [50]. Given the first
input image and the corresponding bounding box of the target object, our model
can track the object across video frames.

Figure 17: Video understanding qualitative examples.

We consider classification and question-answering tasks as open-ended answer
generation and use the Exact Match (EM) to measure the performance. We also
tried to formulate the classification task as multiple-choice answering and
generate answers by computing the logit for each dataset label and selecting
the one with the highest logit, but the performance boost was quite marginal.
Note that we do not train our model directly on the Kinetics-400 [90]; we
instead train on Kinetics-710, a mixture of three different datasets belonging
to the Kinetics family, that is, Kinetics-400, 600, and 700. Our model
achieves top-1 accuracy 79.1 (_vs_. instruction tuning only: 73.8) when
further finetuning on Kinetics-400 for 5 epochs, following [111]. For
Kinetics-Sounds, leveraging both audio and visual inputs largely improves
performance (audio-visual: 89.3 _vs_. video-only: 87.4 _vs_. audio-only:
38.2). For captioning tasks, we use CIDEr [182] as the evaluation metric.
Figure 17 shows the qualitative examples for video understanding tasks.

###  E.8 Embodiment Details

Figure 18: Future state prediction examples on robotic manipulation tasks.
Given the input image and instructions, our model can successfully generate
the target state after the prompt instruction.

| L1 | L2 | L3 | L4 | Avg.
---|---|---|---|---|---
VIMA [87] | 81.5 | 81.5 | 78.7 | 48.6 | 72.6
VIMA-Gato [152] | 57.0 | 53.9 | 45.6 | 13.5 | 42.5
VIMA-Flamingo [5] | 47.4 | 46.0 | 40.7 | 12.1 | 36.6
VIMA-GPT [19] | 46.9 | 46.9 | 42.2 | 12.1 | 37.0
Unified-IO 2${}_{\texttt{L}}$ | 66.9 | 63.8 | 57.5 | 12.6 | 50.2
Unified-IO 2${}_{\texttt{XL}}$ | 70.3 | 69.8 | 64.5 | 13.1 | 54.2
Unified-IO 2${}_{\texttt{XXL}}$ | 71.3 | 70.4 | 68.0 | 15.5 | 56.3

Table 19: Evaluations on VIMA-Bench [87]

In VIMA-Bench [87], there are 4 levels of evaluation protocols: L1 object
placement, L2 novel combination, L3 novel object, and L4 novel task. Results
and comparisons are shown in Table 19. The inputs for the autoregressive
transformer model VIMA [87] are object tokens consisting of cropped images and
bounding boxes; image patch tokens encoded by ViT for VIMA-Gato [152]; image
patch tokens encoded by ViT, further downsampled by a perceiver module for
VIMA-Flamingo [5]; and single image token encoded by ViT for VIMA-GPT [19].
The output of those baselines is all next-step action prediction. Since our
model has to predict all actions at once only with the initial observation,
the task setting is then more challenging than the casual policy learning
baselines. Nevertheless, our models still outperform counterparts that input
image or image patches for all 4 levels and are only behind the object-centric
method [87]. In Figure 18, we show the future state prediction examples on
robotic manipulation tasks. Given the input state image and natural language
prompt, our model can successfully synthesize the target image state.

###  E.9 Other Tasks

Figure 19: 3D object detection qualitative examples from Objectron [3].

Figure 20: Image-based 3D view synthesis examples from Objaverse [40].

Figure 16 shows single object tracking examples from the LaSOT [50] dataset.
Note that Unified-IO 2 does not use specific class labels for tracking and
tracks small moving objects such as a table tennis paddle well. Figure 19
presents qualitative examples of 3D object detection from the Objectron
dataset [3]. As outlined in our main paper, Unified-IO 2 exhibits suboptimal
performance in benchmarks for multi-object 3D detection. Additionally, Figure
20 illustrates examples of image-based 3D view synthesis using the Objaverse
dataset [40]. While the model produces coherent results, it faces challenges
in accurately representing relative camera transformations.