PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards a Unified Foundation Model: Jointly Pre-Training Transformers on Unpaired Images and Text

Qing Li, Boqing Gong, Yin Cui, Dan Kondratyuk, Xianzhi Du, Ming-Hsuan Yang, Matthew Brown

arXiv:2112.07074v1cs.CVcs.LG

TL;DR

The paper asks whether one pre-trained transformer can serve both vision-only and text-only tasks, instead of maintaining separate models. ViT-BERT jointly pre-trains on unpaired images and text using teacher distillation and gradient masking. It learns high-quality representations for both task families, while the work’s scope remains limited to image and text and single-modal tasks.

  • Problem

    The paper asks whether a single transformer can be pre-trained for different modalities, addressing the need to maintain separate foundation models for different domains.

    Source: §1 · p. 1

  • Method

    ViT-BERT combines modality-specific tokenizers and task-specific heads with a shared encoder, and jointly pre-trains on unpaired images and text using BERT and ViT teachers and gradient masking.

    Source: §1 · pp. 1–2

  • Results

    Jointly pre-trained ViT-BERT learns high-quality representations for vision-only and text-only tasks, with distillation and gradient masking helping performance approach separately trained models.

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

  • Takeaways & Limitations

    The experiments demonstrate that one jointly pre-trained ViT-BERT can transfer to both vision-only and text-only tasks.

    Source: §1 · p. 2

    The work focuses on images and text and on single-modal tasks, leaving other modalities and multimodal pre-training or downstream reasoning for future work.

    Source: §5 · p. 8

Abstract

from arXiv · show

In this paper, we explore the possibility of building a unified foundation model that can be adapted to both vision-only and text-only tasks. Starting from BERT and ViT, we design a unified transformer consisting of modality-specific tokenizers, a shared transformer encoder, and task-specific output heads. To efficiently pre-train the proposed model jointly on unpaired images and text, we propose two novel techniques: (i) We employ the separately-trained BERT and ViT models as teachers and apply knowledge distillation to provide additional, accurate supervision signals for the joint training; (ii) We propose a novel gradient masking strategy to balance the parameter updates from the image and text pre-training losses. We evaluate the jointly pre-trained transformer by fine-tuning it on image classification tasks and natural language understanding tasks, respectively. The experiments show that the resultant unified foundation transformer works surprisingly well on both the vision-only and text-only tasks, and the proposed knowledge distillation and gradient masking strategy can effectively lift the performance to approach the level of separately-trained models.

1. Introduction

The paper asks whether one transformer can serve tasks across modalities, rather than requiring separate foundation models. It proposes ViT-BERT for unpaired image and text pre-training, using distillation and gradient masking to improve joint training; evaluations report strong transfer to vision-only and text-only tasks.

1. Introduction

  • The paper frames unification as a response to maintaining separate foundation models for different domains and asks whether one transformer can serve multiple modalities.The stated goal is a model that can be fine-tuned for tasks in different modalities.Source: §1 · p. 1
  • ViT-BERT combines modality-specific tokenizers, a shared transformer encoder, and task-specific heads, with most parameters and computation placed in the shared encoder.The model is jointly pre-trained on unpaired images and text, then fine-tuned for downstream tasks.Source: §1 · p. 1
  • Knowledge distillation uses separately trained BERT and ViT teachers for supervision, while gradient masking balances text and image pre-training updates.The mask selects important text-gradient parameters and leaves remaining parameters for image-pre-training updates; its sparsity increases gradually during training.Source: §1 · p. 2
  • Jointly pre-trained ViT-BERT works well on vision-only and text-only tasks, with distillation and gradient masking bringing performance toward separately trained BERT and ViT.The evaluation fine-tunes on image classification and GLUE natural language understanding tasks.Source: §1 · p. 2

2. Related Work

The related work covers foundation models in language and vision, multimodal models, and multi-task learning. The paper distinguishes its approach by jointly pre-training a lightly modality-specialized transformer for cross-modal downstream tasks.

2. Related Work

  • The paper’s related-work discussion positions its contribution as a unified foundation model for a wide range of tasks across modalities.The section covers transformer foundation models and multi-task learning.Source: §2 · p. 2, §2.2 · p. 2

2.1. Transformer-based Foundation Models

  • Text-pretrained transformers transfer to language tasks, while vision transformers have been applied to classification, detection, and segmentation.The cited vision examples include Image Transformer, iGPT, ViT [12], DETR [4], and Swin Transformer [27].Source: §2.1 · p. 2
  • Earlier multimodal transformers pretrained on image-text pairs address visual question answering, captioning, and retrieval, often using an object-detection backbone for image regions.Examples include VisualBERT [24], VL-BERT, ViLBERT [30], LXMERT, and UNITER [7].Source: §2.1 · p. 2
  • Foundation models are described as transformers trained on generic data and fine-tuned on smaller downstream datasets.The passage gives Wikipedia, JFT-300M, and Conceptual Captions as examples of generic pre-training datasets.Source: §2.1 · p. 2

