Diagnosing Transformers: Illuminating Feature Spaces for Clinical Decision-Making

TL;DR

Diagnosing Transformers presents SUFO, a pipeline that fuses supervised probing, unsupervised similarity analysis, feature dynamics, and outlier analysis to interpret fine-tuned transformer features in clinical NLP. The authors compare general-domain, mixed-domain, and domain-specific pre-trained models on pathology reports and MedNLI, showing that domain-specific models may overfit minority classes while mixed-domain models offer robustness; in-domain pre-training speeds feature disambiguation but the final performance depends on data closeness and diversity. A key finding is that fine-tuning induces substantial sparsification in the feature space, enabling effective outlier detection and expert validation of failure modes. Overall, SUFO provides actionable insights for model selection, evaluation, and deployment in medicine and other high-stakes domains, supporting safer and more trustable use of transformers in practice.

Abstract

Pre-trained transformers are often fine-tuned to aid clinical decision-making using limited clinical notes. Model interpretability is crucial, especially in high-stakes domains like medicine, to establish trust and ensure safety, which requires human engagement. We introduce SUFO, a systematic framework that enhances interpretability of fine-tuned transformer feature spaces. SUFO utilizes a range of analytic and visualization techniques, including Supervised probing, Unsupervised similarity analysis, Feature dynamics, and Outlier analysis to address key questions about model trust and interpretability. We conduct a case study investigating the impact of pre-training data where we focus on real-world pathology classification tasks, and validate our findings on MedNLI. We evaluate five 110M-sized pre-trained transformer models, categorized into general-domain (BERT, TNLR), mixed-domain (BioBERT, Clinical BioBERT), and domain-specific (PubMedBERT) groups. Our SUFO analyses reveal that: (1) while PubMedBERT, the domain-specific model, contains valuable information for fine-tuning, it can overfit to minority classes when class imbalances exist. In contrast, mixed-domain models exhibit greater resistance to overfitting, suggesting potential improvements in domain-specific model robustness; (2) in-domain pre-training accelerates feature disambiguation during fine-tuning; and (3) feature spaces undergo significant sparsification during this process, enabling clinicians to identify common outlier modes among fine-tuned models as demonstrated in this paper. These findings showcase the utility of SUFO in enhancing trust and safety when using transformers in medicine, and we believe SUFO can aid practitioners in evaluating fine-tuned language models for other applications in medicine and in more critical domains.

Figures (14)

• Figure 1: Layer-wise RSA comparing the pre-trained and fine-tuned versions of the models across four pathology classification tasks.
• Figure 2: Illustration of clustering algorithm: fine-tuned BERT on Path-PG. (a) The projection onto the first two PCs. (b) A similar projection for Path-SV. Notice that the scales of the two PCs are drastically different and a naive clustering based on the Euclidean metric may not capture the variation in PC2. (c) The $3$ clusters obtained solely from the projection onto PC1 with each red bar denoting the boundaries of a single cluster. (d) $3$ clusters similarly obtained on PC2. (e) The final set of $3$ clusters obtained by forming all possible combinations of clusters from (c) and (d) and selecting the three largest.
• Figure A1: The first two PCs in the fine-tuned last layer classification token feature spaces of all the models explain on average 95% of the dataset variance across the 4 tasks.
• Figure A2: Filling back in the first two PCs, at the last two steps, $k = 767$ and $k = 768$, yields significant model performance gain.
• Figure A3: Path-PG: BERT