PaperlayerView exact revision on arXiv

Source-linked AI summary

PolyViT: Co-training Vision Transformers on Images, Videos and Audio

Valerii Likhosherstov, Anurag Arnab, Krzysztof Choromanski, Mario Lucic, Yi Tay, Adrian Weller, Mostafa Dehghani

arXiv:2111.12993v1cs.CVcs.LG

TL;DR

Separate task-specific models are common, motivating the question of whether one transformer can share parameters across modalities and datasets. PolyViT co-trains tasks across images, video, and audio; it reports strong benchmark performance with substantially fewer parameters and generalizing representations.

  • Problem

    Current state-of-the-art methods train separate models with different parameters for each task, raising the question of whether one transformer can share parameters across modalities and datasets.

    Source: §1 · p. 1

  • Method

    PolyViT co-trains one model on multiple classification tasks across image, video, and audio modalities, processing one task at a time for each input.

    Source: §1 · p. 1

  • Results

    PolyViT reaches state-of-the-art results on three video and two audio datasets, while multimodal co-training remains competitive and learns representations that generalize across modalities.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The model can add new tasks by learning an additional output head, while reusing shared parameters across tasks.

    Source: §6 · p. 8

    PolyViT does not co-train on large-scale upstream datasets, improve inference time, or fuse modalities such as video and audio for a joint prediction.

    Source: §6 · p. 8

Abstract

from arXiv · show

Can we train a single transformer model capable of processing multiple modalities and datasets, whilst sharing almost all of its learnable parameters? We present PolyViT, a model trained on image, audio and video which answers this question. By co-training different tasks on a single modality, we are able to improve the accuracy of each individual task and achieve state-of-the-art results on 5 standard video- and audio-classification datasets. Co-training PolyViT on multiple modalities and tasks leads to a model that is even more parameter-efficient, and learns representations that generalize across multiple domains. Moreover, we show that co-training is simple and practical to implement, as we do not need to tune hyperparameters for each combination of datasets, but can simply adapt those from standard, single-task training.

1. Introduction

PolyViT co-trains a shared transformer across image, video, and audio classification tasks, aiming to replace separate task-specific models. The approach remains competitive while reducing parameters and supporting practical deployment.

1. Introduction

  • PolyViT shares model parameters across tasks and modalities, beyond using a common architecture for different modalities.Its overview describes a shared transformer encoder alongside modality-specific input tokenizers and task-specific output heads.Source: §1 · p. 1, Figure 1 · p. 2
  • Co-training trains one model on multiple classification tasks, potentially across modalities, while processing one task at a time for each input.The introduction considers settings with as many as nine image, video, and audio tasks.Source: §1 · p. 1
  • Across nine datasets, PolyViT uses 8.3 times fewer parameters with at most a 1.2% accuracy drop versus state-of-the-art single-task baselines.The model learns representations that generalize across multiple tasks and domains.Source: §1 · p. 1
  • Single-modality co-training improves accuracy on each task and achieves state-of-the-art results on five video and audio classification datasets.The authors attribute the gains, especially on smaller datasets, to a regularizing effect that reduces overfitting.Source: §1 · p. 1
  • Adding capabilities can require training an additional linear classifier rather than maintaining a separate model for every task.The paper notes the resulting deployment and online-update advantages of maintaining one model.Source: §1 · pp. 1–2
  • Co-training reuses hyperparameters adapted from standard single-task training and does not increase overall training cost beyond the sum of the single-task steps.The implementation is based on Scenic, and the authors state that code and checkpoints will be released for reproducibility.Source: §1 · p. 2

2. Related work

Prior work spans multi-task learning, cross-modal transformers, and unified task frameworks, but often uses task- or modality-specific components. PolyViT instead shares its transformer parameters across tasks while processing one task per input.

