PaperlayerView exact revision on arXiv

Source-linked AI summary

UniT: Multimodal Multitask Learning with a Unified Transformer

Ronghang Hu, Amanpreet Singh

arXiv:2102.10772v3cs.CVcs.CL

TL;DR

Prior transformer systems generally focused on individual domains or fine-tuned separate parameters for each task, leaving limited support for unified cross-domain multitask learning. UniT addresses this with modality-specific encoders, a shared decoder, and task-specific heads trained end-to-end, jointly learning 7 tasks across 8 datasets with strong performance and compact shared parameters. The experiments also identify dataset overlap handling and architecture-size trade-offs as relevant evaluation conditions.

  • Problem

    Prior transformer multitask efforts had limited cross-domain scope and commonly used task-specific fine-tuning rather than shared parameters across tasks.

  • Method

    UniT encodes each modality separately, applies a shared transformer decoder over encoded inputs, and uses task-specific output heads for end-to-end joint training.

  • Results

    UniT jointly addresses 7 tasks across 8 datasets with strong performance on each task using a compact set of shared parameters.

  • Takeaways & Limitations

    A domain-agnostic transformer can jointly handle visual perception, natural language understanding, and multimodal reasoning within one model.

  • Takeaways & Limitations

    Evaluation removes image overlap across datasets by restricting training and evaluation splits relative to COCO val2017.

Abstract

from arXiv · show

We propose UniT, a Unified Transformer model to simultaneously learn the most prominent tasks across different domains, ranging from object detection to natural language understanding and multimodal reasoning. Based on the transformer encoder-decoder architecture, our UniT model encodes each input modality with an encoder and makes predictions on each task with a shared decoder over the encoded input representations, followed by task-specific output heads. The entire model is jointly trained end-to-end with losses from each task. Compared to previous efforts on multi-task learning with transformers, we share the same model parameters across all tasks instead of separately fine-tuning task-specific models and handle a much higher variety of tasks across different domains. In our experiments, we learn 7 tasks jointly over 8 datasets, achieving strong performance on each task with significantly fewer parameters. Our code is available in MMF at https://mmf.sh.

1. Introduction

Transformers have succeeded across language, vision, and other modalities, but prior work had made limited progress toward one model handling tasks across domains. Existing approaches were typically restricted to one domain or modality and often used task-specific parameters.

  • Transformers have achieved strong results across language, image, video, and audio applications.
  • Prior work had not extensively connected different tasks across domains with transformers.
  • Earlier systems were generally limited to vision-only, language-only, or specific vision-language tasks.
  • Many existing approaches fine-tuned task-specific models instead of sharing parameters across tasks.

2. Related work

Related transformer work spans language, vision, and multimodal reasoning, but most systems train or fine-tune task-specific models. UniT instead targets simultaneous learning of multiple tasks across modalities in one shared transformer.

  • Transformers on language, vision, and multimodal tasks: Transformers have been applied to downstream language tasks through pretrained representations and fine-tuning.
  • Transformers on language, vision, and multimodal tasks: Transformer models support vision tasks including image classification, object detection, and panoptic segmentation.
  • Transformers on language, vision, and multimodal tasks: Vision-language transformers have been developed for multimodal reasoning tasks such as visual question answering and visual entailment.
  • Transformers on language, vision, and multimodal tasks: Most prior transformer applications train or fine-tune a separate model for each task, including downstream fine-tuning after multitask pretraining.
  • Multi-task learning with transformers: Prior multitask learning with transformers commonly focused on a specific domain or modality, though some work explored a single generic model across domains.
  • Contrast to multimodal pretraining: Unlike multimodal pretraining approaches that specialize models through downstream fine-tuning, UniT jointly handles all tasks in one shared model.

3. UniT: Unified Transformer across domains

