Chop Chop BERT: Visual Question Answering by Chopping VisualBERT’s Heads

1 Introduction

Transformer [Vaswani et al., 2017] architecture was designed for translation tasks but has shown good performance across many other tasks, ranging from text summarization [Egonmwan and Chali, 2019; Xu et al., 2020], language modeling [Dai et al., 2019; Ma et al., 2019] to question answering [Shao et al., 2019; Sanh et al., ]. The BERT (Bidirectional Encoder Representations from Transformers) [Devlin et al., 2018] model is then proposed for general pretraining. By fine-tuning the pre-trained model on downstream tasks, BERT further advances the state-of-the-art for multiple NLP tasks. Inspired by it, many variants have been proposed towards the solution of vision-and-language tasks, e.g., Visual Question Answering (VQA [Kafle and Kanan, 2017b]), including VisualBERT [Li et al., 2020], VL-BERT [Su et al., 2020], ViLBERT [Lu et al., 2019], LXMERT [Tan and Bansal, 2019], and UNITER [Chen et al., 2020], etc.

∗ Peng Wang is corresponding author. This work is supported by the National Key RD Program of China (No. 2020AAA0106900), and the National Natural Science Foundation of China (No.61876152, No.U19B2037).

These vision-language BERT models are all Transformerbased, where the most innovative part is the multi-layer and multi-head settings, compared to the classic attention. However, how the heads and layers contribute to those reasoningrequired tasks such as Visual Question Answering, remains a mystery. Here we want to explore the role of heads and layers and focus on the VQA task, hoping to find their patterns that human can 'see' and 'understand'. When faced with highlevel tasks, there are different areas activated in our brain to fulfill different tasks. Our findings reveal that different heads and layers are indeed activated by different question types.

For the first time, we come up with a statistical analysis of Transformer's heads and layers triggered off by 12 types of questions on the Task Driven Image Understanding Challenge (TDIUC) [Kafle and Kanan, 2017a]. The TDIUC dataset is a large VQA dataset with 12 more fine-grained categories proposed to compensate for the bias in distribution of different question types of VQA 2.0 [Goyal et al., 2017], which provide convenience for our analysis. Our experiments based on the VisualBERT, as for it's general Transformer style architecture without more extra designs. The experiments can be divided into two parts: Manual Chopping and Dynamic Chopping. For manual chopping in Section 4.2, we explored the effect of a particular head by removing it, or removing all the other heads in a layer but that one. As the most obvious pattern is found at a layer level, we further experiment by chopping one layer at a time. As there are 12 kinds of questions in TDIUC and 12 layers in VisualBERT, we get a result matrix with the size 12 × 12. The echelon shape in this matrix reveal different layers play different roles across 12 types of question, with higher-level reasoning questions demanding higher-level Transformer layers gradually. For Dynamic Chopping, we chop the layers simultaneously ( Section 4.3) under the control of instance-conditioned layer scores, which are generated by a dynamic chopping module. We further find that removing some layers can decrease parameters but increase the accuracy of models.

Overall, the major contribution of this work is to delve into the roles of Transformer heads and layers in different types of visual demanding questions for the first time. We hope our interesting finding, especially the presented echelon shape of experiment results matrices, which shows a clear tendency of Transformer to gradually rely on higher-level layers when reaching higher-level reasoning questions, can inspire further investigation of the inner mechanism of self-attention layers.

2 RelatedWork & Background

2.1 Transformer in VQA

Transformer-based models, such as VisualBERT [Li et al., 2020] and ViLBERT [Lu et al., 2019], have gained popularity most recently, as they are well suited to pre-training and can be easily transferred to other similar tasks. VisualBERT consists of a stack of Transformer layers that implicitly align elements of an input text and regions in an associated image with self-attention. ViLBERT proposes to learn joint representation of images and text using a BERT-like architecture with separate Transformers for vision and language that can attend to each other. LXMERT [Tan and Bansal, 2019] is a large-scale Transformer model that consists of three encoders: an object relationship encoder, a language encoder, and a cross-modality encoder. UNITER [Chen et al., 2020] adopts Transformer to leverage its self-attention mechanism designed for learning contextualized representations.

