PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Task Self-Training for Learning General Representations

Golnaz Ghiasi, Barret Zoph, Ekin D. Cubuk, Quoc V. Le, Tsung-Yi Lin

arXiv:2108.11353v1cs.CV

TL;DR

Computer vision still lacks a single general model that works well across tasks, and broad multi-task annotation is difficult to scale. MuST trains one student from pseudo-labels generated by independent task-specific teachers. It matches or outperforms supervised and self-supervised approaches on most transfer tasks, and can improve already strong checkpoints.

  • Problem

    Learning one general vision model remains challenging, while collecting varied annotations for the same images is time-consuming and hard to scale.

    Source: §1 · p. 1

  • Method

    MuST independently trains task-specific teachers, has them pseudo-label a shared dataset, and jointly trains one multi-task student on those labels.

    Source: §1 · pp. 1–2

  • Results

    5 of 6 transfer tasks: MuST outperforms or matches supervised and self-supervised learning, except classification.

    Source: §6 · p. 8

  • Takeaways & Limitations

    MuST is a scalable pre-training method for general feature representations and can improve checkpoints trained with billions of examples.

    Source: §1 · p. 2, §6 · p. 8

    Teacher and student training hold model capacity and data augmentation constant, with dataset and labels designed as their main differences.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Despite the fast progress in training specialized models for various tasks, learning a single general model that works well for many tasks is still challenging for computer vision. Here we introduce multi-task self-training (MuST), which harnesses the knowledge in independent specialized teacher models (e.g., ImageNet model on classification) to train a single general student model. Our approach has three steps. First, we train specialized teachers independently on labeled datasets. We then use the specialized teachers to label an unlabeled dataset to create a multi-task pseudo labeled dataset. Finally, the dataset, which now contains pseudo labels from teacher models trained on different datasets/tasks, is then used to train a student model with multi-task learning. We evaluate the feature representations of the student model on 6 vision tasks including image recognition (classification, detection, segmentation)and 3D geometry estimation (depth and surface normal estimation). MuST is scalable with unlabeled or partially labeled datasets and outperforms both specialized supervised models and self-supervised models when training on large scale datasets. Lastly, we show MuST can improve upon already strong checkpoints trained with billions of examples. The results suggest self-training is a promising direction to aggregate labeled and unlabeled training data for learning general feature representations.

1. Introduction

MuST uses pseudo-labels from independently trained task specialists to train one multi-task student, addressing the difficulty of assembling large multi-task vision datasets. The introduction reports gains across transfer tasks, including improvements over supervised and self-supervised baselines and strong pretrained checkpoints.

1. Introduction

  • Scaling multi-task vision learning is difficult because collecting varied annotations for the same images is time-consuming and hard to scale.The authors propose pseudo-labeling as a way to enable large-scale multi-task feature learning.Source: §1 · p. 1
  • A student trained only on pseudo-labels preserves most of its specialized teacher’s transfer-learning performance, suggesting pseudo-labels can distill supervised knowledge.The teachers generate labels on unlabeled images.Source: §1 · p. 2
  • MuST trains independent task-specific teachers, uses them to pseudo-label a shared dataset, then jointly trains one multi-task student.The shared dataset can carry labels from teachers trained on different datasets and tasks.Source: §1 · p. 2
  • A shared ResNet-FPN backbone with task-specific heads absorbs knowledge from four task teachers and matches or exceeds them on transfer tasks.The teachers cover classification, semantic segmentation, object-box detection, and depth estimation.Source: §1 · p. 2
  • MuST representations outperform SimCLR by a large margin on segmentation and depth estimation.The paper evaluates transfer on six computer-vision tasks spanning recognition and 3D geometry estimation.Source: §1 · p. 2
  • 5 of 6 transfer tasks: MuST outperforms supervised JFT pre-training except classification, and it improves ALIGN checkpoints on detection, segmentation, and depth.MuST also outperforms SimCLR on ImageNet and JFT and by a large margin on segmentation and depth.Source: §1 · p. 2

2. Related Work