2. Related work

  • Multi-task models have often shown lower accuracy than single-task models, with deficits increasing as tasks or task relatedness grow, and joint training commonly requires loss calibration.Prior remedies include gradient normalization or surgery and adaptive loss weights.Source: §2 · p. 2
  • PolyViT performs one task per input and alternates SGD updates across task batches; prior work often adds task-specific parameters to improve additional tasks.This setting was called “single-tasking of multiple tasks” in computer vision [52], while NLP uses “multi-task learning”.Source: §2 · p. 2
  • NLP work has co-trained unified models by converting tasks to a common framework that effectively concatenates datasets.The cited approaches are presented as related co-training setups, rather than as the same multimodal parameter-sharing setting.Source: §2 · p. 2
  • Perceiver handles different modalities but trains separate networks per task, whereas MBT fuses modalities for one task and also uses separate parameters across tasks.PolyViT instead shares almost all model parameters among tasks.Source: §2 · pp. 2–3
  • UniT shares only its decoder across vision-and-language tasks, while Bain et al. train one image-video transformer on a single retrieval task and dataset.The authors contrast both setups with simultaneous training across tasks and modalities using a shared transformer backbone.Source: §2 · p. 3
  • Earlier multimodal transformer approaches include cross-modal self-supervision using visual and audio features, with low-rank approximations to reduce parameters, and heterogeneous image-text models.The paper says the reported results of the heterogeneous approach were not competitive with the state of the art.Source: §2 · pp. 2–3

3. Preliminaries

The preliminaries define the paper’s input modalities and tasks, then describe transformer architectures used for images, audio, and video. Their related tokenization and initialization choices motivate a shared multimodal encoder.

3. Preliminaries

  • A modality is an input type, and a task is a modality paired with a dataset’s class set; the paper studies images, audio, and video classification.Each task corresponds directly to a dataset, such as ImageNet-1K for images or Kinetics 400 for video.Source: §3 · p. 3

3.1. Vision Transformers and extensions

  • ViT splits an image into non-overlapping patches, projects each patch into a token, and adds a class token and positional embeddings before encoding.The patch projection can also be viewed as a 2D convolution, and the transformer encoder processes the token sequence through sequential layers.Source: §3.1 · p. 3
  • AST applies the ViT architecture to log-mel spectrogram inputs, which can be tokenized like images.The best AST model described here was initialized from ViT models pretrained on large image datasets.Source: §3.1 · p. 3
  • Unfactorized ViViT extends ViT tokenization from 2D patches to 3D spatio-temporal tubelets, then projects them into tokens.The projection can be interpreted as a 3D convolution.Source: §3.1 · p. 3
  • ViT, ViViT, and AST commonly use large-scale image-pretrained initialization, with positional embeddings interpolated for the target sequence length.ViViT’s 3D embedding projection is initialized from ViT’s 2D projection; these architectural similarities support a shared encoder with separate input tokenizers.Source: §3.1 · p. 4

4. Co-training ViT on images, audio and video

PolyViT shares a transformer across image, audio, and video tasks while retaining modality-specific tokenization and task-specific output heads. Its co-training procedure offers several task schedules and can reuse standard single-task training settings and pretrained ViT weights.

4.1. PolyViT architecture

  • Modality-specific class tokens, input embeddings, and positional embeddings feed the shared backbone, while each task has a separate output head.Input embeddings, positional embeddings, and class tokens are shared across tasks within each modality; output heads are the only task-specific parameters.Source: §4.1 · p. 4
  • Optional modality-adaptor layers increase capacity by placing modality-specific transformer layers after tokenization and sharing the remaining layers.With L_adapt adaptor layers, L_shared = L − L_adapt layers are shared across modalities and tasks.Source: §4.1 · p. 4
  • Sharing the transformer encoder across tasks reduces parameters by approximately the number of tasks without changing inference time, since each forward pass handles one task.This comparison is with standard single-task training when all transformer layers are shared.Source: §4.1 · p. 4

4.2. Co-training procedure

  • Each SGD step samples one task and minibatch, then updates parameters; single-task minibatches allow standard baseline hyperparameters without additional tuning.This also avoids padding examples from different token-length modalities to a common sequence length.Source: §4.2 · p. 4
  • Task-by-task uses random contiguous task blocks, alternating cycles tasks, uniform sampling assigns equal task shares, and weighted sampling follows baseline step counts.The schedules specify how training steps are allocated among tasks.Source: §4.2 · pp. 4–5, Figure 2 · p. 5

