W hat’s in Your Head? E mergent Behaviour in Multi-Task Transformer Models (2021)

Abstract: The primary paradigm for multi-task training in natural language processing is to represent the input with a shared pre-trained language model, and add a small, thin network (head) per task. Given an input, a target head is the head that is selected for outputting the final prediction. In this work, we examine the behaviour of non-target heads, that is, the output of heads when given input that belongs to a different task than the one they were trained for. We find that non-target heads exhibit emergent behaviour, which may either explain the target task, or generalize beyond their original task. For example, in a numerical reasoning task, a span extraction head extracts from the input the arguments to a computation that results in a number generated by a target generative head. In addition, a summarization head that is trained with a target question answering head, outputs query-based summaries when given a question and a context from which the answer is to be extracted. This emergent behaviour suggests that multi-task training leads to non-trivial extrapolation of skills, which can be harnessed for interpretability and generalization.

## introduction the typical framework for training a model in natural language processing to perform multiple tasks is to have a shared pre-trained language model (lm), and add a small, compact neural network, often termed head, on top of the lm, for each task @xcite @xcite . the heads are trained in a supervised manner, each on labelled data collected for the task it performs @xcite . at inference time, the output is read out of a selected target head, while the outputs from the other heads are discarded (figure 1 ). what is the nature of predictions made by nontarget heads given inputs directed to the target head? one extreme possibility is that the pre- in this work, we test whether such interactions occur in multi-task transformer models, and if nontarget heads decode useful information given inputs directed to the target head. we show that multihead training leads to a steering effect, where the target head guides the behaviour of the non-target head, steering it to exhibit emergent behaviour, which can explain the target head's predictions, or generalize beyond the task the non-target head was trained for. we study the "steering effect" in three multihead models (figure 2 ). in a numerical reading comprehension task @xcite , the model is given a question and paragraph and either uses an extractive head to output an input span, or a generative head to generate a number using arithmetic operations over numbers in the input (figure 2 , left). treating the extractive head as the non-target head, we observe that it tends to output the arguments to the arithmetic operation performed by the decoder, and that successful argument extraction is correlated with higher performance. moreover, we perform interventions @xcite , where we modify the representation based on the output of the extractive head, and show this leads to predictable changes in the behaviour of the generative head. thus, we can use the output of the non-target head to improve interpretability. we observe a similar phenomenon in multi-hop question answering (qa) model @xcite , where a non-target span extraction head outputs supporting evidence for the answer predicted by a classification head (figure 2 , center). this emerging interpretability is considerably different from methods that explicitly train models to output explanations @xcite . beyond interpretability, we observe non-trivial extrapolation of skills when performing multi-task training on extractive summarization @xcite and multi-hop qa (figure2, right). specifically, a head trained for extractive summarization outputs supporting evidence for the answer when given a question and a paragraph, showing that multi-task training steers its behaviour towards query-based summarization. we show this does not happen in lieu of multi-task training. to summarize, we investigate the behaviour of non-target heads in three multi-task transformer models, and find that without any dedicated training, non-target heads provide explanations for the predictions of target heads, and exhibit capabilities beyond the ones they were trained for. this extrapolation of skills can be harnessed for many applications. for example, teaching models new skills could by done by training on combinations of tasks different from the target task. this would be useful when labeled data is not available or hard to collect. also, training an additional head that extracts information from the input could be applied as a general practice for model interpretability. ## multi-head transformer models the prevailing method for training models to perform nlp tasks is to add parameter-thin heads on top of a pre-trained lm, and fine-tune the entire network on labeled examples @xcite . given a text input with n tokens @xmath0, . . . , x n , the model first computes contextualized representations @xmath1, . . . , h @xmath2. these representations are then fed into the output heads, with each head o estimating the probability p ψo (y | h) of the true output y given the encoded input h and the parameters ψ o of o. the head that produces the final model output, termed the target head, is chosen either deterministically, based on the input task, or predicted by an output head classifier p(o | x). predictions made by nontarget heads are typically ignored. when p(o | x) is deterministic it can be viewed as an indicator function for the target head. training multi-head transformer models is done by marginalizing over the set of output heads o, and maximizing the probability for a head o, we denote by s o the set of examples (x, y) such that y is in the output space of o, and by so the other training examples. the sets s o and so may consist of examples from different tasks (e.g., question answering and text summarization), or of examples from the same task but with different output formats (e.g., yes/no vs. span extraction questions). our goal is to evaluate the predictions of o on examples from so , for which another head o is the target head, and the relation between these outputs and the predictions of o . in the next sections, we will show that the predictions of o interact with those of o . we will denote by o t the target head, and by o s the steered head. ## overview: experiments & findings this section provides an overview of our experiments, which are discussed in detail in §4, §5, §6. given a model with a target head o t and a steered head o s , our goal is to understand the behaviour of o s on inputs where o t provides the prediction. to this end, we focus on head combinations, where o s is expressive enough to explain the outputs of o t , but unlike most prior work aiming to explain by examining model outputs @xcite @xcite , o s is not explicitly trained for this purpose. concretely, our analysis covers three settings, illustrated in figure 2 and summarized in the first setting (figure 2 left, and §4) considers a model with generative and extractive heads, trained on the drop dataset @xcite for numerical reasoning over text. surprisingly, we observe that the arguments for the arithmetic computation required for the generative head to generate its answer often emerge in the outputs of the extractive head. the second setting (figure 2 middle, and §5) considers a model with a classification head outputting 'yes'/'no' answers, and a span extraction head, trained on the hotpotqa dataset @xcite for multi-hop reasoning. the outputs of the extractive head once again provide explanations in the form of supporting facts from the input context. the last setting (figure 2 right, and §6) considers a model with two extractive heads, one for span extraction and another for (sentence-level) extractive summarization. each head is trained on a different dataset; hotpotqa for span extraction and cnn/dailymail for summarization @xcite . we find that the summarization head tends to extract the supporting facts given inputs from hotpotqa, effectively acting as a query-based summarization model. we now present the above settings. ## setting 1: emerging computation arguments in span extraction we start by examining a combination of generative and extractive heads (figure 2 , left), and analyze the spans extracted from the input when the generative head is selected to output the final answer. ## setting 2: emerging supporting facts in span extraction we now consider a combination of an extractive head and a classification head (figure 2 , middle). ## setting 3: emerging query-based summaries in §4 and §5, we considered models with output heads trained on examples from the same data distribution. would the steering effect occur when output heads are trained on different datasets? we now consider a model trained to summarize text and answer multi-hop questions (figure 2 , right). ## related work transformer models with multiple output heads have been widely employed in previous works (hu 4 wilcoxon signed-rank test with p-value 0.001. and singh, 2021; @xcite @xcite . to the best of our knowledge, this is the first work that analyzes the outputs of the non-target heads. previous work used additional output heads to generate explanations for model predictions @xcite @xcite . specifically, recent work has explored utilization of summarization modules for explainable qa @xcite . in the context of summarization, xu and lapata (2020) have leveraged qa resources for training query-based summarization models. hierarchies between nlp tasks have also been explored in multi-task models not based on transformers @xcite @xcite . contrary to previous work, the models in this work were not trained to perform the desired behaviour. instead, explanations and generalized behaviour emerged from training on multiple tasks. a related line of research has focused on developing probes, which are supervised network modules that predict properties from model representations @xcite @xcite . a key challenge with probes is determining whether the information exists in the representation or is learned during probing @xcite @xcite . unlike probes, steered heads are trained in parallel to target heads rather than on a fixed model. moreover, steered heads are not designed to decode specific properties from representations, but their behaviour naturally extends beyond their training objective. our findings also relate to explainability methods that highlight parts from the input via the model's attention @xcite , and extract rationales through unsupervised training @xcite . the emerging explanations we observe are based on the predictions of a head rather than on internal representations. ## conclusions and discussion we show that training multiple heads on top of a pre-trained language model creates a steering effect, where the target head influences the behaviour of another head, steering it towards capabilities beyond its training objective. in three multi-task settings, we find that without any dedicated training, the steered head often outputs explanations for the model predictions. moreover, modifying the input representation based on the outputs of the steered head can lead to predictable changes in the target head predictions. our findings provide evidence for extrapolation of skills as a consequence of multi-task training, opening the door to new research directions in interpretability and generalization. future work could explore additional head combinations, in order to teach models new skills that can be cast as an extrapolation of existing tasks. in addition, the information decoding behaviour observed in this work can serve as basis for developing general interpretability methods for debugging model predictions. a natural question that arises is what head combinations lead to a meaningful steering effect. we argue that there are two considerations involved in answering this question. first, the relation between the tasks the heads are trained on. the tasks should complement each other (e.g. summarization and question answering), or the outputs of one task should be expressive enough to explain the outputs of the other task, when applied to inputs of the other task. for example, extractive heads are particularly useful when the model's output is a function of multiple input spans. another consideration is the inputs to the heads. we expect that training heads with similar inputs (in terms of length, language-style, etc.) will make the underlying language model construct similar representations, thus, increasing the probability of a steering effect between the heads. 1 2 3 4 5 6 7 8 9 10 11 12 13 b distribution of extracted spans by msegenbert b.1 distribution of extracted and annotated spans from drop figure 5 shows the number of extracted spans by msegenbert compared to the annotated spans, for a subset of 400 examples from the development set of drop. on average, msegenbert extracts a similar number of spans per example (2.12) compared to the spans extracted by annotators (1.95). however, msegenbert tends to over-predict one span compared to the annotated examples.