Rationalizing Transformer Predictions via End-To-End Differentiable Self-Training (2024)

Abstract: We propose an end-to-end differentiable training paradigm for stable training of a rationalized transformer classifier. Our approach results in a single model that simultaneously classifies a sample and scores input tokens based on their relevance to the classification. To this end, we build on the widely-used three-player-game for training rationalized models, which typically relies on training a rationale selector, a classifier and a complement classifier. We simplify this approach by making a single model fulfill all three roles, leading to a more efficient training paradigm that is not susceptible to the common training instabilities that plague existing approaches. Further, we extend this paradigm to produce class-wise rationales while incorporating recent advances in parameterizing and regularizing the resulting rationales, thus leading to substantially improved and state-of-the-art alignment with human annotations without any explicit supervision.

## introduction neural networks are increasingly prevalent across a wide range of applications, driving significant advancements in fields such as natural language processing, computer vision, and beyond. due to the black-box nature of these networks, this widespread use comes with an increased demand for interpretability @xcite , as understanding the basis for the decisions made by these models is crucial for their reliable and ethical deployment. this need has become especially clear with the increasing use of notoriously uninterpretable large language models, which have the potential to quickly lose a user's trust after only few confidently incorrect predictions @xcite . one possible mitigation is the use of encoderonly models, which lend themselves more readily to classical interpretability approaches designed for general neural network classifiers while still providing state-of-the-art performance due to continuous improvements in model structure @xcite and training paradigms @xcite . while a variety of explainability methods exist, that usually assign scores to input tokens indicating their importance for a classification @xcite , these methods often suffer from several drawbacks, including high computational cost, difficult-to-interpret explanations, and potentially even unfaithful representations of the model's decision-making process. in this study, we close this gap by developing a rationalized transformer predictor that generates faithful and interpretable explanations in addition to its decisions within the same forward pass. as a foundation for our approach, we build upon the existing and commonly used three-player game proposed by @xcite . in this framework, a selector model chooses a subset of the input as rationale, while a predictor and a complement predictor model are trained to infer the correct label from either the tokens included in the rationale or the tokens not included in the rationale, respectively. the selector model is then trained to maximally aid the predictor in predicting the correct label while preventing the complement predictor from doing the same, thus ensuring that all tokens indicative of the correct label are included in the rationale. while the general three-player game is sensible, the actual realizations that are proposed often have several limitations, including being not end-to-end differentiable due to a stochastic sampling process in the forward pass, showing interlocking dynamics that might prevent convergence to a suitable solution, and having no guarantee of providing a rationale that actually explains the prediction (compare section 2.3 for a more detailed discussion). for this reason, we propose a new take on this three-player game that is not susceptible to these drawbacks. we achieve this by making use of a single unified model that is trained as a standard classifier on the complete unaltered input, while simultaneously predicting class-wise importance scores for each input token in the same forward pass, which are then trained using self-training to mark spans that the model itself considers important for the specific class. our proposed rationalized transformer predictor (rtp) simplifies and enhances the common three-player structure in several ways, including 1) using only a single model to fulfill all three roles of the three-player game, thus enabling classification and rationale prediction in a single forward pass 2) training the rationales to explain the predictor, but avoiding training the predictor on the rationales, which ensures that the rationales faithfully explain the predictions 3) creating rationalized inputs in continuous fashion to enable fully differentiable training and avoid sampling 4) creating class-wise rationales and 5) using a parameterization that maximizes similarity with human rationale annotations. we evaluate our method on two benchmarks for explainable ai and compare it with existing post-hoc explanation methods as well as methods leveraging standard multi-player procedures. we show that our method achieves state-of-the-art performance on both tasks, demonstrating previously unseen alignment with human rationales in combination with high rationale faithfulness. ## method we propose a new method for end-to-end differentiable training of a rationalized transformer predictor (rtp). in the following, plain letters (e.g., x) denote scalars, while bold letters (e.g., x) denote vectors or tensors. we assume a text classification problem with label set y, and a training set consisting of texts x 0 , ..., x n with corresponding ground truth vectors y 0 , ..., y n . ## experiments we evaluate our method with regards to matching human evidence annotations for text classifications, as well as with regards to the faithfulness of the explanations with regards to the classifier. for details regarding the model and the training and prediction procedures, see appendix a. ## results the results for the evaluation on the inas dataset and the movie review dataset are displayed in ## discussion the rtp model demonstrated strong performance across all evaluated metrics. comparing it specifically to the marc method, it outperformed it in every metric related to measuring agreement with human annotations and most faithfulness metrics. this is notable since the rtp can be seen as a neural network parameterized version of the marc approach, which originally optimized mask parameters for each sample individually instead of training a neural network to directly predict them from the input. another well-performing method, especially with regards to token-level evaluation, is the l2e-marc method. the l2e framework @xcite trains a neural network on pre-calculated rationales created by a post-hoc explainer. even though it only saw rationales produced by the marc method, it manages to outperform it on all metrics measuring token-level agreement with human rationales. these two results indicate, that training to explain many different samples leads to better generalization, which we attribute to reduced overfitting to one specific input. this effect is crucial for the rtp, since it performs input optimization with respect to specific neural network outputs, which has been shown to generally lead to unexpected and uninterpretable artifacts @xcite . the marc method successfully mitigated this issue by combining con-strained optimization with heavy regularization, but artifacts (i.e., unexpected spans included in the rationale) are still to be expected. in the case of the rtp, training on many samples further reduces this issue, since these unwanted gradient signals will generally not match between different samples, so that the neural network mainly adapts to the wanted signal that is consistent within larger parts of the training set, and that correspond to features that are generally indicative of the respective class. ## conclusion we presented a new method for training a rationalized transformer predictor and demonstrated its strong performance on two natural language processing benchmarks. since our proposed training scheme is not invasive to the general training process and does not produce significant overhead during prediction, we believe that this approach has the potential to facilitate wider adoption and availability of rationalized predictors. given that transformers are widely used in other modalities like images @xcite and audio data @xcite , we hypothesize that our approach can be extended to these modalities and potentially lead to results of similar quality.