LegoNet: Piecing Together and Breaking Apart Sub-Networks for Scalable Multi-task Learning

Abstract

Despite considerable progress in general-purpose vision models, most efforts focus on designing a new unified structure that can handle different types of input and supervision. In contrast, we believe each vision task requires its specific designed module to use different forms of perception. For example, a feature pyramid network is commonly used in segmentation but not in classification. We present LegoNet, a general Multi-Task Learning (MTL) framework that is assembled with many small sub-networks from different vision tasks, similar to how Lego pieces can be pieced together into larger structures. By leveraging this property, LegoNet can borrow design elements from single-task models and combine them to create a scalable multi-task model. We demonstrate its efficiency on mainstream vision datasets such as ImageNet, COCO, and ADE20K, and show it achieves comparable results to state-of-the-art single-task models. Moreover, like a Lego creation capable of dynamically piecing together or breaking apart pieces, our model exhibits scalability in both its model capacity and adaptability to a multitude of tasks. It can remove sub-networks and decompose into high-performing components for efficient adaptation, or add sub-networks for learning new tasks in a continuous learning scenario. On downstream tasks, it can be fine-tuned with fewer training parameters, fewer model parameters, and even transformed to a low computation shape. These functions can be controlled and combined to meet various demands of downstream applications.