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

Zitian Chen, Mingyu Ding, Yikang Shen, Wei Zhan, Erik Learned-Miller, Chuang Gan

Keywords: multi-task learning; continous learning; efficient adaptation

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.

Primary Area: transfer learning, meta learning, and lifelong learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 2023