2.2 Multi-head &Multi-layer Transformer

Transformer [Vaswani et al., 2017] architecture relies entirely on the self-attention mechanism. The self-attention Attn is calculated with a set of queries {qi}Ni=1 , a set of keys {ki}Ni=1 and a set of values {vi}Ni=1 :

where N is the number of key-value pairs, d is the dimension of keys and values, and 1/√d is a scaling factor.

With a total of H heads, multi-head attention can be regarded as an ensemble of each head h's self-attention:

A Transformer encoder contains a total of L layers, with each one including two sub-layers: a multi-head selfattention layer and a fully connected feed forward layer. All of the queries and key-value pairs in a same self-attention layer come from the output of the previous layer in the encoder. Therefore, each position in a layer of encoder can attend to all the positions in the previous layer. In our paper, we set L = 12.

For each layer, there is a multi-head attention layer. The multi-head attention in layer l can be written as:

3 ChoppingMethod

We perform chopping upon our base model - VisualBERT, which is a general Transformer style architecture without extra designs. Viewing one head as an element, we get an L × H = 12 × 12 grid-like structure for the model. The purpose of chopping is to disable a part of the model and make it unable to participate in the forward process. For heads and layers, we have separate strategies to reach this end, as they have different characteristics in cooperating. Heads work in parallel while layers work in a sequential manner.

3.1 Manual Chopping

In manual chopping, whether to chop a component or not is controlled by our signal manually. The signals are given at two levels: heads and layers.

Manual Chopping of Heads

The importance of each head is investigated thoroughly. Every time, We simply modified the multi-head attention in a specific layer l by multiplying one head attention by a binary mask variable αlh (h = 1, 2, ..., H):

where αlh can be either 0 or 1. The Transformer model we used have 12 layers with 12 heads in each layer, there are totally 144 αlh . For each time of removing one head, only one α is set to 0 and all the others are set to 1 (in Figure 1 (a)). As for removing all heads but one, only one α is set to 1 and all the others are set to 0 in a specific layer (in Figure 1 (b)). To study the contribution of each head, we exhaust each αlh in the above two experiments.

Manual Chopping of Layers

As the layers are in a serial manner, to chop off one specific Layer l, we directly skip it and connect the output of Layer l-1 to the input of Layer l+1. The detail is shown in Figure 1 (c). We chop each layer in the Transformer one time.

3.2 Automatic Dynamic Chopping

In contrast to manual chopping, we hope the model can automatically learn how important and relevant a specific layer is. Therefore, we design a dynamic chopping module to explore the effect of each layer l by learning a score Sl for each layer, then remove layers with scores below a given threshold. An instance level layer score is calculated by:

whereW is a linear layer, hidden statel is the hidden state output by layer l and we only take out the first element of this feature, which is also called [CLS] feature. 12 head scores qlh are averaged and put into a Sigmoid function to yield a value between 0 and 1, measuring the degree of importance for a specific layer. The score matrices will change when dealing with different instances.

Then we use this layer score Sl to determine whether to chop a layer or not. If Sl is below a threshold, then the Layer l will be chopped at the forward time. To chop the Layer l, we directly skip it and connect the output of Layer l - 1 to the input of the Layer l + 1.

The dynamic chopping module is trained independently without modifying any parameters in the Transformer model. The maximal learning rate is 1e-3 and the batch size is 480. This module is trained with a binary cross-entropy loss (for answer prediction) and an additional loss (L1 norm of head scores).

4 Experiments

4.1 Experimental Settings

We thoroughly study the effect of Transformer's heads and layers on the VQA task when questions are in different types. All of our experiments are conducted on the Task Driven Image Understanding Challenge (TDIUC) [Kafle and Kanan, 2017a] dataset, a large VQA dataset. This dataset was proposed to compensate for the bias in distribution of different question types of VQA 2.0 [Goyal et al., 2017]. It is larger than VQA 2.0 and divides questions into 12 more fine-grained categories based on the task they solve, which provide us convenience to do analysis by question types.