4.3. Initialization of PolyViT

  • PolyViT is initialized from an ImageNet-21K-pretrained ViT; adaptor layers reuse its first layers, shared layers its last layers, and output heads start randomly.Modality-adaptor layers are allowed to diverge during training after initialization.Source: §4.3 · p. 5

5. Experiments

Experiments evaluate PolyViT across nine image, video, and audio classification tasks, as well as on linear probes for transfer to new datasets. Co-training performs well across tasks and modalities, while the selected weighted schedule supports strong benchmark results and parameter savings.

5.1. Experimental Setup

  • The experiments co-train nine tasks with one dataset per task, spanning varied domains and training-set sizes; AudioSet is evaluated with mAP and other tasks with accuracy.The task hyperparameters and single-task baselines use settings reported in prior work.Source: §5.1 · p. 5

5.2. Selecting the best task sampling schedule for co-training

  • Weighted sampling achieves the highest accuracy on 8 of 9 tasks and the second-highest on the remaining task, and its hyperparameters can be tuned in single-task settings then reused.Pets is sampled for 500 of 417,000 total steps, or 0.11% of SGD updates, yet achieves its highest accuracy with this schedule.Source: Table 1 · p. 6, §5.2 · p. 6

5.3. Co-training with PolyViT

  • Single-modality co-training leads on 7 of 9 datasets; multimodal PolyViT stays within 1.2% of single-task baselines while using 8.3 times fewer parameters.With half the transformer layers modality-specific, it improves over single-task baselines on 4 datasets, is at most 0.5% worse on the others, and uses 4.3 times fewer parameters.Source: Table 2 · p. 6, §5.3 · p. 7

5.4. Evaluating learned representations with linear probes

  • Multimodal PolyViT transfers across all 11 linear-evaluation tasks, while single-modality models generally transfer poorly across modalities; video-trained models still transfer well to images.The evaluation trains only a linear classifier on a frozen encoder.Source: §5.4 · pp. 7–8, Table 3 · p. 7

5.5. State-of-the-art performance with single-modality co-training

  • PolyViT surpasses prior state of the art on both audio datasets and all three video datasets, with gains of 2.8% on VGGSound and 1.8% on Kinetics-400, using fewer parameters.The audio comparison uses MBT under matched settings; the video comparison uses ViViT under the same experimental settings.Source: Table 4 · p. 8, §5.5 · p. 8, Table 5 · p. 8

6. Conclusion

PolyViT co-training yields state-of-the-art results on video and audio benchmarks while reducing parameters, and multi-modal co-training learns representations that generalize across modalities. The authors also identify limits in upstream-data coverage, inference, and cross-modal fusion.

6. Conclusion

  • Co-training on a single modality achieves state-of-the-art results on three video and two audio datasets while reducing parameters linearly versus single-task models.The paper also reports accuracy improvements when training for the same total number of steps.Source: §6 · p. 8
  • Multi-modal co-training is more parameter-efficient, learns representations that generalize across modalities, and supports adding tasks by training an additional output head.The authors also report that co-training is practical without tuning hyperparameters jointly across datasets.Source: §6 · p. 8
  • The method does not co-train on large-scale upstream datasets such as ImageNet-21K and C4.The authors identify these datasets as future-work targets.Source: §6 · p. 8
  • PolyViT does not improve inference time because it processes only one task at a time.Source: §6 · p. 8
  • The model does not fuse modalities such as video and audio to make a prediction.The authors identify multi-modal fusion as future work.Source: §6 · p. 8

Appendix

The appendix points readers to additional experimental details across the paper's experiments. The broader-impact discussion notes that applications and training-data biases can affect the suitability and social implications of classification models.

Appendix

  • Appendices A–E provide additional experimental settings, details, or results for Sections 5.1–5.5.Source: Appendix · p. 12

