Task-Customized Self-Supervised Pre-training with Scalable Dynamic Routing

TL;DR

This paper introduces Scalable Dynamic Routing (SDR), a self-supervised learning framework that trains multiple sub-networks on different data subsets, enabling task-specific models that outperform traditional unified models on various downstream tasks.

Contribution

The paper proposes SDR, a novel SSL paradigm that trains multiple sub-nets for task-specific pre-training, improving transfer performance and efficiency across diverse downstream tasks.

Findings

SDR trains 256 sub-nets on ImageNet simultaneously.

SDR achieves state-of-the-art accuracy on 11 downstream classification tasks.

SDR outperforms a unified model trained on full ImageNet.

Abstract

Self-supervised learning (SSL), especially contrastive methods, has raised attraction recently as it learns effective transferable representations without semantic annotations. A common practice for self-supervised pre-training is to use as much data as possible. For a specific downstream task, however, involving irrelevant data in pre-training may degenerate the downstream performance, observed from our extensive experiments. On the other hand, for existing SSL methods, it is burdensome and infeasible to use different downstream-task-customized datasets in pre-training for different tasks. To address this issue, we propose a novel SSL paradigm called Scalable Dynamic Routing (SDR), which can be trained once and deployed efficiently to different downstream tasks with task-customized pre-trained models. Specifically, we construct the SDRnet with various sub-nets and train each sub-net with only one subset of the data by data-aware progressive training. When a downstream task arrives, we route among all the pre-trained sub-nets to get the best along with its corresponding weights. Experiment results show that our SDR can train 256 sub-nets on ImageNet simultaneously, which provides better transfer performance than a unified model trained on the full ImageNet, achieving state-of-the-art (SOTA) averaged accuracy over 11 downstream classification tasks and AP on PASCAL VOC detection task.