Attention-based Interpretability with Concept Transformers

1 INTRODUCTION . The spectacular gains in accuracy of recent large-scale machine learning models like deep neural networks have generally come at the cost of a loss of transparency into their functioning . This “ black box ” aspect severely limits their applicability in safety-critical domains , such as medical diagnostics , healthcare , public infrastructure safety , visual inspection for civil engineering , to name just a few , where it is essential for decisions to be corroborated by robust domain-relevant knowledge . In recent years , approaches focusing on explaining black box models have emerged , mostly with the goal of providing post-hoc explanations in terms of a set of relevant features used by the underlying model to make predictions ( Ribeiro et al. , 2016 ; Selvaraju et al. , 2017 ; Lundberg & Lee , 2017 ; Smilkov et al. , 2017 ) . While widely used , such post-hoc explainability methods have been criticized for operating on low-level features such as pixel values , or sensory signals that are combined in unintelligible ways and do not correspond to high-level concepts that humans easily understand ( Kim et al. , 2018 ; Alvarez-Melis & Jaakkola , 2018 ; Kindermans et al. , 2019 ; Su et al. , 2019 ) . To overcome these limitations and sidestep the potential perils resulting from a misuse of posthoc explainability of black box models , some researchers have been vocally advocating for the use of inherently interpretable models ( Rudin , 2019 ) that in particular would generate decisions based on human-understandable categories ( i.e. , concepts ) grounded in domain expertise rather than raw features ( Barbiero et al. , 2021 ; Ghorbani et al. , 2019 ; Kim et al. , 2018 ; Yeh et al. , 2020 ; Koh et al. , 2020 ; Goyal et al. , 2019 ; Kazhdan et al. , 2020 ; Chen et al. , 2020 ; Alvarez-Melis & Jaakkola , 2018 ; Li et al. , 2018 ; Chen et al. , 2019 ) . For example , to identify a bird species , a model should focus on morphologically meaningful concepts , such as the shape , size and colors of beak , feathers or wings , rather than focusing on raw pixels , and combine them in ways that a domain expert ( in this case an ornithologist ) would reckon as intelligible to produce a classification . In addition , using high-level concepts emulates a human ’ s thinking process ( i.e. , structured into familiar concepts ) and provides insights into the model ’ s reasoning in a human-understandable way . The chasm between post-hoc explainability vs. inherently interpretable models closely reflects a related ongoing discussion in the NLP community on the interpretation of attention mechanisms ( Bahdanau et al. , 2014 ) , and in particular on the interpretability of attention weights over input tokens , with researchers on one end of the debate claiming that attention provides interpretability , while others claim that “ Attention is not explanation ” ( Jain & Wallace , 2019 ) . While the debate over what degree of interpretability that can be ascribed to attention weights is still not settled ( Wiegreffe & Pinter , 2019 ) , it is arguable that in many situations attention is not a “ fail-safe indicator ” ( Serrano & Smith , 2019 ) , particularly when decisions rely on the interaction of multiple interacting tokens as is typically the case in deep architectures . Conversely then , a way to guarantee the interpretability of attention weights would be to make sure that they are not being processed by downstream operations that renders their relation to the decision outputs uninterpretable . This is indeed something that had been proposed in the past , in particular in architectures that preserve the interpretability of “ relevance scores ” ( akin to attention weights ) by acting on them only through a restricted class of intelligible “ aggregation functions ” such as additive models ( Alvarez-Melis & Jaakkola , 2018 ) , which are a common functional elements in interpretable and white-box models ( Caruana et al. , 2015 ) . In this paper , we propose the ConceptTransformer ( CT ) , a transformer-based module ( Vaswani et al. , 2017 ) for classification tasks , that can be used to enhance an arbitrary deep learning classifier with domain knowledge in the form of plausible cross-attention weights between input features and highlevel interpretable concepts . The CT can be used as a drop-in replacement for the classifier head of any deep learning architecture . The resulting model can then be trained end-to-end without any additional overhead on the training pipeline , except a modification of the loss function that enforces plausibility of the explanation . Importantly , the CT was specifically conceived to provide explanations that guarantee faithfulness by design and plausibility by construction . Faithfulness is defined as the degree to which the explanation reflects the decision and aims therefore to ensure that the explanations are indeed explaining the model ’ s operation ( Lakkaraju et al. , 2019 ; Guidotti et al. , 2018 ) . In our model this is achieved by enforcing a linear relation between the transformer value vectors that represent the concepts and their contribution to the classification log-probabilities . Plausibility refers to how convincing the interpretation is to humans ( Guidotti et al. , 2018 ; Carvalho et al. , 2019 ) . In the CT architecture , plausibility is achieved by construction by supervising the attention heads of the cross-attention mechanism to conform with inputs-concepts-outputs relations derived by domain knowledge . We validate our approach on three image benchmark datasets , MNIST Even/Odd ( Barbiero et al. , 2021 ) , CUB-200-2011 ( Welinder et al. , 2010 ) , and aPY ( Farhadi et al. , 2009 ) . On these datasets we will examine how the faithfulness and plausibility of the CT explanations are practically translated into domain-relevant explanations behind particular output decisions or diagnostic insights about ensuing wrong classifications . We will also quantify the benefit of domain-expert knowledge in terms of statistical efficiency by showing that providing domain-relevant explanations to our CT model tends to improve the performance of the downstream classification , in particular in low-data regime . This for instance translates in an 8-9 % improvement in accuracy on the bird classification CUB-200-2011 dataset when CT is trained in conjunction with part location annotations . We note in addition that one of the strengths of our CT model is its versatility that allows it to be effortlessly applied to other data modalities as well by combining it with deep learning classifiers which are then rendered interpretable with no appreciable overhead or change in their training pipeline . This is in stark contrast to other inherently interpretable models that are often specifically designed for the domain at hand , and require adhoc multi-stage training procedures . CT on the other hand , is differentiable and can be flexibly included in any end-to-end training pipeline that uses backpropagation , as we will showcase by combining it with a host of different deep learning backbones ranging from convolutional architectures like Residual Networks ( He et al. , 2016 ) to more modern Vision Transformer ( Dosovitskiy et al. , 2020 ) and hybrid Compact Convolutional Transformer models ( Hassani et al. , 2021 ) . 2 RELATED WORK . In recent years , there have been significant advancements towards designing interpretable models that quantify the importance of individual features with respect to the prediction output . One general approach is post-hoc analysis , in which one interprets a trained model by fitting explanations to the classification outputs ( Alvarez-Melis & Jaakkola , 2018 ; Ribeiro et al. , 2016 ; Lundberg & Lee , 2017 ) . In particular for CNNs , popular techniques are activation maximization ( van den Oord et al. , 2016 ; Nguyen et al. , 2016 ; Yosinski et al. , 2015 ) and saliency visualization ( Selvaraju et al. , 2017 ; Smilkov et al. , 2017 ; Sundararajan et al. , 2017 ) . However , these post-hoc methods do not actually explain how the underlying model reached a particular classification outcome . In contrast , attention-based interpretable techniques aim to expose which parts of a given input a network focuses on , and therefore deems to be most relevant , when making a decision . Examples of attention models are Zhang et al . ( 2014 ) ; Zhou et al . ( 2016 ; 2018 ) ; Zheng et al . ( 2017 ) ; Fu et al . ( 2017 ) ; Girshick ( 2015 ) ; Girshick et al . ( 2014 ) ; Huang et al . ( 2016 ) . The problem with these models is that they focus on low-level individual features when providing an explanation . Such features are often not intuitive for humans , are typically noisy and non-robust , or can be misleading when interpreted afterwards ( Kim et al. , 2018 ) . One of the recent advancements in the field of interpretability was to design methods that explain predictions with high-level human understandable concepts ( Ghorbani et al. , 2019 ; Kim et al. , 2018 ; Yeh et al. , 2020 ; Koh et al. , 2020 ; Goyal et al. , 2019 ; Kazhdan et al. , 2020 ; Chen et al. , 2020 ; Barbiero et al. , 2021 ; Li et al. , 2018 ; Chen et al. , 2019 ) – either by identifying common activation patterns in the last nodes of the neural network corresponding to human understandable categories or constraining the network to learn such concepts . For instance , TCAV ( Kim et al. , 2018 ) proposes to define concepts from user-annotated examples in which concepts appear . Others propose prototypesbased explanation models ( Li et al. , 2018 ; Chen et al. , 2019 ) , but they typically require specialized convolutional architectures to ensure feature extraction . In particular , ProtoPNet ( Li et al. , 2018 ) uses previously learned prototypes to focus attention on various parts of an image . This architectural design implies that object-level ( global ) concepts can not be easily incorporated , and since prototypes are not learned together with the attention model , explanations based on these prototypes may lack faithfulness . SENN ( Alvarez-Melis & Jaakkola , 2018 ) proposes a network that transforms inputs into interpretable basic features , generates concept relevance scores and then aggregates concepts with relevance scores to explain predictions . While it is out-of-the-box interpretable , it lacks concept localization . Barbiero et al . ( 2021 ) proposed a differentiable approach that allows the extraction of logic explanations from neural network models using First-Order Logic . The approach relies on an entropy-based criterion to automatically identify the most relevant concepts that have contributed to a particular classification output . In our approach , high-level concepts are defined with a set of related dimensions , and can be partspecific or global . Such concepts are typically readily available in many domains and can be used to enhance the performance of the learning task while offering explainability at no additional cost for the network . Obtained explanations are plausible and guaranteed to be faithful , since concepts participate in the model computation . Finally , CT also allows in some cases to discover the presence concepts that were not annotated . 3 APPROACH . The ConceptTransformer module . The ConceptTransformer ( CT ) is a transformer-based module designed to be used as classifier head in a deep learning architecture that generates classification outputs using cross-attention ( Vaswani et al. , 2017 ) between input features and a set of embeddings representing domain-relevant concepts . Fig . 1 shows the case where the inputs to the CT are embeddings of P visual patches of an input image that are linearly projected and concatenated into the query matrix Q ∈ RP×dm of a query-key-value cross-attention mechanism whose corresponding key matrix K ∈ RC×dm is the linearly projected concatenation of the embeddings representing the C concepts . In addition , the concepts are linearly projected with a value projection matrix and concatenated to result in the value matrix V ∈ RC×dm . Cross-attention then outputs an attention weight αpc = softmax ( 1√ dm QK > ) pc with p = 1 , . . . , P , c = 1 , . . . , C , between each patch-concept pair , which are combined into an attention map matrix A = [ αpc ] ∈ RP×C . The final output of the CT is the product obtained by multiplying the attention map A , the value matrix V and an output matrix O ∈ Rdm×nc that projects onto the ( unnormalized ) nc logits over the output classes , and averaging over patches : logiti = 1 P P∑ p=1 [ AV O ] pi with i = 1 , . . . , nc . ( 1 ) Notice that here for simplicity we described a single-head attention model , but in our experiments we will be using a multi-head version ( Vaswani et al. , 2017 ) . Equation 1 says that , given an input x to the network , the conditional probability of output class i is Pr ( i|x ) = softmaxi ( ∑C c=1 βc γc ( x ) ) with βc with components ( βc ) i = [ V O ] ci , ( 2 ) and γc ( x ) are positive relevance scores that depend on x through the averaged attention weights : γc ( x ) = 1 P ∑P p=1 αpc . The output of the CT is therefore essentially a simple multinomial logistic regression model over positive variables γc ( x ) that measures the contribution of each concept . Notice that this result follows from the linear relation between the value vectors and the classification logits , which itself comes from the design choices of computing outputs from the value matrix V through the linear projection V O , and aggregating patch contributions by averaging . Faithful concept-based explanations by design . The CT was conceived as a drop-in replacement for the classifier head of an arbitrary deep learning classifier to provide concept-based explanations of the outputs that are guaranteed to be faithful by design . We formalize this statement as follows : Proposition 1 Each concept relevance score γc ( x ) in Equation 2 is a faithful explanation of the output . More specifically , the probability of choosing the preferred output ic = argmaxi ( βc ) i of concept c ( assuming it ’ s unique ) is guaranteed to decrease if γc ( x ) is externally set to zero . Moreover , the correlation between γc ( x ) and Pr ( ic|x ) is strictly positive . Proof Proof of Proposition 1 is provided in Appendix A . Note that the last statement in the Proposition above is a corollary of the first one , and it specifically shows that CT is guaranteed to satisfy the technical definitions of faithfulness given for instance by Alvarez-Melis & Jaakkola ( 2018 ) . Training and plausibility by construction . As mentioned , the CT is a differentiable transformerbased module that can be embedded in a deep learning architecture trained end-to-end with backpropagation . In addition , the fact that it exposes attention weights over concept tokens that can be user defined gives us the freedom to shape these attention weights according domain-expert knowledge relevant for the problem under consideration . This can be done by explicitly guiding the attention heads to attend to concepts in the input that are a priori known to be informative to correctly classify the input . In practice this can be achieved by supervising the attention weights at training as for instance proposed by Deshpande & Narasimhan ( 2020 ) as a self-supervised technique for bidirectional language models . In particular , given a desired distribution of attention H provided by domain knowledge ( e.g. , we know which patches in the input image contain which concepts that are relevant to classify the input ) we can force the CT attention weights A by adding an “ explanation cost ” term to the loss function that is proportional to Lexpl = ||A − H||2F , where || · ||F is the Frobenius norm . The final loss used to train the architecture then becomes L = Lcls + λ Lexpl , where Lcls denotes the original classification loss , Lexpl the additional explanation loss , and the constant λ ≥ 0 controls the relative contribution of the explanation loss to the total loss . Notice that setting λ = 0 essentially amounts to just minimizing the classification loss and disregarding the prior domain knowledge as imparted into CT by guiding the attention heads .

Summary: The paper proposes a new version of transformer called concept transformer with the aim of improving the interpretability of its attention by computing cross attention between the inputs features and a set of concepts. This should make the model both more explainable, in the sense that it is easier for the human to interpret the attention weights of meaningful concepts, and also more faithful, in the sense that the attention scores given to particular concepts directly impact the final prediction of the model.  Update: Following the author response I'd like to keep my positive score reflecting my view of the paper.