2.2. Multi-task Learning

  • Multi-task learning shares parameters and computation across tasks, with prior examples spanning single-modality and multimodal settings.Examples include Mask-RCNN for detection, segmentation, and pose estimation, and MT-DNN [25] for language understanding.Source: §2.2 · p. 2
  • Prior multimodal models include heterogeneous architectures, UniT with modality-specific encoders and a shared decoder, and VATT trained on aligned video, audio, and text triplets.The paper notes that VATT [1] requires aligned data triplets, while UniT [18] uses separate visual and text encoders.Source: §2.2 · p. 2
  • ViT-BERT targets a foundation model adaptable to tasks across modalities, with multi-task learning at pre-training rather than task-specific training.Its modality-specific image patch projections and text embeddings are lighter than the separate encoders used in UniT [18].Source: §2.2 · p. 3

3. Method

ViT-BERT uses modality-specific input processing and task-specific output heads around a shared transformer encoder. Joint pre-training on images and text uses teacher distillation and gradient masking to address noisy, potentially conflicting learning signals.

3.1.1 Modality-specific Tokenizer

  • Images are split into patches, flattened, projected, and combined with a class token and positional embeddings; text is tokenized as in BERT with word, position, and segment embeddings.The model does not use paired image-text inputs for pre-training.Source: §3.1.1 · p. 3

3.1.2 Shared Transformer Encoder

  • The shared encoder stacks self-attention and MLP blocks, and its final token embeddings feed the task-specific heads.Layer normalization precedes each attention or MLP layer.Source: §3.1.2 · p. 3

3.1.3 Task-specific Heads

  • Each task uses a two-layer MLP classifier over the final encoder output to predict its output classes.The class-token embedding is used for most tasks, while masked language modeling uses the masked-token embedding.Source: §3.1.3 · p. 3

3.2. Joint Pre-Training

  • Joint pre-training faces noisy optimization over many steps and potentially conflicting gradients, so the method adds distillation supervision and gradient masking.These challenges are framed as especially severe for text pre-training.Source: §3.2 · p. 3

3.2.1 Knowledge Distillation

  • Separately pre-trained BERT and ViT serve as teacher models, and distillation combines ground-truth cross-entropy with teacher-student KL divergence.The ratio α balances the two losses, and temperature τ is a distillation hyperparameter.Source: §3.2.1 · pp. 3–4

3.2.2 Gradient Masking

  • Gradient masking combines text and image gradients by retaining important text-gradient parameters and leaving the rest for image-pretraining updates.The mask is based on text-gradient magnitude and is gradually made sparser; the procedure is inspired by iterative magnitude pruning.Source: §3.2.2 · pp. 4–5

4. Experiments and Results

The experiments evaluate joint pre-training on image and text datasets, then fine-tune ViT-BERT on vision-only and text-only tasks. Results, ablations, and qualitative analyses examine performance and how representations and gradient masks develop.

4.1.1 Image Modality

  • Image pre-training uses 1.3M ImageNet images; fine-tuning covers ImageNet-1K, CIFAR-10/100, Oxford-IIIT Pets, and Oxford Flowers-102.Image fine-tuning uses 384×384 resolution, compared with 224×224 during pre-training.Source: §4.1.1 · p. 5

4.1.2 Text Modality

  • Text pre-training uses BooksCorpus and English Wikipedia with masked language modeling and next sentence prediction; fine-tuning covers five GLUE tasks.The five tasks are MNLI, QQP, QNLI, SST-2, and RTE.Source: §4.1.2 · p. 5

4.2. Model Variants

  • ViT-BERT uses a 12-layer shared transformer; comparisons include BERT, ViT, Joint, randomly initialized, and variants without knowledge distillation or gradient masking.For cross-modal fine-tuning of BERT or ViT, modality-specific embeddings are randomly initialized.Source: §4.2 · p. 5

4.3. Pre-training Setting

  • Models are jointly pre-trained for 1M steps, mixing batches of 4,096 images and 512 text documents; default settings use α=1, τ=1, and β=50%.The default iterative mask update interval is 10K steps, with a 10% mask-ratio change per iteration.Source: §4.3 · pp. 5–6

4.4.1 Fine-tuning on Downstream Tasks

  • 86% average accuracy across all tasks for ViT-BERT versus 83.3% for Joint; vision averages are 89.0% versus ViT’s 89.3%, and text averages are 83.1% versus BERT’s 88.1%.The table reports accuracy on downstream text-only and vision-only tasks.Source: §4.4.1 · p. 6, Table 1 · p. 6