Prior work shows that features learned through single-task pre-training may not generalize across downstream tasks. MuST instead uses self-training to combine information from multiple tasks in one model.

Self-training and multi-task learning

  • MuST targets general representations by applying self-training to multi-task learning, whereas prior self-training work focused on single tasks.The approach builds on work showing that pseudo-label pre-training can improve transfer learning.Source: §2 · p. 3

Representation learning

  • Single-task supervised pre-training may not produce features general enough for varied downstream applications, and task-aligned pre-training can perform best for detection and semantic segmentation.Prior work also found that Instagram pre-training may benefit classification while harming localization.Source: §2 · p. 3

Multi-task learning

  • The ResNet-FPN design assigns top-level features to classification, multi-scale features to detection and instance segmentation, and high-resolution features to pixel-wise tasks.The cited figure distinguishes these task-specific feature scales by color.Source: Figure 3 · p. 3

3. Method

MuST transfers specialized teachers’ knowledge to unlabeled or partially labeled data, then trains a shared-backbone student on supervision spanning multiple tasks. The paper evaluates representations by fine-tuning them on downstream tasks.

3.1. Specialized Teacher Models

  • Four teachers provide classification, detection, segmentation, and depth pseudo labels, drawing on models trained on medium- or large-scale datasets or an open-source depth checkpoint.Detection and segmentation labels use a 0.5 score threshold; classification uses soft labels, while depth predictions are used directly.Source: §3.1 · pp. 3–4

3.2. Multi-Task Student Model

  • The student shares a backbone across classification, object detection, and pixel-wise prediction, with task-specific heads and separate heads for tasks within the same category.ResNet-FPN features are routed by task: C5 for classification, pyramid features for detection, and fused P2 features for pixel-wise prediction.Source: §3.2 · p. 4
  • Every image receives supervision for all tasks from supervised or pseudo labels, and MuST combines the task losses using weights that vary by experiment.ImageNet weights follow a batch-size and learning-rate scaling rule except for depth; JFT learns task weights during training.Source: §3.2 · p. 4
  • MuST supports cross-dataset training by using available supervised labels and generating pseudo labels for absent tasks; it uniformly samples across datasets and treats supervised and pseudo labels equally.This avoids task-specific undersampling caused by dataset-size imbalance because every task is labeled on every image.Source: §3.2 · pp. 4–5

3.3. Transfer Learning

  • Representational quality is evaluated by end-to-end fine-tuning on six downstream tasks, including surface-normal prediction, which no specialized teacher was trained for.Learning rates and training steps are swept to support fair comparisons between representations.Source: §3.3 · p. 5

4. Experiments

Across downstream tasks, MuST improves representations over supervised and self-supervised baselines, and adding pseudo labels from more tasks improves performance. Scaling to larger unlabeled datasets and bootstrapping from ALIGN checkpoints also yield reported gains.

4.2. Learning with Multi-Task Self-Training

  • Adding each new type of pseudo label yields strong improvement across all six downstream tasks.The tested additions are depth, segmentation, and detection pseudo labels alongside supervised ImageNet classification labels.Source: §4.2 · p. 5, Table 3 · p. 6
  • On ImageNet, MuST improves over state-of-the-art self-supervised learning on 4/6 downstream tasks and supervised learning on 5/6.The comparison uses the same pre-training dataset, ImageNet.Source: §4.2 · p. 5, Table 2 · p. 6
  • MuST achieves the best transfer-learning performance on 6/6 downstream tasks versus supervised multi-task training on ImageNet, COCO, and Objects365.MuST trains on supervised and pseudo labels across the three datasets.Source: §4.2 · p. 6, Table 5 · p. 7

4.3. Scaling Multi-Task Self-Training

  • On JFT’s 300M images, MuST using only unlabeled images outperforms supervised JFT representations on 5/6 downstream tasks and SimCLR on 5/6.The pseudo-label teachers indirectly use about 3.7M supervised images, compared with 300M labeled JFT images.Source: §4.3 · p. 6, Table 6 · p. 7
  • Fine-tuning an ALIGN checkpoint with MuST improves transfer performance on four of six downstream tasks, while CIFAR-100 and DIODE surface-normal performance is slightly worse.The paper also reports improvements over the ALIGN EfficientNet-L2 checkpoint on four downstream tasks.Source: §4.3 · pp. 7–8, Figure 5 · p. 8

