Disentangling Representations of Text by Masking Transformers

1 INTRODUCTION AND MOTIVATION . Large-scale pretrained models such as ELMo ( Peters et al. , 2018 ) , BERT ( Devlin et al. , 2018 ) , and XLNet ( Yang et al. , 2019 ) have come to dominate in modern natural language processing ( NLP ) . Such models rely on self-supervision over large datasets to learn general-purpose representations of text that achieve competitive predictive performance across a spectrum of downstream tasks ( Liu et al. , 2019 ) . A downside of such learned representations is that it is not obvious what information they are encoding , which hinders model robustness and interpretability . The opacity of representations produced by models such as BERT has motivated a line of NLP research on designing probing tasks as a means of uncovering what properties of input texts are encoded into token- and sentencelevel representations ( Rogers et al. , 2020 ; Linzen et al. , 2019 ; Tenney et al. , 2019 ) . In this paper we investigate whether we can uncover disentangled representations from pretrained models . That is , rather than mapping inputs onto a single vector that captures arbitrary combinations of features , our aim is to extract a representation that factorizes into distinct , complementary properties of the input . An advantage of explicitly factorizing representations is that it aids interpretability , in the sense that it becomes more straightforward to determine what factors of variation inform predictions in downstream tasks . Further , disentangled representations may facilitate increased robustness under distributional shifts by capturing a notion of invariance : If syntactic changes in a sentence do not affect the representation of semantic features , and vice versa , then we can hope to learn models that are less sensitive to any incidental correlations between these factors . A general motivation for learning disentangled representations is to try and minimize — or at least expose — model reliance on spurious correlations , i.e. , relationships between ( potentially sensitive ) attributes and labels that exist in the training data but which are not casually linked ( Kaushik et al. , 2020 ) . This is particularly relevant in the context of large pretrained models like BERT , as we do not know what the representations produced by such models encode . To date , most research on disentangled representations has focused on applications in computer vision ( Locatello et al. , 2019b ; Kulkarni et al. , 2015 ; Chen et al. , 2016 ; Higgins et al. , 2016 ) , where there exist comparatively clear independent factors of variation such as size , position , color , and orientation , which have physical grounding and can be formalized in terms of actions of symmetry subgroups ( Higgins et al. , 2018 ) . A challenge in learning disentangled representations of text is that it is more ambiguous what factors of variation should admit invariance . Still , we may hope to disentangle particular properties for certain applications ( e.g. , sentiment , or perhaps protected demographic information ( Elazar & Goldberg , 2018 ) ) , and there are also general properties of language that we might hope to disentangle , e.g. , syntax and semantics ( Chen et al. , 2019 ) . In this paper we ask whether complementary factors of variation might already be captured by pretrained models , and whether it is possible to uncover these by identifying appropriate “ subnetworks ” . The intuition for this hypothesis is that generalization across a sufficiently large and diverse training set may implicitly necessitate representations that admit some notion of invariance , as the many factors of variation in the training data give rise to a combinatorial explosion of possible inputs . Intriguing prior work ( Radford et al. , 2017 ) examining the correlation between sentiment and individual nodes within pretrained networks offers some additional support for this intuition . To test this hypothesis , we propose to use masking as a mechanism to isolate representations of individual factors . Recent work on lottery tickets ( Frankle & Carbin , 2018 ) suggest that overparameterized networks are redundant , in that a network reduced to a small subset of weights set to “ winning ” initial values can achieve predictive performance similar to the full network . Building on this intuition , we hypothesize that it might be possible to uncover a representation for a factor of interest by starting with a pretrained representation and simply masking out weights or hidden units that correlate with other factors of variation . We use BERT ( Devlin et al. , 2018 ) as an archetypal pretrained transformer to test two variants of this basic idea . In the first variant we learn binary masks for all weight matrices in the model ; in the second we derive masks for all hidden units ( intermediate representations ) . To learn these masks we minimize a triplet loss that encourages the resultant representations for instances that are similar with respect to an aspect of interest to be relatively near to one another , independent of other factors . Our approach of uncovering existing subnetworks within pretrained models that yield disentangled representations differs substantially from prior work on disentangling representations in NLP , which have either relied on adversarial debiasing approaches ( Elazar & Goldberg , 2018 ; Barrett et al. , 2019 ) or variational auto-encoders ( Chen et al. , 2019 ; Esmaeili et al. , 2019 ) . We evaluate masking in the context of two tasks . The first is a setting in which we aim to disentangle a representation of features for a target task from that of information encoding a secondary , non-target attribute ( e.g. , this might be sensitive information , or simply an unrelated factor ) . In the second we follow prior work in attempting to induce representations of semantics and syntax , respectively . In both settings , our surprising finding is that masking alone often outperforms previously proposed approaches ( which learn or finetune networks in their entirety ) . While a small amount of masking generally suffices to achieve disentanglement , we can further increase sparsity by combining masking with weight pruning . The main contributions of this paper are as follows . ( 1 ) We propose a novel method of disentangling representations in NLP : Masking weights or hidden units within pretrained transformers ( here , BERT ) . ( 2 ) We empirically demonstrate that we are indeed able to identify sub-networks within pretrained transformers that yield disentangled representations that outperform existing approaches ( which finetune all model weights ) . ( 3 ) We show that masking can be combined with weight-pruning techniques to learn sparse subnetworks for each factor of interest . 2 METHODS . We are interested in learning a disentangled representation that maps inputs x ( text ) onto vectors z ( a ) and z ( b ) that encode two distinct factors of variation . To do so , we will learn two sets of masks M ( a ) and M ( b ) that can be applied to either the weights or the intermediate representations in a pretrained model ( in our case , BERT ) . We estimate only the mask parameters and do not finetune the weights of the pretrained model . To learn the masks M ( a ) and M ( b ) , we assume access to triplets ( x0 , x1 , x2 ) in which x0 and x1 are similar with respect to aspect a but dissimilar with respect to aspect b , whereas x0 and x2 are similar with respect to aspect b but dissimilar with respect to aspect a . In some of our experiments ( i.e. , when disentangling the sentiment from the genre in movie reviews ) we further assume that we have access to class labels y ( a ) ∈ { 0 , 1 } and y ( b ) ∈ { 0 , 1 } for the aspects of interest . 2.1 MASKING WEIGHTS AND HIDDEN ACTIVATIONS . Figure 1 illustrates the two forms of masking that we consider in our approach . We depict a single linear layer of the model , h = ( h ( a ) , h ( b ) ) are the input activations , W are the weights in the pretrained model,1 and h′ = ( h′ ( a ) , h′ ( b ) ) are the output activations . We augment each layer of the original network with two masks M = ( M ( a ) , M ( b ) ) , applied in two ways : 1 . Masking Weights When training BERT with masked weights , masks M ( a ) and M ( b ) have the same shape as weights W , and output activations are computed using the masked tensor of weights h′ = h · ( W ◦M ) . ( 1 ) 2 . Masking Hidden Activations When training BERT with masked hidden units , the masksM ( a ) and M ( b ) have the same shape as the intermediate ( hidden ) activations h ( a ) and h ( b ) . Output activations are computed by applying the original weights W to masked inputs h′ = ( h ◦M ) ·W . ( 2 ) In both methods , we follow Zhao et al . ( 2020 ) and only mask the last several layers of BERT , leaving the bottom layers ( including the embedding layer ) unchanged . In practice we find masking the last 9 layers to work the best . 2.2 TRIPLET LOSS . To learn the masks , we assume that we have access to supervision in the form of triplets , as introduced above . Passing ( x0 , x1 , x2 ) through our model yields two representations of each instance : ( z ( a ) 0 , z ( b ) 0 ) , ( z ( a ) 1 , z ( b ) 1 ) , ( z ( a ) 2 , z ( b ) 2 ) , for which we define the loss L ( a ) triplet = max ( ‖z ( a ) 0 − z ( a ) 1 ‖ − ‖z ( a ) 0 − z ( a ) 2 ‖+ , 0 ) , ( 3 ) L ( b ) triplet = max ( ‖z ( b ) 0 − z ( b ) 2 ‖ − ‖z ( b ) 0 − z ( b ) 1 ‖+ , 0 ) , ( 4 ) Ltriplet = 1 2 ( L ( a ) triplet + L ( b ) triplet ) . ( 5 ) Here is a hyperparameter specifying a margin for the loss , which we set to = 2 in all experiments . 2.3 SUPERVISED LOSS . In some settings we may have access to more direct forms of supervision . For example , when learning representations for the genre and sentiment in a movie review , we have access to class 1We omit the bias term in the linear layer , for which no masking is performed . labels y ( a ) and y ( b ) for each aspect . To make full use of such supervision when available , we add two linear classification layers C ( a ) and C ( b ) to our model and compute the classification losses L ( a ) cls = CrossEntropy ( C ( a ) ( z ( a ) ) , y ( a ) ) , L ( b ) cls = CrossEntropy ( C ( b ) ( z ( b ) ) , y ( b ) ) , ( 6 ) Lcls = 1 2 ( L ( a ) cls + L ( b ) cls ) . ( 7 )

Summary: The paper proposes a procedure to extract disentangled representations from pretrained BERT models. In particular, the paper proposes learning binary masks over BERT weights (or, as an alternative, over BERT activations) such that the resulting representations correspond to the desired aspect representations. The model requires additional supervision (binary labels or example triplets) and training (for the masks but not the BERT weights). The experiments aim to perform disentangling to ensure that (1) the learned representation does not “leak” a potentially sensitive attribute, and (2) the downstream classifier’s performance is good across all subgroups formed by the attributes. The experiments show that the proposed method outperforms baselines such as unmasked BERT, unmasked-but-finetuned BERT, and unmasked-but-adversarially-finetuned BERT.