TDIUC has 12 question types, including both classical computer vision tasks and novel high-level vision tasks which require varying degrees of image understanding and reasoning. The total number of unique answers is 1, 618, which is then set as the size of our answer vocabulary. Experiments are conducted on 4 NVIDIA GeForce 2080Ti GPUs with a batch size of 480. We choose a general and popular pretraining model VisualBERT [Li et al., 2020], which is built on top of PyTorch. This model consists of visual embeddings & text embeddings, a one-stream Transformer and a classifier specific for VQA tasks. There are L = 12 layers of encoder, and each layer has H = 12 heads. The hidden state dimension is 768. We load the model pre-trained on COCO Caption [Chen et al., 2015] dataset, then finetune it with a leaning rate of 5e - 5 on the TDIUC dataset. The image features are extracted from a ResNet-101 [He et al., 2016] based Faster R-CNN [Ren et al., 2015] model pre-trained on Visual Genome [Krishna et al., 2017].

4.2 Effect of Manual Chopping

Removing One Head at One Time

First, we follow the previous work [Michel et al., 2019] to study if there is a specific role for a particular attention head h in different question types. To study the contribution of a particular head, we mask a single head (i.e., multiplying Attlh with zero-valued binary mask) in Transformer at a time, then evaluate the model's performance. If one head is important in a certain question type, the accuracy after the removal of this head will decrease dramatically. On the other hand, the increase in accuracy after removal can tell that this head is insignificant or even plays a negative role in this question type. In Table 1, we show the relative difference of accuracy ((Accnew - Accorg)/Accorg) on each type when only one head is removed. In each layer, there are 12 heads and we show the one who has the largest absolute value of relative difference to get a clue about the maximal influence of one layer. In most of the question types, the largest absolute value of relative difference out of 12 heads is still a small number, meaning some heads in this layer can be removed to get a more efficient model. In question types such as "Activity Recognition" and "Utilities & Affordances", removal of a specific head in more layers can bring a dramatic decrease in accuracy, which can illustrate that more layers (both low-level and high-level) are needed here. Observing the percentage of change in accuracy for each head (can be found in Appendix A), we found that in the majority of question types, removing one head does not hurt the accuracy a lot. Moreover, removal of some heads can even bring the improvement of accuracy, which means heads are not equally important in answering a specific question type.

Removing All Heads but One

It is also a question whether more than one head h is needed in a specific layer. We mask all heads in a layer except for a single one and show the relative difference of accuracy on each question type in Table 2. In each layer, we only show the largest value of relative difference out of 12 heads (max((Accinew -Acciorg )/Acciorg ), i ∈ [1, 12]). If there exists one head that can replace or even surpass all heads in a layer, this largest value will be a positive value. For some layers in a specific question type, one head is indeed sufficient at test time.

On the other hand, in some layers even the largest value of relative difference is a negative number, which means the high accuracy comes from the cooperative work of multiple heads. In question types Color Attributes and Activity Recognition, nearly none of a single head can replace a whole layer in all 12 layers. Besides, observing the relative difference of accuracy for each head (in Appendix B), we find that most of the heads in one layer tend to have the same level of influence. To further investigate whether there are some specific roles for different layers, we conduct experiments by removing a whole layer at a time below.

Removing One Layer

To investigate the effect of layers, we remove one layer in Transformer at a time. The dramatically decrease of accuracy illustrates that a layer plays a vital role in a question type. The more the accuracy decrease, the more important a layer is. While a positive relative difference tells that a layer is useless and can be removed. From the results in Figure 2, removing most of the layers cause a more or less decrease of accuracy, while other layers bring a slight increase of accuracy.