4.4.2 Ablation Study

  • Knowledge distillation and gradient masking together yield 86% average accuracy across tasks, 2.7 points above Joint; KD alone improves SST-2 from 80.5% to 91.1%.Text accuracy favors α=0.7, while vision favors α=1; τ=1 works best for both. Gradient-based masks outperform random masks, and β around 0.5 balances task accuracies.Source: §4.4.1 · p. 6, §4.4.2 · p. 7, Figure 2 · p. 7

4.5. Qualitative Study

  • During joint training, upper transformer layers develop co-activation for image and text, while lower layers retain separate modality responses.Across mask updates, top blocks become sparser than bottom blocks; block 11 is sparser than block 0 by iteration 5.Source: §4.5 · pp. 7–8, Figure 3 · p. 8, Figure 4 · p. 8

5. Conclusions and Discussion

ViT-BERT is a single model jointly pre-trained on unpaired images and text, with a shared transformer encoder intended to support both domains. The authors report high-quality representations for vision-only and text-only tasks, while limiting this work to images, text, and single-modal tasks.

5. Conclusions and Discussion

  • ViT-BERT jointly pre-trains one transformer on unpaired images and text, using minimal modality- and task-specific parameters to encourage a shared encoder.The paper also reports that qualitative studies show highly shared representations across images and text in the encoder's top blocks.Source: §5 · p. 8
  • The authors report that ViT-BERT learns high-quality representations for both vision-only and text-only tasks, and that ablations validate knowledge distillation and gradient masking.The paper also reports highly shared image-text representations in the transformer's top blocks.Source: §5 · p. 8
  • The study covers images and text, but not other modalities such as video or audio, and evaluates only vision-only or text-only tasks.The authors identify multimodal pre-training and downstream tasks such as visual question answering and image captioning as possible future additions.Source: §5 · p. 8
  • Because predictions rely on training-set statistics, the model's predictions will reflect biases in those data, potentially including gender or racial biases.Source: §5 · p. 8

A.1. Pre-training

The joint-training baseline combines image and text pre-training losses with a weighted sum. Its loss ratio is selected using masked language modelling and image-classification validation accuracy.

A.1. Pre-training

  • The joint-training baseline trains the model with a weighted sum of the text and image pre-training losses.Source: §A.1 · p. 11
  • λ=0.2 is selected for the joint-training baseline using validation-set masked language modelling and image-classification accuracy.Source: §A.1 · p. 11
  • Table 3's “Avg” is the average of masked language modelling and image-classification accuracy; “NSP” means Next Sentence Prediction.Source: Table 3 · p. 11

A.2. Fine-tuning

For downstream fine-tuning, the authors replace the pre-training output heads with a randomly initialized two-layer MLP sized for the target task. Vision and text tasks use separate optimizers and learning-rate sweeps.

A.2. Fine-tuning

  • Downstream fine-tuning replaces the original output heads with a randomly initialized two-layer MLP that outputs the target task's number of classes.Source: §A.2 · p. 11
  • Vision-only fine-tuning uses SGD with momentum 0.9 and searches learning rates from 0.001 to 0.06.Source: §A.2 · p. 11
  • Text-only fine-tuning uses AdamW and searches learning rates from 5e-5 to 2e-5.Source: §A.2 · p. 11

B. Qualitative Examples

Figure 6 presents qualitative examples of ViT-BERT and baseline predictions across vision-only and text-only tasks.

B. Qualitative Examples

ImageNet-1K

The passages present text-task examples: a QNLI question, a premise–hypothesis pair, two related questions, and a movie-review sentence.

ViT-BERT

  • A QNLI example asks who translated a French version of the Scriptures prepared around 1294.The paragraph identifies Guyard de Moulin as a Roman Catholic priest but does not state the answer to the question.Source: ViT-BERT · p. 12
  • A premise–hypothesis pair compares a claim about a setting coming alive with fewer visitors against a statement about visiting at sundown or out of season.Source: ViT-BERT · p. 12
  • Two related questions ask for a favorite Prison Break character and the best-looking actor in the show.Source: ViT-BERT · p. 12
  • A movie-review example says the film lacks wit, feeling, and believability to compensate for its coarseness and banality.Source: ViT-BERT · p. 12

ViT-BERT

Figure 6 presents predictions from the model and baselines across vision-only and text-only tasks, with GT denoting ground truth.

ViT-BERT

  • Figure 6 compares model and baseline predictions across vision-only and text-only tasks; GT denotes ground truth.Source: Figure 6 · p. 12

Related papers

Loading 2112.07074v1…