Align before Fuse: Vision and Language Representation Learning with Momentum Distillation

1 Introduction . Vision-and-Language Pre-training ( VLP ) aims to learn multimodal representations from large-scale image-text pairs that can improve downstream Vision-and-Language ( V+L ) tasks . Most existing VLP methods ( e.g . LXMERT [ 1 ] , UNITER [ 2 ] , OSCAR [ 3 ] ) rely on pre-trained object detectors to extract region-based image features , and employ a multimodal encoder to fuse the image features with word tokens . The multimodal encoder is trained to solve tasks that require joint understanding of image and text , such as masked language modeling ( MLM ) and image-text matching ( ITM ) . While effective , this VLP framework suffers from several key limitations : ( 1 ) The image features and the word token embeddings reside in their own spaces , which makes it challenging for the multimodal encoder to learn to model their interactions ; ( 2 ) The object detector is both annotation-expensive and compute-expensive , because it requires bounding box annotations during pre-training , and highresolution ( e.g . 600⇥1000 ) images during inference ; ( 3 ) The widely used image-text datasets [ 4 , 5 ] are collected from the web and are inherently noisy , and existing pre-training objectives such as MLM may overfit to the noisy text and degrade the model ’ s generalization performance . We propose ALign BEfore Fuse ( ALBEF ) , a new VLP framework to address these limitations . We first encode the image and text independently with a detector-free image encoder and a text encoder . Then we use a multimodal encoder to fuse the image features with the text features through crossmodal attention . We introduce an intermediate image-text contrastive ( ITC ) loss on representations from the unimodal encoders , which serves three purposes : ( 1 ) it aligns the image features and the text features , making it easier for the multimodal encoder to perform cross-modal learning ; ( 2 ) it 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . improves the unimodal encoders to better understand the semantic meaning of images and texts ; ( 3 ) it learns a common low-dimensional space to embed images and texts , which enables the image-text matching objective to find more informative samples through our contrastive hard negative mining . To improve learning under noisy supervision , we propose Momentum Distillation ( MoD ) , a simple method which enables the model to leverage a larger uncurated web dataset . During training , we keep a momentum version of the model by taking the moving-average of its parameters , and use the momentum model to generate pseudo-targets as additional supervision . With MoD , the model is not penalized for producing other reasonable outputs that are different from the web annotation . We show that MoD not only improves pre-training , but also downstream tasks with clean annotations . We provide theoretical justifications on ALBEF from the perspective of mutual information maximization . Specifically , we show that ITC and MLM maximize a lower bound on the mutual information between different views of an image-text pair , where the views are generated by taking partial information from each pair . From this perspective , our momentum distillation can be interpreted as generating new views with semantically similar samples . Therefore , ALBEF learns vision-language representations that are invariant to semantic-preserving transformations . We demonstrate the effectiveness of ALBEF on various downstream V+L tasks including image-text retrieval , visual question answering , visual reasoning , visual entailment , and weakly-supervised visual grounding . ALBEF achieves substantial improvements over existing state-of-the-art methods . On image-text retrieval , it outperforms methods that are pre-trained on orders of magnitude larger datasets ( CLIP [ 6 ] and ALIGN [ 7 ] ) . On VQA and NLVR2 , it achieves absolute improvements of 2.37 % and 3.84 % compared to the state-of-the-art method VILLA [ 8 ] , while enjoying much faster inference speed . We also provide quantitative and qualitative analysis on ALBEF using Grad-CAM [ 9 ] , which reveals its ability to perform accurate object , attribute and relationship grounding implicitly . 2 Related Work . 2.1 Vision-Language Representation Learning . Most existing work on vision-language representation learning fall into two categories . The first category focuses on modelling the interactions between image and text features with transformerbased multimodal encoders [ 10 , 11 , 12 , 13 , 1 , 14 , 15 , 2 , 3 , 16 , 8 , 17 , 18 ] . Methods in this category achieve superior performance on downstream V+L tasks that require complex reasoning over image and text ( e.g . NLVR2 [ 19 ] , VQA [ 20 ] ) , but most of them require high-resolution input images and pre-trained object detectors . A recent method [ 21 ] improves inference speed by removing the object detector , but results in lower performance . The second category focuses on learning separate unimodal encoders for image and text [ 22 , 23 , 6 , 7 ] . The recent CLIP [ 6 ] and ALIGN [ 7 ] perform pre-training on massive noisy web data using a contrastive loss , one of the most effective loss for representation learning [ 24 , 25 , 26 , 27 ] . They achieve remarkable performance on image-text retrieval tasks , but lack the ability to model more complex interactions between image and text for other V+L tasks [ 21 ] . ALBEF unifies the two categories , leading to strong unimodal and multimodal representations with superior performance on both retrieval and reasoning tasks . Furthermore , ALBEF does not require object detectors , a major computation bottleneck for many existing methods [ 1 , 2 , 3 , 8 , 17 ] . 2.2 Knowledge Distillation . Knowledge distillation [ 28 ] aims to improve a student model ’ s performance by distilling knowledge from a teacher model , usually through matching the student ’ s prediction with the teacher ’ s . While most methods focus on distilling knowledge from a pre-trained teacher model [ 28 , 29 , 30 , 31 , 32 ] , online distillation [ 33 , 34 ] simultaneously trains multiple models and use their ensemble as the teacher . Our momentum distillation can be interpreted as a form of online self-distillation , where a temporal ensemble of the student model is used as the teacher . Similar ideas have been explored in semi-supervised learning [ 35 ] , label noise learning [ 36 ] , and very recently in contrastive learning [ 37 ] . Different from existing studies , we theoretically and experimentally show that momentum distillation is a generic learning algorithm that can improve the model ’ s performance on many V+L tasks . 3 ALBEF Pre-training . In this section , we first introduce the model architecture ( Section 3.1 ) . Then we delineate the pretraining objectives ( Section 3.2 ) , followed by the proposed momentum distillation ( Section 3.3 ) . Lastly we describe the pre-training datasets ( Section 3.4 ) and implementation details ( Section 3.5 ) . 3.1 Model Architecture . As illustrated in Figure 1 , ALBEF contains an image encoder , a text encoder , and a multimodal encoder . We use a 12-layer visual transformer ViT-B/16 [ 38 ] as the image encoder , and initialize it with weights pre-trained on ImageNet-1k from [ 31 ] . An input image I is encoded into a sequence of embeddings : { vcls , v1 , ... , vN } , where vcls is the embedding of the [ CLS ] token . We use a 6-layer transformer [ 39 ] for both the text encoder and the multimodal encoder . The text encoder is initialized using the first 6 layers of the BERTbase [ 40 ] model , and the multimodal encoder is initialized using the last 6 layers of the BERTbase . The text encoder transforms an input text T into a sequence of embeddings { wcls , w1 , ... , wN } , which is fed to the multimodal encoder . The image features are fused with the text features through cross attention at each layer of the multimodal encoder . 3.2 Pre-training Objectives . We pre-train ALBEF with three objectives : image-text contrastive learning ( ITC ) on the unimodal encoders , masked language modeling ( MLM ) and image-text matching ( ITM ) on the multimodal encoder . We improve ITM with online contrastive hard negative mining . Image-Text Contrastive Learning aims to learn better unimodal representations before fusion . It learns a similarity function s = gv ( vcls ) > gw ( wcls ) , such that parallel image-text pairs have higher similarity scores . gv and gw are linear transformations that map the [ CLS ] embeddings to normalized lower-dimensional ( 256-d ) representations . Inspired by MoCo [ 24 ] , we maintain two queues to store the most recent M image-text representations from the momentum unimodal encoders . The normalized features from the momentum encoders are denoted as g0v ( v0cls ) and g 0 w ( w 0 cls ) . We define s ( I , T ) = gv ( vcls ) > g0w ( w 0 cls ) and s ( T , I ) = gw ( wcls ) > g0v ( v 0 cls ) . For each image and text , we calculate the softmax-normalized image-to-text and text-to-image similarity as : pi2tm ( I ) = exp ( s ( I , Tm ) /⌧ ) PM m=1 exp ( s ( I , Tm ) /⌧ ) , pt2im ( T ) = exp ( s ( T , Im ) /⌧ ) PM m=1 exp ( s ( T , Im ) /⌧ ) ( 1 ) where ⌧ is a learnable temperature parameter . Let yi2t ( I ) and yt2i ( T ) denote the ground-truth one-hot similarity , where negative pairs have a probability of 0 and the positive pair has a probability of 1 . The image-text contrastive loss is defined as the cross-entropy H between p and y : Litc = 1 2 E ( I , T ) ⇠D ⇥ H ( yi2t ( I ) , pi2t ( I ) ) + H ( yt2i ( T ) , pt2i ( T ) ) ⇤ ( 2 ) Masked Language Modeling utilizes both the image and the contextual text to predict the masked words . We randomly mask out the input tokens with a probability of 15 % and replace them with the special token [ MASK ] 1 . Let T̂ denote a masked text , and pmsk ( I , T̂ ) denote the model ’ s predicted probability for a masked token . MLM minimizes a cross-entropy loss : Lmlm = E ( I , T̂ ) ⇠DH ( y msk , pmsk ( I , T̂ ) ) ( 3 ) where ymsk is a one-hot vocabulary distribution where the ground-truth token has a probability of 1 . Image-Text Matching predicts whether a pair of image and text is positive ( matched ) or negative ( not matched ) . We use the multimodal encoder ’ s output embedding of the [ CLS ] token as the joint representation of the image-text pair , and append a fully-connected ( FC ) layer followed by softmax to predict a two-class probability pitm . The ITM loss is : Litm = E ( I , T ) ⇠DH ( yitm , pitm ( I , T ) ) ( 4 ) where yitm is a 2-dimensional one-hot vector representing the ground-truth label . We propose a strategy to sample hard negatives for the ITM task with zero computational overhead . A negative image-text pair is hard if they share similar semantics but differ in fine-grained details . We use the contrastive similarity from Equation 1 to find in-batch hard negatives . For each image in a mini-batch , we sample one negative text from the same batch following the contrastive similarity distribution , where texts that are more similar to the image have a higher chance to be sampled . Likewise , we also sample one hard negative image for each text . The full pre-training objective of ALBEF is : L = Litc + Lmlm + Litm ( 5 )

Summary: This paper proposes a new vision-language representation learning framework, ALBEF, which uses image-text contrastive (ITC), masked language modeling (MLM), and image-text matching (ITM) losses with momentum distillation to learn SOTA representations. The ITC and MLM losses are adapted to be convex combinations of the original loss and the KL divergence between the predicted probability distribution (of the image or text or masked tokens) and the soft pseudo targets obtained from the momentum model.   The pipeline is evaluated on image-text retrieval, visual entailment, NLVR, VQA, and grounding and shows improved performance on all tasks compared to prior work, and with fewer training samples in comparison to several prior methods. Qualitative visualizations also verify the quality of the learned representations, with convincing heat maps suggesting well grounded, semantically aligned embeddings.