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

Qing Li, Boqing Gong+5 more|arXiv (Cornell University)|Dec 14, 2021

Advanced Neural Network Applications9 citations

TL;DR

This paper proposes ViT-BERT, a unified foundation transformer that jointly pre-trains on unpaired images and text using modality-specific tokenizers, a shared transformer encoder, and task-specific heads. It introduces knowledge distillation from pre-trained BERT and ViT teachers and a gradient masking strategy to balance multi-modal learning, achieving performance close to separate models on both vision and NLP benchmarks.

ABSTRACT

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.

Motivation & Objective

• To explore the feasibility of training a single transformer model for both vision and natural language tasks.
• To address the challenge of joint pre-training on unpaired image and text data without modality alignment.
• To minimize performance degradation in vision and text tasks when sharing a single model architecture.
• To develop techniques that enable effective multi-modal pre-training with conflicting gradients and limited supervision.
• To validate that a unified model can achieve performance comparable to separately pre-trained BERT and ViT on downstream tasks.

Proposed method

• Design a unified transformer architecture with modality-specific tokenizers for images and text, a shared transformer encoder, and task-specific heads.
• Use separately pre-trained BERT and ViT as teacher models to provide knowledge distillation signals during joint training.
• Apply a gradient masking strategy that selectively updates parameters based on gradient magnitude, gradually sparsifying the mask to balance image and text pre-training.
• Train the model end-to-end on unpaired image-text data (e.g., Wikipedia) using a combined loss from both modalities.
• Gradually increase mask sparsity during training, favoring parameters with higher gradient magnitude for text pre-training.
• Fine-tune the unified model on downstream vision and NLP tasks to evaluate transferability.

Experimental results

Research questions

• RQ1 Can a single transformer model be effectively pre-trained on unpaired images and text to serve as a unified foundation model for both vision and language tasks?
• RQ2 How can knowledge distillation from pre-trained BERT and ViT improve joint pre-training performance?
• RQ3 Can gradient masking effectively reconcile conflicting optimization signals from image and text pre-training objectives?
• RQ4 To what extent does the unified model match the performance of separately trained BERT and ViT on downstream tasks?
• RQ5 Does the shared representation space in the transformer encoder support cross-modal semantic alignment?

Key findings

• The jointly pre-trained ViT-BERT model achieves strong performance on both image classification and natural language understanding tasks, approaching the performance of separately trained BERT and ViT.
• Knowledge distillation from pre-trained BERT and ViT significantly improves the quality of representations learned during joint pre-training.
• The gradient masking strategy effectively balances learning signals from image and text pre-training, reducing performance degradation.
• The top layers of the transformer encoder develop shared, high-level representations activated by both images and text, indicating cross-modal semantic alignment.
• The mask sparsity increases faster in deeper transformer blocks, suggesting that higher-level features are more amenable to parameter sharing between modalities.
• Qualitative analysis shows that neurons in deeper blocks co-activate for semantically similar concepts across images and text, such as 'cat' in an image and the word 'cat'.