UniT combines modality-specific encoders with a domain-agnostic transformer decoder and task-specific heads, supporting vision-only, language-only, and multimodal tasks. The model is trained and evaluated across seven tasks and eight datasets, with shared or separate decoders explored.

  • 3. UniT: Unified Transformer across domains: UniT uses separate image and text encoders followed by a shared or task-specific transformer decoder and simple task-specific heads.
  • 3.1. Image encoder: The image pathway applies a CNN backbone and transformer encoder to convert an image feature map into contextualized visual hidden states.
  • 3.1. Image encoder: A learned task embedding is added to the image encoder so its outputs can contain task-specific information.
  • 3.2. Text encoder: The text pathway uses pretrained BERT to encode tokenized inputs into a sequence of textual hidden states, with a learned task embedding prefixed to the input.
  • 3.2. Text encoder: Keeping only BERT’s [CLS] hidden vector works nearly equally well while saving computation.
  • 3.3. Domain-agnostic UniT decoder: The decoder consumes one encoded modality for unimodal tasks or concatenated image and text representations for multimodal tasks.
  • 3.3. Domain-agnostic UniT decoder: A task-specific query sequence produces decoded states through self-attention and cross-attention, using either one shared decoder or separate decoders.
  • 3.4. Task-specific output heads: Detection uses class, box, and optional attribute heads, whereas VQA, entailment, and language tasks use classifiers over the first top-layer decoder state.

4. Experiments

UniT is evaluated as a shared-decoder model across detection, VQA, language understanding, and visual entailment, using joint training over multiple datasets. It achieves strong cross-task performance while exposing trade-offs involving decoder sharing, task interactions, fine-tuning, and model capacity.

  • Experimental setup: UniT is tested on object detection, VQA, GLUE language understanding, and SNLI-VE visual entailment across multiple dataset combinations.Experiments include COCO and Visual Genome detection, VQAv2, four GLUE tasks, and SNLI-VE.
  • Detection and VQA: Shared-decoder joint training outperforms separately trained single-task models on all three datasets in the final detection-and-VQA experiment.The comparison is reported for Table 1 line 4 versus line 1.
  • Detection and VQA: The three-dataset shared-decoder variant outperforms single-task models after extending training, addressing the longer schedule required for detection and shared multimodal behavior.Detection requires structural outputs and object-relation modeling, while VQA requires multimodal fusion and reasoning.
  • Detection and VQA: Joint training on two detection datasets usually benefits both datasets, while Visual Genome detection benefits VQA more than COCO detection.The authors attribute the VQA difference to Visual Genome’s more diverse object annotations and broader visual-concept coverage.
  • Multiple domains: UniT jointly performs 7 tasks across 8 datasets with 8× fewer parameters than task-specific fine-tuned models; the final model has 201M parameters.The evaluated tasks span vision-only, vision-and-language, and language-only settings.
  • Multiple domains: Vision-and-language tasks consistently benefit from joint training, whereas unimodal tasks generally favor separate training under the evaluated settings.The authors relate this pattern to stronger unimodal representations benefiting multimodal reasoning, while noting iteration and parameter-count differences.
  • Multiple domains: Per-task fine-tuning notably improves object-detection mAP but has moderate or sometimes negative effects on other tasks, while increasing parameters eightfold and reducing generality.The trade-off also includes longer training.
  • Multiple domains: UniT approaches domain-specific transformer baselines with one generic model and shared hyperparameters across all 8 datasets, trained end-to-end in one step.The comparison includes DETR for detection, VisualBERT for vision-language reasoning, and BERT for GLUE tasks.

5. Conclusion

UniT applies a transformer framework across domains to jointly handle multiple tasks within one unified encoder-decoder model. It addresses 7 tasks across 8 datasets with shared parameters and strong performance, supporting a step toward general-purpose agents.

  • Conclusion: UniT jointly handles 7 tasks across 8 datasets in a single unified encoder-decoder model with shared parameters.The tasks span visual perception, natural language understanding, and multimodal reasoning.
  • Conclusion: The model learns all tasks in a single training step while achieving strong performance on each task with a compact parameter set.The conclusion characterizes the architecture as domain-agnostic.
  • Conclusion: UniT takes a step toward building general-purpose intelligence agents for applications across visual perception, language understanding, and multimodal reasoning.

