An Efficient General-Purpose Modular Vision Model via Multi-Task Heterogeneous Training

TL;DR

This paper introduces a scalable, multi-task vision transformer model trained on heterogeneous datasets, achieving high performance across diverse tasks with modularity for efficient adaptation and continual learning.

Contribution

It proposes a modified mixture-of-experts vision transformer capable of multi-task learning on diverse datasets, addressing heterogeneity challenges and enabling efficient downstream task adaptation.

Findings

Achieves comparable results to single-task models on multiple vision tasks.

Demonstrates strong generalization and modularity for downstream applications.

Enables efficient fine-tuning with fewer parameters and less computation.

Abstract

We present a model that can perform multiple vision tasks and can be adapted to other downstream tasks efficiently. Despite considerable progress in multi-task learning, most efforts focus on learning from multi-label data: a single image set with multiple task labels. Such multi-label data sets are rare, small, and expensive. We say heterogeneous to refer to image sets with different task labels, or to combinations of single-task datasets. Few have explored training on such heterogeneous datasets. General-purpose vision models are still dominated by single-task pretraining, and it remains unclear how to scale up multi-task models by leveraging mainstream vision datasets designed for different purposes. The challenges lie in managing large intrinsic differences among vision tasks, including data distribution, architectures, task-specific modules, dataset scales, and sampling strategies. To address these challenges, we propose to modify and scale up mixture-of-experts (MoE) vision transformers, so that they can simultaneously learn classification, detection, and segmentation on diverse mainstream vision datasets including ImageNet, COCO, and ADE20K. Our approach achieves comparable results to single-task state-of-the-art models and demonstrates strong generalization on downstream tasks. Due to its emergent modularity, this general-purpose model decomposes into high-performing components, efficiently adapting to downstream tasks. We can fine-tune it with fewer training parameters, fewer model parameters, and less computation. Additionally, its modularity allows for easy expansion in continual-learning-without-forgetting scenarios. Finally, these functions can be controlled and combined to meet various demands of downstream tasks.