Results show that lower-level layers are enough for simple questions, such as questions belong to Absurd, Object Presence, etc.; higher-level layers are important in questions where more reasoning abilities are needed such as questions belong to Sentiment Understanding, Activity Recognition, etc.The 12 × 12 result matrix in Figure 2 also shows an interesting echelon shape that gradually changes according the level of reasoning. Questions belonging to Other Attributes (e.g., 'What shape is the clock?') and Object Utilities and Affordances (e.g., 'What object can be used to break glass?') may ask the materials or shapes of objects, which needs strong reasoning ability and common sense knowledge to answer them; Activity Recognition (e.g., 'What is the girl doing?') requires scene and activity understanding abilities; Positional Reasoning (e.g., 'What is to the left of the man on the sofa?') and Sentiment Understanding (e.g., 'How is she feeling?') are two kinds of well known complicated questions requiring multiple reasoning and understanding abilities. Removing highlevel layers will dramatically hurt accuracy in the abovementioned question types.

We also observe some other interesting phenomena. In question type Scene Recognition (e.g., 'What room is this?') removal of the first 5 layers causes decrease of accuracy in roughly a same manner, reflecting that they are all important for scene understanding. In question type Object Presence (e.g., 'Is there a cat in the image?') Layer 4 plays the most important role. Layer 4 is also important in other questions classes related to objects such as Subordinate Object Recognition, Counting (where the model needs to count the number of objects) and Color Attributes (locate the object firstly then recognize its color). The question type Absurd (i.e., Nonsensical queries about the image) is the easiest class, and it can be answered easily by determining whether the object mentioned in the question appears in the image, which also depends on objects and the removal of Layer 4 impacts on the accuracy greatly. Layer 3 and Layer 5 play a critical role in question type Counting; Layer 2 and Layer 5 are important in question type Color Attributes, which illustrates that some layers cooperate to work out special reasoning problems.

4.3 Effect of Dynamic Chopping

In our previous section, we analyze the effect of 'removing one head', 'removing all heads but one within a single layer' and 'directly removing one layer'. Based on the observations, we consider a single layer as a basic module in Transformer and want to further explore what would happen if we chop two or more layers at the same time. In the dynamic chopping module, a score is learned for each layer and we set several thresholds to determine whether to remove a layer or not, then we record the change in accuracy.

The mean of layer scores is collected in each type and the results can be found in Figure 3. The distribution of layer scores is similar to that of relative difference in manually removing one layer (Figure 2). When removing a layer hurts the accuracy a lot, the dynamic chopping module tends to learn a higher score for this layer.

We set several thresholds to find the best trade-off between accuracy and model size. Layers whose score is lower than the threshold value are removed. All the results can be found in Table 3 without fine-tuning. The percentage of parameters under the control of different thresholds are shown in Figure 4 and the detailed value of parameters can be found in Appendix D. For lower-level questions, the chopping curve is steeper, meaning it is easier to chop more layers. However, for higher-level questions, the chopping curve is smoother, as the chopping of layers needs stricter thresholds.

With larger thresholds, fewer parameters are kept by the VisualBERT model. When we set threshold to (0.3), the model achieves higher accuracy on half of the question types with roughly 76.92% of parameters reserved. When we reserve layers whose score is larger than 0.5, only 55.85% of parameters are reserved, with only 1% loss of accuracy. When we set threshold to (0.7), more parameters will be removed, with a significant impact on performance. However, to remove 50% layers randomly yields a much worse performance. The results show that our dynamic chopping module successfully predict scores to control which layer can be chopped efficiently. We also observe that one layer cannot replace the whole Transformer (Appendix C), and the result shows leaving only one layer in almost all question types cause over 50% of accuracy drop. All the results illustrate that not all layers can be reduced without significantly influencing performance. A qualitative example is also shown in Figure 5. After chopping insignificant layers under the threshold (0.5), the surviving layers can generate the true answer. Please refer to Appendix E and Appendix F for more examples and the analysis about failure cases.

5 Conclusion

In this paper, we conduct extensive experiments on a Visual Question Answering dataset categorized by question types and explore the roles of heads and layers in a Vision-andLanguage pre-trained Transformer by ablating these heads and layers separately. Interesting results as shown in an echelon shape reveal higher-level layers are required for higherlevel questions. A dynamic chopping module is further designed to automatically learn efficiently chopped model. We hope our findings to advance the understanding of Transformer structure under different level of visual reasoning question types and inspire the investigation of more accurate and more efficient visual reasoning models.