Intermediate Entity-based Sparse Interpretable Representation Learning (2022)

Abstract: Interpretable entity representations (IERs) are sparse embeddings that are “human-readable” in that dimensions correspond to fine-grained entity types and values are predicted probabilities that a given entity is of the corresponding type. These methods perform well in zero-shot and low supervision settings. Compared to standard dense neural embeddings, such interpretable representations may permit analysis and debugging. However, while fine-tuning sparse, interpretable representations improves accuracy on downstream tasks, it destroys the semantics of the dimensions which were enforced in pre-training. Can we maintain the interpretable semantics afforded by IERs while improving predictive performance on downstream tasks? Toward this end, we propose Intermediate enTity-based Sparse Interpretable Representation Learning (ItsIRL). ItsIRL realizes improved performance over prior IERs on biomedical tasks, while maintaining “interpretability” generally and their ability to support model debugging specifically. The latter is enabled in part by the ability to perform “counterfactual” fine-grained entity type manipulation, which we explore in this work. Finally, we propose a method to construct entity type based class prototypes for revealing global semantic properties of classes learned by our model. Code for pre-training and experiments will be made publicly available.

## introduction deep pre-trained models yield sota performance on a range of nlp tasks, but do so by learning and exploiting dense continuous representations of inputs which complicate model interpretation. that is, the dimensions in learned representations have no a priori semantics, and consequently are not directly human readable. indeed, this has inspired an entire line of work on "probing" dense representations to recover the implicit knowledge stored within them @xcite . an alternative is to design architectures that explicitly imbue embeddings with semantics. to this end, recent work has proposed learning highdimensional sparse interpretable entity representations (iers) for general and biomedical domains @xcite . iers are composed of a transformerbased @xcite entity typing model with a corresponding fine-grained static type system that accepts an entity mention and its context, and outputs individual probabilities that the mention is an instance of the respective types. these embeddings may then be used as features for downstream tasks. iers afford a variety of model transparency (dimensions have semantics) which may facilitate model debugging and/or instill confidence in model outputs. for example, if one defines a linear layer on top of entity-type representations, learned coefficients are interpretable as weights assigned to specific entity types. one could learn rules or manually debug models by reviewing incorrect predictions and inspecting the corresponding induced representations to identify potentially systematic erroneous type assignments. in addition to providing this type of interpretability, iers have been shown to perform comparatively well in zero-and few-shot settings @xcite . a limitation of iers is that they do not naturally permit fine-tuning, because doing so destroys the semantically meaningful entity typing representations learned during pre-training. this requirement is a limitation because fine-tuned models will in general achieve stronger predictive performance when supervision is available. in this work we aim to improve the predictive performance of iers without sacrificing their interpretability. specifically, we propose intermediate entity-based sparse interpretable representation learning (itsirl). we show that this model outperforms prior iers by a substantial margin on experiments over biomedical datasets -a domain where interpretability is often paramount -while providing natural mechanisms for model debugging by virtue of the representational semantics inherent to the architecture. we then propose a counterfactual analysis of our intermediate interpretable layer to measure the effect of entity type manipulation on downstream predictions. this intervention is made possible by virtue of the model design. using manually constructed, class-specific entity type sets we show that this intervention can be used to fix errors made by the proposed itsirl model automatically, ultimately allowing the model to outperform dense (uninterpretable) models in terms of test accuracy. we then propose a method in which we combine entity types over classes on training data to create positive and negative class prototypes that can be used to better understand the "global" semantics learned by itsirl for downstream tasks. our specific contributions are as follows: the decoding step can be fine-tuned for specific tasks. ## background: interpretable entity representations model we first review the ier model architecture. much of the material and notation here comes directly from @xcite . let @xmath0, ..., w n ) denote a sequence of input context words, @xmath1, ..., w j ) denote an entity mention span in s (over positions i through j), and @xmath2, 1] |t | denote a vector whose values are predicted probabilities corresponding to fine-grained entity types t from a predefined type system. given a labeled dataset @xmath3, s, t * ) (1) , ..., (m, s, t * ) (k) } the iers' objective is to estimate parameters θ of a function f θ that maps the mention m and its context s to a vector t that captures salient features (fine-grain types) of the entity mention within its context. the entity embedding t whose individual dimensions have explicit semantics can then be used directly as input for downstream tasks using standard similarity measures (e.g., dot products). note that fine-tuning these representations would destroy their interpretability because dimensions would no longer be readable as the probability of the input representing specific entity types. the model f θ that produces these embeddings is depicted as the "encoder" in figure 1 . first, a bert-based encoder @xcite maps inputs m and s to an intermediate dense vector representation. the encoder input is a token sequence @xmath4, where the mention m and context s are segmented into wordpiece tokens @xcite . the vector output [cls] token serves as a d-dimensional dense mention and context representation: h [cls] = bertencoder(x) ∈ r d . the key ingredient of iers is a type embedding layer, which projects this intermediate representation to a vector whose dimensions correspond to the entity types in t using a single linear layer with parameters @xmath5. finally, each dimension (individually) is passed through the sigmoid function, yielding the predicted probabilities that form the interpretable entity representation t (the "intermediate layer" in figure 1 ). more concisely: @xmath6. to estimate parameters we optimize the sum of binary cross-entropy losses entity types t over training examples d. ## intermediate entity-based sparse interpretable representation learning we modify the ier model just described as follows: this proposed model architecture -which we have called itsirl -is depicted in figure 1. during pre-training, we adopt a loss l that combines entity typing loss over the sparse intermediate interpretable layer l et and the reconstruction loss of the output representation l recon in contrast to prior ier work in which sparse entity type representations were used directly for downstream tasks, here we pass the intermediate interpretable representation into a feed forward decoder network that produces a new representation which is used for prediction. this choice leads to differences in interpretability between iers and our proposed architecture. we explore this in section 5, along with how these intermediate predicted entity types affect task performance and how user or automated mechanisms to manipulate (i.e., up or down weight) these intermediate types affects performance. this approach in some ways resembles concept bottleneck models @xcite ; reviewed further in section 6). however, these methods generally use low dimensional, human-labeled concept supervision to guide learning for a single task. by contrast, in our approach we exploit large-scale, possibly noisy entity type supervision to learn to induce interpretable representations which might be useful across tasks, i.e., for general pre-training. we could pre-train such models in a few ways: (i) train them end-to-end, or, (ii) use existing ier models as points of initialization. in the latter case, we freeze the ier model originally trained using only l et and train/update the rest of the model weights using only l recon as the loss on our pretraining data. for our experiments we use the publicly avail-able biomedical ier model checkpoint, entity type system, and pre-training data from @xcite . the model checkpoint is based on an underlying pubmedbert model @xcite ). the type system contains 68,304 entity types and the training data consists of 37,357,141 triples of the form (mention, context, [list of entity types]) derived from pubmed linked wikipedia pages where entity types are wikipedia categories. ## experimental setup we evaluate the proposed itsirl architecture on two biomedical benchmark tasks: entity label classification for cancer genetics @xcite and sentence similarity regression for the biosses dataset found in the blurb benchmark @xcite . ## entity type counterfactual manipulation and global explainability we have claimed that (sparse) entity type representations permit "interpretability", but this is an illdefined term in general. here we demonstrate that itsirl provides a specific type of "interpretability" in that it can help facilitate model understanding and error analysis via "counterfactual" entity type manipulation, made possible by the intermediate entity type layer. specifically, we consider the cancer genetics classification task @xcite , and focus on revealing learned global structure of classes. we then show how manipulating predicted types on erroneous test cases affects the itsirl model's performance. ## related work in this work we introduced an architecture with an encoder that uses supervision from a pre-defined our use of the intermediate interpretable entity layer to represent classes for global explainability is reminiscent of work for learning prototypes for images @xcite , timeseries @xcite or text @xcite , however in our case constructing the prototypes of each class is done post-hoc and as such the prototypes are used for analysis rather than classification or learning. additionally, our method is interpretable at the vector component level whereas the latent representations used for constructing prototypes are not. also, our pre-trained representations are not tied to a classification task like prototypes and as such can be used for various different tasks. our model could be viewed as including an internal probing task which tests a models' ability to induce type information by measuring the accuracy of a probe @xcite @xcite . however, probing is usually a post-hoc means of revealing the information implicitly stored within internal dense output representations, whereas our model was defined and pre-trained in such a way as to explicitly provide intermediate interpretable entity type representations. ## conclusions in this work we proposed intermediate entitybased sparse interpretable representation learning (itsirl), an extension to the iers architecture which provides an intermediate interpretable layer whose decoded dense representation output can be fine-tuned and leveraged for performance on downstream tasks. empirically we show the model substantially outperforms prior iers work on two diverse benchmark biomedical tasks. to demonstrate the utility of the kind of interpretability afforded by itsirl, we proposed a counterfactual entity type manipulation analysis which allows for modeling debugging. this is a finegrained, human interaction inquiry made possible by the proposed model architecture and pretraining scheme. using coarse class type sets, we show this technique can allow itsirl to surpass performance against dense non-interpretable models. this analysis establishes that entity type manipulation works intuitively as expected in itsirl, which is important for future work on methods for flagging when a predicted answer should be inspected and possibly manipulated at the entity type level.