5. Discussion

MuST performs better than SimCLR on ImageNet and JFT, and exceeds supervised JFT pre-training on five of six tasks. The discussion presents it as a scalable approach that combines labeled and unlabeled data, with transfer gains also reported for ALIGN checkpoints.

5. Discussion

  • At ImageNet scale, self-supervised learning can outperform supervised learning, while supervised learning performs better at JFT scale.The discussion raises whether more unlabeled data could compensate for the self-supervised approach’s JFT-scale gap.Source: §5 · p. 8
  • MuST improves transfer learning performance of the ALIGN EfficientNet-L2 checkpoint on four downstream tasks.The table caption describes the checkpoints as versatile and competitive with state-of-the-art models, but does not provide specific values.Source: Table 7 · p. 8
  • MuST outperforms SimCLR on both ImageNet and JFT, and supervised JFT pre-training on five of six tasks, except image classification.The discussion identifies these results as evidence of MuST’s potential to learn general features that improve with more unlabeled data.Source: §5 · p. 8
  • Both self-supervised learning and self-training scale with unlabeled data, while the discussion reports better general-feature performance for self-training at present.It identifies combining self-supervised learning and self-training as a promising direction.Source: §5 · p. 8
  • MuST can combine labeled and unlabeled data, while self-supervised learning relies on pretext-task generalization and may not perform well on downstream tasks.The discussion says pseudo labels are easier to design when downstream tasks of interest are known in advance, and MuST generalizes to unseen surface-normal prediction given the teacher tasks used in this paper.Source: §5 · p. 8

6. Conclusion

The paper presents MuST as a scalable pre-training method for general visual representations. Its transfer results favor or match supervised and self-supervised approaches on five of six tasks, and it improves already strong checkpoints.

6. Conclusion

  • MuST outperforms or matches supervised and self-supervised learning on five of six transfer tasks, except classification.The evaluation covers six datasets spanning visual recognition, localization, and 3D geometry prediction.Source: §6 · p. 8
  • MuST is presented as a scalable multi-task self-training method for learning general representations.The paper compares it with supervised and self-supervised approaches on ImageNet and JFT.Source: §6 · p. 8
  • MuST can improve upon already strong checkpoints trained with billions of examples.The paper concludes that multi-task self-training is a scalable pre-training method able to learn general feature representations.Source: §6 · p. 8

A. Details of Training and Evaluation Datasets

The training setup draws on classification, detection, segmentation, and depth sources, alongside JFT as a large-scale test setting. The depth teacher combines five datasets and is trained to be invariant to depth range and scale.

A.1. Training Datasets

  • The appendix introduces five datasets used to train teacher models.Source: §A.1 · p. 9
  • ImageNet provides classification training data with 1.2M images and 1,000 classes; its images are center-cropped and have one primary object each.ImageNet is one of the datasets used to train teachers.Source: §A.1 · p. 9
  • Objects365 supplies object-detection training data with 600k images across 365 classes.Source: §A.1 · p. 9
  • COCO contains 118k images with multiple label types, and the experiments use its panoptic-segmentation labels.Source: §A.1 · p. 9
  • MiDaS, used to generate depth pseudo labels, is trained on five datasets—DIML Indoor (220k images), MegaDepth (130k), ReDWeb (3,600), WSVD (1.5M), and 3D movies (75k)—and is invariant to depth range and scale across them.The paper says this produces robust pseudo labels.Source: §A.1 · p. 9
  • JFT is a large-scale multi-label classification dataset with 300M labeled images, used to test MuST and self-supervised methods at scale.Source: §A.1 · p. 9

A.2. Evaluation Datasets

The evaluation datasets cover classification, detection, semantic segmentation, depth estimation, and surface normals. Surface-normal prediction is held out as a training task to test representation generality.

