RelViT: Concept-guided Vision Transformer for Visual Relational Reasoning

1 INTRODUCTION . Deep neural networks have achieved great success in visual recognition . However , their ability for visual relational reasoning , i.e . reasoning with entities and their relationships in a visual scene , still falls short of human-level performances , especially in real-world domains . The challenges of common visual relational reasoning tasks , e.g . HICO and GQA benchmarks ( Chao et al. , 2015 ; Hudson & Manning , 2019 ) are manifested in three aspects : 1 ) object-centric learning : to identify objects ( including humans ) as well as their visual properties ; 2 ) relational reasoning : to infer all pairwise relationships between the object entities ; and 3 ) systematic generalization : to reason with visual entities and relations on novel object-relation combinations and extrapolate to longer reasoning hops ( Bahdanau et al. , 2018 ; Hupkes et al. , 2020 ) . While existing models have leveraged pre-trained object detectors ( Ren et al. , 2015 ; Jiang et al. , 2020 ) and/or explicit symbolic reasoning methods ( Yi et al. , 2018 ) to tackle these challenges , they leave ample space for improvement . More recently , vision transformers ( ViTs ) have become the new paradigm for visual recognition and have made great strides in a broad range of visual recognition tasks ( Dosovitskiy et al. , 2020 ; Wang et al. , 2021a ; Liu et al. , 2021 ) . Several properties of ViT make it a compelling model choice for visual relational reasoning . First , the self-attention mechanism in ViT offers a strong relational inductive bias , explicitly modeling the relations between input entities . Second , the design of image as patches facilitates the learning of object-centric representations , as evidenced by recent works , e.g . DINO and EsViT ( Caron et al. , 2021 ; Li et al. , 2021 ) , that demonstrate ViTs trained with self-supervised learning ( SSL ) capture objects in the image without label annotations . To investigate the efficacy of the ViT backbone for visual relational reasoning , in particular on systematic generalization , we introduce new systematic splits to canonical benchmarks and compare the ViT backbone with the CNN backbone . Results on GQA show that switching to ViTs in MCAN model ( Yu et al. , 2019 ) brings an immediate 11 % gain in accuracy . However , the performance gap between the original GQA testing split and the new systematic split remains considerable ( 15 % in accuracy ) for both backbones . It suggests that generic ViTs can still be improved to tackle the reasoning task , especially on systematic generalization . Recent works have shown that neural networks can learn representations with better generalization , by learning certain auxiliary tasks of predicting human-specified concepts ( Hill et al. , 2020 ; Koh et al. , 2020 ) . A natural question emerges : can we exploit these concepts to improve the reasoning ability of ViTs ? Our approach is to make better use of concepts ( e.g . the labels in the original training dataset ) in the ViT training for better relational reasoning . To this end , we first introduce a novel concept-feature dictionary , where each key is a concept and its value is a queue of image features with the same concept , as shown in Figure 1 . It allows dynamic and flexible training-time image feature retrieval during training . Based on this dictionary , we then augment the canonical ViT training pipeline with two auxiliary tasks : To facilitate high-level reasoning about relationships , we design a global task that helps cluster images with the same concept together to produce semantically consistent relational representations . To learn better object-centric representations , we develop a local task that guides the model to discover object-centric semantic correspondence across images ( Liu et al. , 2010 ) . Thanks to the plug-and-play feature of our concept-feature dictionary , our auxiliary tasks can be easily incorporated into the existing ViT training pipeline without additional input pre-processing . We term the resulting model concept-guided vision transformer ( or RelViT for short ) . We evaluate our method on two standard visual relational reasoning benchmarks : HICO and GQA . Beyond the original independent and identically distributed ( I.I.D . ) training-testing split , we introduce a systematic split for each dataset to examine the ability of systematic generalization , i.e. , recognizing novel object-relation combinations . Our results show that RelViT significantly outperforms previous approaches . On HICO , it improves the best baseline by 16 % , 43 % , and 7 % on the original non-systematic and two new systematic splits , respectively , as shown in Figure 2 . On GQA , it further closes the gap of overall accuracy between models using visual backbone feature only and models using additional bounding box features ( obtained from pre-trained object detectors ) by 13 % and 18 % on the two splits . We also show that our method is compatible with various ViT variants and robust to hyperparameters . Finally , our qualitative inspection indicates that RelViT does improve ViTs on learning relational and object-centric representations . Our main contributions are summarized as follows : • We propose RelViT , by incorporating visual relational concepts to the ViT training with the newlyintroduced concept-guided global and local auxiliary tasks , where a concept-feature dictionary is proposed to enable dynamic and flexible image feature retrieval with the concept keys . • In extensive experiments on the original non-systematic and new systematic split of the HICO and GQA datasets , we demonstrate the advantages of RelViT over various strong baselines for visual relational reasoning . • We perform ablation studies on RelViT to show the contributions of its key components , its compatibility to various ViT architectures , and its robustness to hyper-parameters . We provide qualitative results to confirm our improved learning of relational and object-centric representations . 2 METHODOLOGY . 2.1 BACKGROUND . Vision transformers . Here we briefly review the architecture of multi-staged ViTs ( Dosovitskiy et al. , 2020 ) . Given an image I ∈ RH×W×C , a ViT model g first tokenizes the input into N image tokens ( patches ) with a resolution of ( T , T ) : tokenize ( I ) = [ t1 , · · · , tN ] , ti ∈ RT 2×C , N = HW/T 2 , where ( H , W ) and C denotes the original resolution and number of channel of the image , respectively . Then in each stage , a patch embedding and a multi-head self attention ( MHSA ) module is applied to these tokens to produce input for the next stage . The final output of ViT g ( I ) is a sequence of tokens [ z1 , · · · , zN ] that correspond to the aforementioned input tokens . For global prediction tasks , e.g . image categorization , a summary of the input image can be obtained by either inserting an extra [ CLS ] token to the input sequence of image tokens or performing an extra pooling operation over the output tokens ( Zhai et al. , 2021 ) . Self-supervised learning with DINO and EsViT . Our method is developed upon the recently proposed self-supervised learning ( SSL ) approach self-distillation with no labels ( DINO ) ( Caron et al. , 2021 ) and its follow-up EsViT ( Li et al. , 2021 ) . As shown in Figure 1 , their main idea is to encourage the output consistency between a teacher gt and a student network gs , parameterized by θt and θs , respectively . Given an input image I , both networks map it to a probability distribution Pt ( I ) = ht ( gt ( I ) ) and Ps ( I ) = hs ( gs ( I ) ) via an extra projection head h ( · ) . The teacher and student network will be updated alternatively by following these two rules : ( 1 ) For the student network : θs ← argminθs LGlobal , where LGlobal = −Pt ( I ) logPs ( I ) ; ( 2 ) For the teacher network , θt is updated using an exponential moving average ( EMA ) on θs : θt ← λθt + ( 1− λ ) θs , where λ controls the updating momentum . In practice , multiple views of the input image I will be generated via data augmentation and the teacher and student networks will receive different views , preventing the task from being trivialized . EsViT further extends the image-level loss LGlobal to patch-level by applying dense SSL ( Wang et al. , 2021c ) for learning correspondence between the different views , enhancing the performance on dense prediction . Readers are encouraged to refer to Caron et al . ( 2021 ) and Li et al . ( 2021 ) for more details about these two works . 2.2 RELVIT . RelViT is a concept-guided ViT that makes better use of the concepts in the ViT training for better relational reasoning . In this section , we first introduce a concept-feature dictionary to store and retrieve image features with their concept keys . We then augment the canonical ViT training pipeline with two auxiliary tasks : a global level task and a local level task , both are concept-guided by resorting to the concept-feature dictionary . Intuitively , the global task help cluster images with the same concept together to produce semantically consistent relational features , while the local task guides the model to discover object-centric semantic correspondence across images . Concept-feature dictionary . We assume the total number of concepts is M , and the set of all concepts C = { c1 , · · · , cM } . A concept-feature dictionary is denoted by D = { ( c1 , Q1 ) , · · · , ( cM , QM ) } , where each concept ci is associated with a queue Qi of image features . During training , each image I may come with multiple concepts , which we denote by CI ⊂ C. For instance , there may exist several human-object interactions in an image from the HICO dataset , each of which may correspond to a concept . As shown in Figure 1 , whenever a new image-concept pair ( I , CI ) comes , we uniformly draw a concept code c from CI , pick up the queue Q from the dictionary that corresponds to c , and then retrieve the image feature f from Q . Meanwhile , we pass the input image I to the teacher network gt to get the new image feature f ′ = gt ( I ) , and enqueue it to Q . Note that if Q is full already , we first need to dequeue the oldest image feature from Q . During training , we use the retrieved image feature f for the two auxiliary tasks below , rather than the input image feature f ′ . Furthermore , the sampling strategy , i.e . how to retrieve image feature f from Q , plays an important role in the overall performance of our method . We consider the following two sampling strategies : Algorithm 1 RelViT : Concept-guided Vision Transformer Input : A set of training images with concepts { ( I1 , C1 ) , · · · } , an image augmentation function aug ( · ) , momentum update factor λ , loss weight α , a concept-feature dictionary D , teacher and student ViT gt and gs , parameterized by θt and θs , respectively . 1 : for ( Ii , Ci ) in { ( I1 , C1 ) , · · · } do 2 : I ( 1 ) i , I ( 2 ) i = aug ( Ii ) , aug ( Ii ) 3 : Uniformly draw a concept code c ∼ Ci . 4 : Retrieve Q from D with c. 5 : if Q is not empty then 6 : Sample feature f ∼ Q , following some sampling tactics . 7 : Laux = LGlobal ( f , gs ( I ( 2 ) i ) ) + LLocal ( f , gs ( I ( 2 ) i ) ) 8 : Insert feature gt ( I ( 1 ) i ) into Q ; if it is full , remove the oldest feature . 9 : else 10 : Laux = LGlobal ( gt ( I ( 1 ) i ) , gs ( I ( 2 ) i ) ) + LLocal ( gt ( I ( 1 ) i ) , gs ( I ( 2 ) i ) ) 11 : end if 12 : Update θs with the loss function L = Lmain + αLaux . 13 : Update θt using an EMA : θt ← λθt + ( 1− λ ) θs . 14 : end for • Uniform sampling . Each image feature is drawn with equal probability from the queue , i.e . suppose we have N features in the queue , then the probability of each feature being sampled is 1/N . This tactic encourages the diversity of the retrieved image features , benefiting the overall performance . However , some older features in the queue may largely fall behind the current model if the teacher network gt is updated quickly , eliciting unstable training . • “ Most-recent ” sampling . The sampling probability mass is allocated based on the freshness of image features , and the most recent feature has the highest chance to be retrieved . Specifically , suppose we have N features in the queue Q ( |Q| > = N ) . Then for the i-th newest feature f , we define its weight wi = N − i + 1 . Finally , the probability of the i-th newest feature being sampled is wi/ ∑N j=1 wj . This tactic ensures we retrieve more up-to-date features and thereby stabilizes the learning . But it may hurt the overall performance due to a lack of feature diversity , as the chance of older features being sampled is small . Note that the feature queue is empty at the beginning of training . In this case , we simply use the input image feature f ′ for the auxiliary tasks , and also enqueue it to Q that corresponds to the concept of the input image , as shown in Algorithm 1 . As we can show in the next , now our proposed global and local tasks reduce to DINO ( Caron et al. , 2021 ) and EsViT ( Li et al. , 2021 ) , respectively . Concept-guided global task . Suppose we have two views { I ( 1 ) , I ( 2 ) } of an image I , the main idea of our concept-guided global task is to replace I ( 1 ) in the DINO loss ( Caron et al. , 2021 ) with the image feature f sampled from the concept-feature dictionary , which becomes LGlobal = −ht ( f ) log hs ( gs ( I ( 2 ) ) ) , ( 1 ) where ht and hs are the projection head of the teacher and student network , respectively , and gs is the student network . Intuitively , minimizing the global loss is equivalent to encouraging the similarity of any two different image features with the same concept . Hence , it can help produce more semantically consistent relational representations , in particular when the concepts stored in the concept-feature dictionary are themselves relational . Similar inter-class representation learning techniques have been explored before ( Wang et al. , 2017 ; Caron et al. , 2018 ) . However , these approaches require a rather complex pre-processing stage , e.g . the images have to be split in terms of the concept before training , making them not directly applicable to existing training pipelines . Rather , with our proposed concept-feature dictionary that dynamically saves & retrieves image features from the running storage , our concept-guided global task becomes a plug-n-play task to existing training pipelines . Besides , our dictionary also allows replaying features relationally , which can help tackle the long-tailed data distribution issue ( Shen et al. , 2018 ) in many real-world visual reasoning tasks . Concept-guided local task . As we mentioned earlier , our concept-guided local task aims at facilitating object-centric learning , by the means of correspondence learning ( Liu et al. , 2010 ; Wang et al. , 2019 ) . Recent studies have unveiled the possibility of learning correspondence with SSL ( Wang et al. , 2021c ; Li et al. , 2021 ) . However , only low-level correspondence between two augmented ( e.g . rotated ) views of an image can be discovered , while the semantic information of objects is missing . To remedy this , we bring concepts to these methods , endowing them the capability of learning semantic correspondence from images . Specifically , suppose we have two views { I ( 1 ) , I ( 2 ) } of an image I , and we also tokenize the image feature into a sequence of N local image tokens . Then at the output of ViT , we obtain gt ( I ( 1 ) ) = [ z ( 1 ) 1 , · · · , z ( 1 ) N ] and gs ( I ( 2 ) ) = [ z ( 2 ) 1 , · · · , z ( 2 ) N ] , where z denotes the local feature . Prior work , such as EsViT ( Li et al. , 2021 ) , relies on the local features gt ( I ( 1 ) ) and gt ( I ( 2 ) ) for the local task . Instead , we replace gt ( I ( 1 ) ) with the image feature f retrieved from the concept-feature dictionary with the concept of I . We then split f into multiple local features , i.e . f = [ z ( f ) 1 , · · · , z ( f ) N ] and our conceptguided local loss becomes LLocal = − 1 N N∑ i=1 ht ( z ( f ) j ? ) log hs ( z ( 2 ) i ) , j ? = argmax j CosineDistance ( z ( f ) j , z ( 2 ) i ) , ( 2 ) where ht ( · ) , hs ( · ) are the projection heads that map local features to probability distributions1 . Intuitively , it greedily matches the output between two local regions that have minimal feature distance – bootstrapping the object-level semantic correspondence among images with the same concept . Overall loss . By combining the global and local tasks , we add an auxiliary task loss Laux to the main loss Lmain ( e.g . cross-entropy loss of the reasoning task ) . The eventual objective is L = Lmain + αLaux , Laux = LGlobal + LLocal , ( 3 ) where a trade-off weight α is added for better flexibility . As we mentioned above , our method will reduce to EsViT , a baseline with non-concept auxiliary tasks , when we use the current input features gt ( I ( 1 ) ) instead of f sampled from our dictionary for computing LGlobal and LLocal .

Summary: This work studies visual relationship reasoning and proposes a model based on the recent contrastive learning methods DINO and EsViT. Compared with the two predecessors, a feature dictionary is proposed to provide the online updated concept features for the teacher model instead. Correspondingly, two auxiliary tasks are introduced to drive semantic concept learning. Experiments on two large benchmarks show the efficacy of the proposed method. And extensive ablations are also made to probe the model components.