Broader Impact

  • The implications of perceptual classification applications depend on the intentions of the people building and using them.Source: Broader Impact · p. 12
  • Biases in training datasets may make models unsuitable for some applications, and classification decisions may affect social groups differently.Source: Broader Impact · p. 12

A. Experimental set-up: additional details

The experimental details specify learning-rate and warmup choices, dataset splits, and modality-specific augmentation and regularization.

A. Experimental set-up: additional details

  • Co-training sums task-specific warmup steps and uses one momentum state; video and audio learning rates are reused from prior work, while image learning rates are selected on validation data.The image learning-rate grid is {0.03, 0.1, 0.3}.Source: A. Experimental set-up: additional… · p. 12
  • Validation splits use 2% of CIFAR 10/100, 10% of Pets, 1% of ImageNet-1k, and 20% of RESISC45 training data; video and audio use standard splits.Source: A. Experimental set-up: additional… · p. 12
  • Image tasks use no augmentation, while video and audio follow prior preprocessing; audio also uses Mixup with α = 0.3 and stochastic depth with p = 0.3.Stochastic depth is applied along both audio adaptor and shared layers.Source: A. Experimental set-up: additional… · p. 12

B. Selecting the best task sampling schedule: additional experimental details

The accumulated schedule uses the smallest learning rate across tasks, while the task-by-task schedule uses a randomized task order.

B. Selecting the best task sampling schedule: additional experimental details

C. Co-training with PolyViT: additional experimental details and results

This section details PolyViT’s large-model evaluation and experimental setup, including how input sequences and task-level results are reported.

C. Co-training with PolyViT: additional experimental details and results

  • Video and audio test performance uses multiple-crop evaluation.The procedure follows the cited approaches for videos and audio, respectively.Source: C. Co-training with PolyViT:… · p. 12
  • For the Large configuration, audio baselines received validation-set hyperparameter tuning, and those changes were applied to all PolyViT runs.MiniAS used 30 epochs rather than 50 for the Base model; image tasks used separate learning-rate grid searches for baselines and linear probes.Source: C. Co-training with PolyViT:… · p. 12
  • Image-task learning rates were selected from {0.03, 0.1, 0.3} using single-task validation accuracy, then used for baselines and PolyViT variants.The sequence length is one class token plus the number of patches, and shared transformer layers reuse parameters across different sequence lengths.Source: Table 6 · p. 13, Table 7 · p. 13
  • Table 8 reports test accuracy and MiniAS mAP for the Large configuration, with parameter totals summed across models in multi-model rows.The table’s “# models” column indicates rows containing multiple models.Source: Table 8 · p. 13

D. Linear probes: additional experimental details

The linear-probe experiments specify task splits, training choices, and how embeddings are adapted when probing across modalities.

D. Linear probes: additional experimental details

  • For video and audio probes, only output-head parameters are optimized, using hyperparameters from the cited modality-specific work.Image-task training and warmup steps reuse those of RESISC45; co-training uses multiple-crop evaluation for video and audio.Source: D. Linear probes: additional… · p. 13
  • Image tasks use 2% of their training set for validation, while video and audio use their standard train, validation, and test splits.Source: D. Linear probes: additional… · p. 14
  • Cross-modal linear probes adapt patch and positional embeddings differently according to the source and target modalities.For video-to-image or audio probes, the method sums 3D patch embeddings across frames and averages positional embeddings; image- or audio-to-video probes repeat 2D patch embeddings along the frame axis.Source: D. Linear probes: additional… · p. 14
  • Image tasks use no augmentation; video and audio follow the cited preprocessing and augmentation, with Mixup α = 0.3 for S-S v2.Source: D. Linear probes: additional… · p. 14
  • Table 9 lists tasks used for linear probes, and learning-rate grid searches use validation-set performance.Source: Table 9 · p. 14

E. State-of-the-art performance on one modality: additional experimental details and results

The additional video comparison finds PolyViT ahead of a ViViT model trained on the same amount of data, while the section also specifies baseline scope and experimental settings.

E. State-of-the-art performance on one modality: additional experimental details and results

Loading 2111.12993v1…