A.2. Evaluation Datasets

  • Surface-normal prediction is not used as a training task, allowing the evaluation to test task generality.The dataset overview describes five training tasks across the evaluation datasets.Source: §A.2 · p. 9
  • CIFAR-100 provides classification evaluation with 50k images and 100 classes.Source: §A.2 · p. 9
  • PASCAL Detection uses VOC 2007 and 2012 trainval sets, totaling 16.5k images across 20 classes.Source: §A.2 · p. 9
  • PASCAL Segmentation uses the VOC 2012 training set, with 1.5k images and 20 classes.Source: §A.2 · p. 9
  • NYU Depth V2 contains 47,584 training images and 654 validation images for depth estimation.Source: §A.2 · p. 9
  • ADE20k provides segmentation evaluation with 20k images and 150 object and stuff classes across indoor and outdoor scenes.Source: §A.2 · p. 9
  • DIODE contains 16,884 diverse indoor and outdoor images; the evaluation uses its surface-normal labels.Source: §A.2 · p. 9

B.1. Training Teacher Models

The teacher models use ResNet-152 backbones, with task-specific training setups and reported final performance for classification, detection, and segmentation. The depth teacher is a pre-trained MiDaS model.

B.1. Training Teacher Models

  • All teacher models use a ResNet-152 backbone.The teachers generate the pseudo labels used in MuST.Source: §B.1 · p. 9
  • 26.1 AP for Objects365 detection and 53.8 mIoU for COCO semantic segmentation are the reported final performances.The COCO model predicts semantic classes per pixel, not object instances.Source: §B.1 · p. 9
  • The depth teacher is a pre-trained MiDaS model downloaded from its GitHub repository.It generates pseudo labels for depth estimation.Source: §B.1 · p. 9
  • 81.6 top-1 accuracy is the reported final performance of the ImageNet classification teacher.It is trained for 400 epochs with random augmentation and label smoothing.Source: §B.1 · p. 10

B.2. Training Multi-Task Student Models

Student models use a shared training setup, with weight decay varying between ImageNet and JFT experiments.

B.2. Training Multi-Task Student Models

  • Student training uses batch size 256, image size 640, and scale jittering from 0.5 to 2.0.No random augmentation or label smoothing is applied.Source: §B.2 · p. 10
  • Weight decay is 4e-5 in ImageNet experiments and 3e-6 in JFT experiments.Source: §B.2 · p. 10
  • Student training applies no random augmentation or label smoothing.Source: §B.2 · p. 10

B.3. Fine-tuning on Evaluation Datasets

Evaluation fine-tunes pre-trained representations on multiple datasets, initializing the backbone and FPN from a pre-trained model and task-specific heads randomly. The accompanying visualization shows predictions from one multi-task student on ImageNet validation images.

B.3. Fine-tuning on Evaluation Datasets

  • Fine-tuning initializes the ResNet and FPN backbone from a pre-trained model, randomly initializes remaining layers, and uses end-to-end training with cosine learning rates.The authors grid-search learning rates and training steps; weight decay is set to 1e-4 after experiments found it made little difference.Source: §B.3 · p. 10
  • CIFAR-100 evaluation reports top-1 accuracy; the best SimCLR learning rate is 0.16 versus 0.005 for the supervised model.The passage says this higher-learning-rate trend also holds for subsequent tasks.Source: §B.3 · p. 10
  • PASCAL VOC segmentation is evaluated with mIoU after initializing the backbone and FPN from pre-trained layers and replacing the segmentation head.Source: §B.3 · p. 10
  • PASCAL detection reports AP50 with 11 recall points; its model uses pre-trained backbone and FPN layers with randomly initialized heads.Source: §B.3 · p. 10
  • NYU Depth evaluation reports the percentage of predictions within a 1.25 relative ratio of ground truth; DIODE reports angles below 11.25°.NYU uses 47,584 training and 654 validation images; DIODE uses 16,884 training and 446 validation images.Source: §B.3 · p. 10

C. Visualization of Student Model Predictions

Related papers

Loading 2108.11353v1…