B-Cos Aligned Transformers Learn Human-Interpretable Features

TL;DR

Vision and Swin Transformers achieve strong accuracy in computational pathology but lack intrinsic interpretability for clinical use. The paper proposes $B$-cos transformers (BvT and Bwin) that replace all linear transforms with the $B$-cos transformation to enforce weight–input alignment, enabling more interpretable biomedically meaningful features; the method includes two $B$-cos* branches with MaxOut and omits ReLUs, with the final $B$-cos defined as $ \text{B-cos}(x;w) = \max_{i \in \{1,2\}} \text{B-cos*}(x;w_i)$ and $ \text{B-cos*}(x;w) = \|\hat{w}\| \|x\| |c(x,\hat{w})|^B \text{sgn}(c(x,\hat{w}))$ (with $\|\hat{w}\|=1$). The study demonstrates this leads to more uniform representations via CK A, and domain experts rate BvT/Bwin as more trustworthy, with competitive F1 gains on three public datasets and transfer learning boosts. The work suggests that inherently interpretable transformer architectures can bridge the gap between high performance and trustworthiness in pathology AI, with potential for broader adoption in medical imaging.

Abstract

Vision Transformers (ViTs) and Swin Transformers (Swin) are currently state-of-the-art in computational pathology. However, domain experts are still reluctant to use these models due to their lack of interpretability. This is not surprising, as critical decisions need to be transparent and understandable. The most common approach to understanding transformers is to visualize their attention. However, attention maps of ViTs are often fragmented, leading to unsatisfactory explanations. Here, we introduce a novel architecture called the B-cos Vision Transformer (BvT) that is designed to be more interpretable. It replaces all linear transformations with the B-cos transform to promote weight-input alignment. In a blinded study, medical experts clearly ranked BvTs above ViTs, suggesting that our network is better at capturing biomedically relevant structures. This is also true for the B-cos Swin Transformer (Bwin). Compared to the Swin Transformer, it even improves the F1-score by up to 4.7% on two public datasets.

Figures (6)

• Figure 1: Attention maps of ViT and BvT (ours) on the test set of (a) NCT-CRC-HE-100K, (b) TCGA-COAD-20X, and (c) Munich-AML-Morphology. BvT attends to various diagnostically relevant features such as cancer tissue, cells, and nuclei.
• Figure 2: The model architecture of ViT and BvT (ours). We replace all linear transformations in ViT with the B-cos transform and remove all ReLU activation functions.
• Figure 3: Rollout, Attention-Last (Attn-Last), Grad-CAM, LRP, LRP of the second layer (LRP-Second), LRP of the last layer (LRP-Last), and Transformer Attribution (TA) applied on the test set of Munich-AML-Morphology. The image shows an eosinophil, which is characterized by its split, but connected nucleus, large specific granules (pink structures in the cytoplasm), and dense chromatin (dark spots inside the nuclei) Rosenberg_2013_NatRevImm. Across all visualization techniques, BvT focuses on these exact features unlike ViT.
• Figure 4: We compute the central kernel alignment (CKA), which measures the representation similarity between each hidden layer. Since the B-cos transform aligns the weights with the inputs, BvT (ours) achieves a more uniform representation structure compared to ViT (values closer to 1). When trained with the binary cross-entropy loss (BCE) instead of the categorical cross-entropy loss (CCE), the alignment is higher.
• Figure 5: In a blinded study, domain experts ranked models (lower is better) based on whether the models focus on biomedically relevant features that are known in the literature to be important for diagnosis. We then performed the Conover post-hoc test after Friedman with adjusted p-values according to the two-stage Benjamini-Hochberg procedure. BvT ranks above ViT with $p$$<$$0.1$ (underlined) and $p$$<$$0.05$ (bold).