Model Compression for Domain Adaptation through Causal Effect Estimation (2021)

Abstract: Recent improvements in the predictive quality of natural language processing systems are often dependent on a substantial increase in the number of model parameters. This has led to various attempts of compressing such models, but existing methods have not considered the differences in the predictive power of various model components or in the generalizability of the compressed models. To understand the connection between model compression and out-of-distribution generalization, we define the task of compressing language representation models such that they perform best in a domain adaptation setting. We choose to address this problem from a causal perspective, attempting to estimate the average treatment effect (ATE) of a model component, such as a single layer, on the model’s predictions. Our proposed ATE-guided Model Compression scheme (AMoC), generates many model candidates, differing by the model components that were removed. Then, we select the best candidate through a stepwise regression model that utilizes the ATE to predict the expected performance on the target domain. AMoC outperforms strong baselines on dozens of domain pairs across three text classification and sequence tagging tasks.1

## introduction the rise of deep neural networks has transformed the way we represent language, allowing models to learn useful features directly from raw inputs. however, recent improvements in the predictive quality of language representations are often related to a substantial increase in the number of model parameters. indeed, the introduction of the transformer architecture @xcite and attention-based models @xcite @xcite have improved performance on most natural language processing (nlp) tasks, while facilitating a large increase in model sizes. since large models require a significant amount of computation and memory during training and inference, there is a growing demand for compressing such models while retaining the most relevant information. while recent attempts have shown promising results @xcite , they have some limitations. specifically, they attempt to mimic the behavior of the larger models without trying to understand the information preserved or lost in the compression process. in compressing the information represented in billions of parameters, we identify three main challenges. first, current methods for model compression are not interpretable. while the importance of different model parameters is certainly not uniform, it is hard to know a priori which of the model components should be discarded in the compression process. this notion of feature importance has not yet trickled down into compression methods, and they often attempt to solve a dimensionality reduction problem where a smaller model aims to mimic the predictions of the larger model. nonetheless, not all parameters are born equal, and only a subset of the information captured in the network is actually useful for generalization @xcite . the second challenge we observe in model compression is out-of-distribution generalization. typically, compressed models are tested for their in-domain generalization. however, in reality the distribution of examples often varies and is different than that seen during training. without testing for the generalization of the compressed models on different test-set distributions, it is hard to fully assess what was lost in the compression process. the setting explored in domain adaptation provides us with a platform to test the ability of the compressed models to generalize across-domains, where some information that the model has learned to rely on might not exist. strong model performance across domains provides a stronger signal on retaining valuable information. lastly, another challenge we identify in training and selecting compressed models is confidence estimation. in trying to understand what gives large models the advantage over their smaller competitors, recent probing efforts have discovered that commonly used models such as bert @xcite , learn to capture semantic and syntactic information in different layers and neurons across the network @xcite . while some features might be crucial for the model, others could learn spurious correlations that are only present in the training set and are absent in the test set @xcite . such cases have led to some intuitive common practices such as keeping only layers with the same parity or the top or bottom layers @xcite . those practices can be good on average, but do not provide model confidence scores or success rate estimates on unseen data. our approach addresses each of the three main challenges we identify, as it allows estimating the marginal effect of each model component, is designed and tested for out-of-distribution generalization, and provides estimates for each compressed model performance on an unlabeled target domain. we dive here into the connection between model compression and out-of-distribution generalization, and ask whether compression schemes should consider the effect of individual model components on the resulting compressed model. particularly, we present a method that attempts to compress a model while maintaining components that can generalize well across domains. inspired by causal inference @xcite , our compression scheme is based on estimating the average effect of model components on the decisions the model makes, at both the source and target domains. in causal inference, we measure the effect of interventions by comparing the difference in outcome between the control and treatment groups. in our setting, we take advantage of the fact that we have access to unlabeled target examples, and treat the model's predictions as our outcome variable. we then try to estimate the effect of a subset of the model components, such as one or more layers, on the model's output. to do that, we propose an approximation of a counterfactual model where a model component of choice is removed. we train an instance of the model without that component and keep everything else equal apart from the input and output to that component, which allows us to perform only a small number of gradient steps. using this approximation, we then estimate the average treatment effect (ate) by comparing the predictions of the base model to those of its counterfactual instance. since our compressed models are very efficiently trained, we can generate a large number of such models per each source-target domain pair. we then train a regression model on our training domain pairs in order to predict how well a compressed model would generalize from a source to a target domain, using the ate as well as other variables. this regression model can then be applied to new source-target domain pairs in order to select the compressed model that best supports cross-domain generalization. to organize our contributions, we formulate three research questions: in § 6 we directly address each of the three research questions, and demonstrate the usefulness of our method, ate-guided model compression (amoc), to improve model generalization. ## previous work previous work on the intersection of neural model compression, domain adaptation, and causal inference is limited, as our application of causal inference to model compression and our discussion of the connection between compression and cross-domain generalization are novel. however, there is an abundance of work in each field on its own, and on the connection between domain adaptation and causal inference. since our goal is to explore the connection between compression and out-of-distribution generalization, as framed in the setting of domain adaptation, we survey the literature on model compression and the connection between generalization, causality, and domain adaptation. ## causal terminology causal methodology is most commonly used in cases where the goal is estimating effects on real-world outcomes, but it can be adapted to help us understand and explain what affects nlp models @xcite . specifically, we can think of intervening on a model and altering its components as a causal question, and measure the effect of this intervention on model predictions. a core benefit of this approach is that we can estimate treatment effects on model's predictions without the need for manually-labeled target data. borrowing causal methodology into our setting, we treat model components as our treatment, and try to estimate the effect of removing them on our model's predictions. the predictions of a model are driven by its components, and by changing one component and holding everything else equal, we can estimate the effect of this intervention. we can use this estimation in deciding which model component should be kept in the compression process. as the link between model compression and causal inference was not explored previously, we provide here a short introduction to causal inference and its basic terminology, focusing on its application to our use case. we then discuss the connection to pearl's do-operator @xcite and the estimation of treatment effects. imagine we have a model m that classifies examples to one of l classes. given a set c of k model components, which we hypothesize might affect the model's decision, we denote the set of binary variables i @xmath0, 1}|@xmath1, . . . , k}}, where each corresponds to the inclusion of the component in the model, that is, if i c @xmath2. our goal is to assert how the model's predictions are affected by the components in c. as we are interested in the effect on the class probability assigned by m, we measure this probability for an example x, and denote it for a class l as z(m(x)) l and for all l classes as z(m(x)). using this setup, we can now define the ate, the common metric used when estimating causal effects. ate is the difference in mean outcomes between the treatment and control groups, and using do-calculus @xcite we can define it as follows: ## methodology we start by describing the task of compressing models such that they perform well on out-of-distribution examples, detailing the domain adaptation framework we focus on. then, we describe our compression scheme, designed to allow us to approximate the ate and responsible for producing compressed model candidates. finally, we propose a regression model that uses the ate and other features to predict a candidate model's performance on a target domain. this regression allows us to select a strong candidate model. ## additional analysis 7.1 layer importance to further understand the importance of each of bert's layers, we compute the frequency in which each layer appears in the best candidate model, namely, the model with the highest f1 score on the target test set, of every experiment. figure 3 captures the layer frequencies across the different data sets and across the number of removed layers. the plots suggest that the two final layers, layers 11 and 12, are the least important layers with average frequencies of 30.3% and 24.8%, respectively. additionally, in most cases layer 1 is ranked below the other layers. these results imply that the compressed models are able to better recover from the loss of parameters when the external layers are removed. the most important layer appears to be layer 4, with an average frequency of 73.3%. finally, we notice that a large frequency variance exists across the different subplots. such variance supports our hypothesis that the decision of which layers to remove should not be based solely on the architecture of the model. to pin down the importance of a specific layer for a given base model, we utilize a similar regression analysis to that of § 6. specifically, we train a regression model on all compressed candidates for a given source-target domain pair (in all three tasks), adding indicator variables for the exclusion of each layer from the model. this model associates each layer with a regression coefficient, which can be interpreted as the marginal effect of that layer being removed on expected target performance. we then compute for each layer its average coefficient across source-target pairs ( as can be seen in the table, layers that their removal is associated with better model performance are more often not included in the best performing compressed models. indeed, the spearman's rank correlation between the two rankings is as high as 0.924. such analysis demonstrates that the regression model used as part of amoc not only selects high quality candidates, but can also shed light on the importance of individual layers. ## conclusion we explored the relationship between model compression and out-of-distribution generalization. amoc, our proposed algorithm, relies on causal inference tools for estimating the effects of interventions. it hence creates an interpretable process that allows to understand the role of specific model components. our results indicate that amoc is able to produce a smaller model with minimal loss in performance across domains, without any use of target labeled data at test time (q1). amoc can efficiently train a large number of compressed model candidates, that can then serve as training examples for a regression model. we have shown that this approach results in a high quality estimation of the performance of compressed models on unseen target domains (q2). moreover, our stepwise regression analysis indicates that the at e s and at e t estimates are instrumental for these attractive properties (q3). as training and test set mismatches are common, we steered our model compression research towards out-of-domain generalization. besides its realistic nature, this setup poses additional modeling challenges, such as understanding the proximity between domains, identifying which components are invariant to domain shift, and estimating performance on unseen domains. hence, amoc is designed for model compression in the out-of-distribution setup. we leave the design of similar in-domain compression methods for future work. finally, we believe that using causal methods to produce compressed nlp models that can well generalize across distributions is a promising direction of research, and hope that more work will be done in this intersection.