Disentangling Representations of Text by Masking Transformers (2021)

Abstract: Representations from large pretrained models such as BERT encode a range of features into monolithic vectors, affording strong predictive accuracy across a range of downstream tasks. In this paper we explore whether it is possible to learn disentangled representations by identifying existing subnetworks within pretrained models that encode distinct, complementary aspects. Concretely, we learn binary masks over transformer weights or hidden units to uncover subsets of features that correlate with a specific factor of variation; this eliminates the need to train a disentangled model from scratch for a particular task. We evaluate this method with respect to its ability to disentangle representations of sentiment from genre in movie reviews, toxicity from dialect in Tweets, and syntax from semantics. By combining masking with magnitude pruning we find that we can identify sparse subnetworks within BERT that strongly encode particular aspects (e.g., semantics) while only weakly encoding others (e.g., syntax). Moreover, despite only learning masks, disentanglement-via-masking performs as well as — and often better than —previously proposed methods based on variational autoencoders and adversarial training.

## introduction and motivation large pretrained models such as elmo @xcite , bert @xcite , and xl-net @xcite have come to dominate modern nlp. such models rely on self-supervision over large datasets to learn general-purpose representations of text that achieve strong predictive performance across a spectrum of downstream tasks @xcite . a downside of such learned representations is that it is not obvious what information they encode, which hinders model robustness and interpretability. the opacity of embeddings produced by models such as bert has motivated nlp research on designing probing tasks as a means of uncovering the properties of input texts that are encoded in learned representations @xcite @xcite . 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 inputs. explicitly factorizing representations aids interpretability, in the sense that it becomes more straightforward to determine which factors of variation inform predictions in downstream tasks. 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 causally linked @xcite . this is particularly important for large pretrained models like bert, as we do not know what the representations produced by such models encode. here, learning disentangled representations may facilitate increased robustness under distributional shifts by capturing a notion of invariance: if syntactic changes 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. as one example that we explore in this paper, consider the task of identifying tweets that contain hate speech @xcite . recent work shows that models trained over tweets annotated on a toxicity scale exhibit a racial bias: they have a tendency to over-predict that tweets written by users who self-identify as black are "toxic", owing to the use of african american vernacular english (aave; @xcite . in principal, disentangled representations would allow us to isolate relevant signal from irrelevant or spurious factors (such as, in this case, the particular english dialect used), which might in turn reveal and allow us to mitigate unwanted system biases, and increase robustness. to date, most research on disentangled representations has focused on applications in computer vision @xcite @xcite , where there exist comparatively clear independent factors of variation such as size, position, and orientation, which have physical grounding and can be formalized in terms of actions of symmetry subgroups @xcite . a challenge in learning disentangled representations of text is that it is less clear which factors of variation should admit invariance. still, we may hope to disentangle particular properties for certain applications -e.g., protected demographic information @xcite -and there are general properties of language that we might hope to disentangle, e.g., syntax and semantics @xcite . ## 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 m (a) and m (b) , we assume access to triplets (x 0 , x 1 , x 2 ) in which x 0 and x 1 are similar with respect to aspect a but dissimilar with respect to aspect b, whereas x 0 and x 2 are similar with respect to aspect b but dissimilar with respect to aspect a. in some of our experiments (e.g., when disentangling sentiment from genre in movie reviews) we further assume that we have access to class labels y (a) ∈ {0, 1} and y (b) ∈ {0, 1} for as-pects of interest. in such cases, we build triplets using these labels, defining (x 0 , x 1 , x 2 ) such that y ## experiments we conduct a series of experiments to evaluate the degree to which the proposed masking strategy achieves disentanglement, as compared to existing methods for disentanglement in nlp. as a first illustrative example, we consider a corpus of movie reviews, in which sentiment is correlated with film genre (3.1). we treat this as a proxy for a spurious correlation, and evaluate the robustness of the models to shifts in conditional probabilities of one attribute (sentiment) given another (genre). we then consider a more consequential example: hate speech classification on twitter (3.2). prior work @xcite has shown that models exploit a spurious correlation between "toxicity" and african american vernacular english (aave); we aim to explicitly disentangle these factors in service of fairness. we evaluate whether the model is able to achieve equalized odds, a commonly used fairness metric. finally, following prior work, we investigate disentangling semantics from syntax (insofar as this is possible) in section 3.3. ## discussion we have presented a novel perspective on learning disentangled representations for natural language processing in which we attempt to uncover existing subnetworks within pretrained transformers (e.g., bert) that yield disentangled representations of text. we operationalized this intuition via a masking approach, in which we estimate only binary masks over weights or hidden states within bert, leaving all other parameters unchanged. we demonstrated that -somewhat surprisingly -we are able to achieve a level of disentanglement that often exceeds existing approaches (e.g., a varational auto-encoder on top of bert), which have the benefit of finetuning all model parameters. our experiments demonstrate the potential benefits of this approach. in section 3.1 we showed that disentanglement via masking can yield representations that are comparatively robust to shifts in correlations between (potentially sensitive) attributes and target labels. aside from increasing robustness, finding sparse subnetworks that induce disentangled representations constitutes a new direction to pursue in service of providing at least one type of model interpretability for nlp. finally, we note that sparse masking (which does not mutate the underlying transformer parameters) may offer