A. Hyper-parameters and details of UniT

UniT’s joint training requires careful optimizer handling for task-specific parameters. Updating unused parameters can cause divergence, while skipping those updates stabilizes training in the reported setting.

  • UniT hyper-parameters and dataset sampling probabilities are summarized in Tables A.1 and A.2.
  • Updating unused parameters caused divergence when jointly training COCO detection, VG detection, and VQAv2 with a shared decoder.VQA accuracy stayed around 25% in this divergent setting.
  • Skipping optimizer updates for unused parameters allowed the model to converge properly in the shared-decoder setting.
  • Lowering detection sampling probabilities also avoided VQA divergence but produced lower detection mAP than skipping updates on unused parameters.

B. Multitask learning in UniT

UniT jointly handles diverse tasks in one shared model rather than maintaining separately fine-tuned task-specific models. It achieves comparable performance with substantially fewer parameters, while task interactions produce both benefits and competition.

  • UniT is a joint multi-task model across several domains with comparable performance to per-task models and 8× fewer parameters.
  • A shared UniT model avoids storing eight separately fine-tuned models, which would require 8× the total parameters.
  • The multi-task model performs better on VQAv2 and SNLI-VE but does not outperform separately trained models on pure vision or pure language tasks.
  • UniT keeps object-detection and language knowledge in one shared model instead of losing it through specialization to other tasks.
  • UniT’s ability to jointly solve tasks across domains is presented as a critical step towards general intelligence.
  • Joint training spans diverse tasks, including object detection and sentiment analysis, without requiring strict task compatibility.
  • Task compatibility still varies because joint learning provides shared supervision while tasks compete for finite model capacity.

C. Additional ablation results

Additional ablations show that UniT is sensitive to architecture, initialization, optimization, training duration, and encoder updating. Several configurations reveal trade-offs between detection and language or reasoning performance.

  • Increasing the image encoder hidden size from 256 to 768 leads to noticeably lower detection performance, possibly from overfitting detection features.
  • Initializing the convolutional backbone from ImageNet classification produces lower COCO detection mAP than using a detection-pretrained ResNet-50.The authors suspect the shorter effective COCO training schedule contributes to this result.
  • Using one decoder query slightly lowers SNLI-VE accuracy, while increasing the count to 100 does not improve over the default 25 queries.
  • A learning rate of 1e-4 causes training divergence, whereas 1e-5 lowers COCO detection mAP but raises SNLI-VE and MNLI accuracies.
  • Using 2× more training iterations raises COCO detection mAP but lowers MNLI accuracy, reflecting different training demands across tasks.
  • Initializing from a COCO single-task model raises detection mAP but slightly reduces performance on the other two datasets.
  • Freezing both modality encoders yields significantly lower performance on COCO, SNLI-VE, and MNLI, while updating the text encoder improves the language-involving tasks.
  • The ablations suggest that co-adaptation between encoders and the shared decoder is critical to multi-task training.

D. Learning curves

Learning curves show generally increasing performance across UniT’s multi-task training, with small-dataset language tasks peaking early and declining slightly afterward.

  • Figure D.1 plots validation performance against training iterations for UniT with shared or separate decoders across eight datasets.Single-dataset models are included for reference.
  • Most tasks in the multi-task models improve monotonically during training.
  • SST-2 and QNLI accuracies peak early and slightly decline later, likely because these relatively small datasets are overfit.

E. More visualizations

These visualizations show UniT’s validation learning curves and additional predictions from a shared-decoder model across eight datasets. The same model is applied across tasks and datasets.

  • The same UniT model is applied to each task and dataset represented in the visualizations.
  • The learning curves compare shared-decoder, separate-decoder, and single-task training across the experiments.Panels (a) and (b) cover decoder-sharing variants, while panels (c–j) show single-task training for each task.
  • Additional predictions from the shared-decoder model are shown across eight datasets.These examples correspond to Table 3 line 5 in the main paper.
Loading 2102.10772v3…