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 presents an vision-and-language pretraining framework that utilizes momentum distillation to address the issue of training on noisy web data. More specifically, it presents a self-training approach to learn from pseudo-targets that are generated by a momentum model. In contrast to existing approaches which simply use cross-modal attention to reason about concatenated sequences of image regions and words, this work also leverages contrastive learning to learn more effective representations before using cross-modal attention.