Distilling BlackBox to Interpretable Models for Efficient Transfer Learning

Abstract

Building generalizable AI models is one of the primary challenges in the healthcare domain. While radiologists rely on generalizable descriptive rules of abnormality, Neural Network (NN) models suffer even with a slight shift in input distribution (e.g., scanner type). Fine-tuning a model to transfer knowledge from one domain to another requires a significant amount of labeled data in the target domain. In this paper, we develop an interpretable model that can be efficiently fine-tuned to an unseen target domain with minimal computational cost. We assume the interpretable component of NN to be approximately domain-invariant. However, interpretable models typically underperform compared to their Blackbox (BB) variants. We start with a BB in the source domain and distill it into a mixture of shallow interpretable models using human-understandable concepts. As each interpretable model covers a subset of data, a mixture of interpretable models achieves comparable performance as BB. Further, we use the pseudo-labeling technique from semi-supervised learning (SSL) to learn the concept classifier in the target domain, followed by fine-tuning the interpretable models in the target domain. We evaluate our model using a real-life large-scale chest-X-ray (CXR) classification dataset. The code is available at: https://github.com/batmanlab/MICCAI-2023-Route-interpret-repeat-CXRs.

1. Introduction

Model generalizability is one of the main challenges of AI, especially in high stake applications such as healthcare. While NN models achieve state-of-the-art (SOTA) performance in disease classification [9,17,24], they are brittle to small shifts in the data distribution [7] caused by a change in acquisition protocol or scanner type [22]. Fine-tuning all or some layers of a NN model on the target domain can alleviate this problem [2], but it requires a substantial amount of labeled data and be computationally expensive [12,21]. In contrast, radiologists follow fairly generalizable and comprehensible rules. Specifically, they search for patterns of changes in anatomy to read abnormality from an image and apply logical rules for specific diagnoses. This approach is transparent and closer to an interpretable-by-design approach in AI. We develop a method to extract a mixture of interpretable models based on clinical concepts, similar to radiologists’ rules, from a pre-trained NN. Such a model is more data- and computation-efficient than the original NN for fine-tuning to a new distribution.

Standard interpretable by design method [18] finds an interpretable function (e.g., linear regression or rule-based) between human-interpretable concepts and final output [14]. A concept classifier [19,26] detects the presence or absence of concepts in an image. In medical images, previous research uses TCAV scores [13] to quantify the role of a concept on the final prediction [3,6,23], but the concept-based interpretable models have been mostly unexplored. Recently Posthoc Concept Bottleneck models (PCBMs) [25] identify concepts from the embeddings of BB. However, the common design choice amongst those methods relies on a single interpretable classifier to explain the entire dataset, cannot capture the diverse sample-specific explanations, and performs poorly than their BB variants.

Our Contributions.

This paper proposes a novel data-efficient interpretable method that can be transferred to an unseen domain. Our interpretable model is built upon human-interpretable concepts and can provide sample-specific explanations for diverse disease subtypes and pathological patterns. Beginning with a BB in the source domain, we progressively extract a mixture of interpretable models from BB. Our method includes a set of selectors routing the explainable samples through the interpretable models. The interpretable models provide First-order-logic (FOL) explanations for the samples they cover. The remaining unexplained samples are routed through the residuals until they are covered by a successive interpretable model. We repeat the process until we cover a desired fraction of data. Due to class imbalance in large CXR datasets, early interpretable models tend to cover all samples with disease present while ignoring disease subgroups and pathological heterogeneity. We address this problem by estimating the class-stratified coverage from the total data coverage. We then finetune the interpretable models in the target domain. The target domain lacks concept-level annotation since they are expensive. Hence, we learn a concept detector in the target domain with a pseudo labeling approach [15] and finetune the interpretable models. Our work is the first to apply concept-based methods to CXRs and transfer them between domains.

2. Methodology

Notation.

Assume $f^0:𝒳\to 𝒴$ is a BB, trained on a dataset $𝒳\times 𝒴\times 𝒞$, with $𝒳,𝒴$, and $𝒞$ being the images, classes, and concepts, respectively; $f^0=h^0\circ \Phi$, where $\Phi$ and $h^0$ is the feature extractor and the classifier respectively. Also, $m$ is the number of class labels. This paper focuses on binary classification (having or not having a disease), so $m=2$ and $𝒴\in \left\{\mathrm{0,1}\right\}$. Yet, it can be extended to multiclass problems easily. Given a learnable projection [4,5], $t:\Phi \to 𝒞$, our method learns three functions: (1) a set of selectors $(\pi :𝒞\to \{\mathrm{0,1}\left\}\right)$ routing samples to an interpretable model or residual, (2) a set of interpretable models $(g:𝒞\to 𝒴)$, and (3) the residuals. The interpretable models are called “experts” since they specialize in a distinct subset of data defined by that iteration’s coverage $\tau$ as shown in SelectiveNet [16]. Figure 1 illustrates our method.

Fig. 1.

Schematic view of our method. Note that $f^k(.)=h^k\left(\Phi \right(.\left)\right)$. At iteration $k$, the selector routes each sample either towards the expert $g^k$ with probability ${\pi }^k(.)$ or the residual $r^k=f^{k-1}-g^k$ with probability $1-{\pi }^k(.).g^k$ generates FOL-based explanations for the samples it covers. Note $\Phi$ is fixed across iterations.

2.1. Distilling BB to the Mixture of Interpretable Models

Handling Class Imbalance.

For an iteration $k$, we first split the given coverage ${\tau }^k$ to stratified coverages per class as $\left\{{\tau }_m^k=w_m\cdot {\tau }^k;w_m=N_m/N,\forall m\right\}$, where $w_m$ denotes the fraction of samples belonging to the $m^{th}$ class; $N_m$ and $N$ are the samples of $m^{\mathit{\text{th}}}$ class and total samples, respectively.

Learning the Selectors.

At iteration $k$, the selector ${\pi }^k$ routes $i^{th}$ sample to the expert $\left(g^k\right)$ or residual $\left(r^k\right)$ with probability ${\pi }^k\left({\mathit{c}}_i\right)$ and $1-{\pi }^k\left({\mathit{c}}_i\right)$ respectively. For coverages $\left\{{\tau }_m^k,\forall m\right\}$, we learn $g^k$ and ${\pi }^k$ jointly by solving the loss:

$${\theta }_{s^k}^{*},{\theta }_{g^k}^{*}=\underset{{\theta }_{s^k},{\theta }_{g^k}}{\text{arg min}}{ℛ}^k\left({\pi }^k\left(.;{\theta }_{s^k}\right),g^k\left(.;{\theta }_{g^k}\right)\right)\hspace{1em}\text{s.t.}{\zeta }_m\left({\pi }^k\left(.;{\theta }_{s^k}\right)\right)\ge {\tau }_m^k\forall m,$$ (1)

where ${\theta }_{s^k}^{*},{\theta }_{g^k}^{*}$ are the optimal parameters for ${\pi }^k$ and $g^k$, respectively. ${ℛ}^k$ is the overall selective risk, defined as, ${ℛ}^k\left({\pi }^k,g^k\right)=\sum _m\frac{\frac{1}{N_m}{\sum }_{i=1}^{N_m}{ℒ}_{\left(g^k,{\pi }^k\right)}^k\left({\mathit{x}}_i,{\mathit{c}}_i\right)}{{\zeta }_m\left({\pi }^k\right)}$, where ${\zeta }_m\left({\pi }^k\right)=\frac{1}{N_m}{\sum }_{i=1}^{N_m}{\pi }^k\left({\mathit{c}}_{\mathit{i}}\right)$ is the empirical mean of samples of $m^{th}$ class selected by the selector for the associated expert $g^k$. We define ${ℒ}_{\left(g^k,{\pi }^k\right)}^k$ in the next section. The selectors are neural networks with sigmoid activation. At inference time, ${\pi }^k$ routes a sample to $g^k$ if and only if ${\pi }^k(.)\ge .0.5$.

Learning the Experts.

For iteration $k$, the loss ${ℒ}_{\left(g^k,{\pi }^k\right)}^k$ distills the expert $g^k$ from $f^{k-1},$BB of the previous iteration by solving the following loss:

$${ℒ}_{\left(g^k,{\pi }^k\right)}^k\left({\mathit{x}}_i,{\mathit{c}}_i\right)=\underset{\begin{array}{c}\text{trainable component}\\ \text{for current iteration }k\end{array}}{\underbrace{\ell \left(f^{k-1}\left({\mathit{x}}_i\right),g^k\left({\mathit{c}}_i\right)\right){\pi }^k\left(c_i\right)}}\underset{\begin{array}{c}\text{fixed component trained}\\ \text{in the previous iterations}\end{array}}{\underbrace{\prod _{j=1}^{k-1}\left(1-{\pi }^j\left({\mathit{c}}_i\right)\right)}},$$ (2)

where ${\pi }^k\left({\mathit{c}}_i\right){\prod }_{j=1}^{k-1}\left(1-{\pi }^j\left({\mathit{c}}_i\right)\right)$ is the cumulative probability of the sample covered by the residuals for all the previous iterations from $1,\cdots ,k-1$ (i.e., $\left.{\prod }_{j=1}^{k-1}\left(1-{\pi }^j\left({\mathit{c}}_i\right)\right)\right)$ and the expert $g^k$ at iteration $k$ (i.e., ${\pi }^k{(\mathit{c}}_{\mathit{i}})$).

Learning the Residuals.

After learning $g^k$, we calculate the residual as, $r^k\left(x_i,c_i\right)=f^{k-1}\left(x_i\right)-g^k\left(c_i\right)$ (difference of logits). We fix $\Phi$ and optimize the following loss to update $h^k$ to specialize on those samples not covered by $g^k$, effectively creating a new BB $f^k$ for the next iteration $(k+1)$:

$${ℒ}_f^k\left({\mathit{x}}_j,{\mathit{c}}_j\right)=\underset{\begin{array}{c}\text{trainable component}\\ \text{for iteration }k\end{array}}{\underbrace{\ell \left(r^k\left({\mathit{x}}_j,{\mathit{c}}_j\right),f^k\left({\mathit{x}}_j\right)\right)}}\underset{\begin{array}{c}\text{non-trainable component}\\ \text{for iteration }k\end{array}}{\underbrace{\prod _{i=1}^k\left(1-{\pi }^i\left({\mathit{c}}_j\right)\right)}}$$ (3)

We refer to all the experts as the Mixture of Interpretable Experts (MoIE-CXR). We denote the models, including the final residual, as MoIE-CXR+R. Each expert in MoIE-CXR constructs sample-specific FOLs using the optimization strategy and algorithm discussed in [4].

2.2. Finetuning to an Unseen Domain

We assume the MoIE-CXR-identified concepts to be generalizable to an unseen domain. So, we learn the projection $t_t$ for the target domain and compute the pseudo concepts using SSL [15]. Next, we transfer the selectors, experts, and final residual (${\left\{{\pi }_s^k,g_s^k\right\}}_{k=1}^K$ and $f_s^K$) from the source to a target domain with limited labeled data and computational cost. Algorithm 1 details the procedure.

Algorithm 1.

Finetuning to an unseen domain.

1: Input: Learned selectors, experts, and final residual from source domain: ${\left\{{\pi }_s^k,g_s^k\right\}}_{k=1}^K$ and $f_s^K$ respectively, with $K$ as the number of experts to transfer. BB of the source domain: $f_s^0=h_s^0\left({\Phi }_s\right)$. Source data: ${𝒟}_s=\left\{{𝒳}_s,{𝒞}_s,{𝒴}_s\right\}$. Target data: ${𝒟}_t=\left\{{𝒳}_t,{𝒴}_t\right\}$. Target coverages ${\left\{{\tau }_k\right\}}_{k=1}^K$.

2: Output: Experts ${\left\{{\pi }_t^k,g_t^k\right\}}_{k=1}^K$ and final residual $f_t^K$ of the target domain.

3: Randomly select $n_t\ll N_t$ samples out of $N_t=\left|{𝒟}_t\right|$.

4: Compute the pseudo concepts for the correctly classified samples in the target domain using $f_s^0$, as, $c_t^i=t_s\left({\Phi }_s\left(x_s^i\right)\right)$ s.t., $y_t^i=f_s^0\left(x_t^i\right),i=1\cdots n_t$

5: Learn the projection function $t_t$ for target domain semi-supervisedly [15] using the pseudo labeled samples ${\left\{x_t^i,c_t^i\right\}}_{i=1}^{n_t}$ and unlabeled samples ${\left\{x_t^i\right\}}_{i=1}^{N_t-n_t}$.

6: Complete the triplet for the target domain $\left\{{𝒳}_t,{𝒞}_t,{𝒴}_t\right\}$, where $c_t^i=t_t\left({\Phi }_s\left(x_t^i\right)\right),i=1\cdots N_t$.

7: Finetune ${\left\{{\pi }_s^k,g_s^k\right\}}_{k=1}^K$ and $f_s^K$ to obtain ${\left\{{\pi }_t^k,g_t^k\right\}}_{k=1}^K$ and $f_t^K$ using equations 1, 2 and 3 respectively for 5 epochs. ${\left\{{\pi }_t^k,g_t^k\right\}}_{k=1}^K$ and $\left\{{\left\{{\pi }_t^k,g_t^k\right\}}_{k=1}^K,f_t^K\right\}$ represents MoIE-CXR and MoIE-CXR + R for the target domain.

3. Experiments

We perform experiments to show that MoIE-CXR 1) captures a diverse set of concepts, 2) does not compromise BB’s performance, 3) covers “harder” instances with the residuals in later iterations resulting in their drop in performance, 4) is finetuned well to an unseen domain with minimal computation.

Experimental Details.

We evaluate our method using 220,763 frontal images from the MIMIC-CXR dataset [11]. We use Densenet121 [8] as BB $\left(f^0\right)$ to classify cardiomegaly, effusion, edema, pneumonia, and pneumothorax, considering each to be a separate binary classification problem. We obtain 107 anatomical and observation concepts from the RadGraph’s inference dataset [10], automatically generated by DYGIE++ [20]. We train BB following [24]. To retrieve the concepts, we utilize until the $4^{\mathit{\text{th}}}$ Densenet block as feature extractor $\Phi$ and flatten the features to learn $t$. We use an 80%−10%−10% train-validation-test split with no patient shared across splits. We use 4, 4, 5, 5, and 5 experts for cardiomegaly, pneumonia, effusion, pneumothorax, and edema. We employ ELL [1] as $g$. Further, we only include concepts as input to $g$ if their validation auroc exceeds 0.7. Refer to Table 1 in the supplementary material for the hyperparameters. We stop until all the experts cover at least 90% of the data cumulatively.

Table 1.

MoIE-CXR does not compromize the performance of BB. We provide the mean and standard errors of AUROC over five random seeds. For MoIE-CXR, we also report the percentage of test set samples covered by all experts as “Coverage”. We boldfaced our results and BB.

Model	Effusion	Cardiomegaly	Edema	Pneumonia	Pneumothorax
Blackbox (BB)	0.92	0.84	0.89	0.79	0.91
INTERPRETABLE BY DESIGN
CEM [26]	0.83±1e−4	0.75±1e−4	0.77±2e−4	0.62±4e−4	0.76±3e−4
CBM (Sequential) [14]	0.78±1e−4	0.72±1e−4	0.77±5e−4	0.60±1e−3	0.75±6e−4
CBM + ELL [1,14]	0.81±1e−4	0.72±1e−4	0.79±5e−4	0.62±8e−4	0.75±6e−4
POSTHOC
PCBM [25]	0.88±1e−4	0.81±1e−4	0.82±1e−4	0.72±1e−4	0.85±7e−4
PCBM-h [25]	0.90±1e−4	0.83±1e−4	0.85±1e−4	0.77±1e−4	0.89±7e−4
PCBM + ELL [1,25]	0.90±1e−4	0.82±1e−4	0.85±1e−4	0.75±1e−4	0.85±6e−4
PCBM-h + ELL [1,25]	0.91±1e−4	0.83±1e−4	0.87±1e−4	0.77±1e−4	0.90±1e−4
OURS
MoIE-CXR (Coverage)	${0.93}_{\pm 1\mathbf{e}-4}^{\mathit{\left(}\mathit{0.90}\mathit{\right)}}$	${0.85}_{\pm 1\mathbf{e}-4}^{\mathit{\left(}\mathit{0.96}\mathit{\right)}}$	${0.91}_{\pm 1\mathbf{e}-4}^{\mathit{\left(}\mathit{0.92}\mathit{\right)}}$	${0.80}_{\pm 1\mathbf{e}-4}^{\mathit{\left(}\mathit{0.97}\mathit{\right)}}$	${0.91}_{\pm 2\mathbf{e}-4}^{\mathit{\left(}\mathit{0.93}\mathit{\right)}}$
MoIE-CXR+R	0.91 ±1e−4	0.82 ±1e−4	0.88 ±1e−4	0.78 ±1e−4	0.90 ±2e−4

Baseline.

We compare our method with 1) end-to-end CEM [26], 2) sequential CBM [14], and 3) PCBM [25] baselines, comprising of two parts: a) concept predictor $\Phi :𝒳\to 𝒞$, predicting concepts from images, with all the convolution blocks; and b) label predictor, $g:𝒞\to 𝒴$, predicting labels from the concepts. We create CBM + ELL and PCBM + ELL by replacing the standard classifier with the identical $g$ of MOIE-CXR to generate FOLs [1] for the baseline.

MoIE-CXR Captures Diverse Explanations.

Figure 2 illustrates the FOL explanations. Recall that the experts $\left(g\right)$ in MoIE-CXR and the baselines are ELLs [1], attributing attention weights to each concept. A concept with high attention weight indicates its high predictive significance. With a single $g$, the baselines rank the concepts in accordance with the identical order of attention weights for all the samples in a class, yielding a generic FOL for that class. In Fig. 2, the baseline PCBM + ELL uses left_pleural and pleural_unspec to identify effusion for all four samples. MoIE-CXR deploys multiple experts, learning to specialize in distinct subsets of a class. So different interpretable models in MoIE assign different attention weights to capture instance-specific concepts unique to each subset. In Fig. 2 expert2 relies on right_pleural and pleural_unspec, but expert4 relies only on pleural_unspec to classify effusion. The results show that the learned experts can provide more precise explanations at the subject level using the concepts, increasing confidence and trust in clinical use.

Fig. 2.

Qualitative comparison of MoIE-CXR discovered concepts with the baselines.

MoIE-CXR does not Compromise BB’s Performance. Analysing MoIE-CXR:

Table 1 shows that MoIE-CXR outperforms other models, including BB. Recall that MoIE-CXR refers to the mixture of all interpretable experts, excluding any residuals. As MoIE-CXR specializes in various subsets of data, it effectively discovers sample-specific classifying concepts and achieves superior performance. In general, MoIE-CXR exceeds the interpretable-by-design baselines (CEM, CBM, and CBM + ELL) by a fair margin (on average, at least ~ 10% ↑), especially for pneumonia and pneumothorax where the number of samples with the disease is significantly less (~750/24000 in the testset).

Analysing MoIE-CXR+R:

To compare the performance on the entire dataset, we additionally report MoIE-CXR+R, the mixture of interpretable experts with the final residual in Table 1. MoIE-CXR+R outperforms the interpretable-by-design models and yields comparable performance as BB. The residualized PCBM baseline, i.e., PCBM-h, performs similarly to MoIE-CXR+R. PCBM-h rectifies the interpretable PCBM’s mistakes by learning the residual with the complete dataset to resemble BB’s performance. However, the experts and the final residual approximate the interpretable and uninterpretable fractions of BB, respectively. In each iteration, the residual focuses on the samples not covered by the respective expert to create BB for the next iteration and likewise. As a result, the final residual in MoIE-CXR+R covers the “hardest” examples, reducing its overall performance relative to MoIE-CXR.

Identification of Harder Samples by Successive Residuals.

Figure 3 (a–c) reports the proportional AUROC of the experts and the residuals per iteration. The proportional AUROC is the AUROC of that model times the empirical coverage, ${\zeta }^k$, the mean of the samples routed to the model by the respective selector $\left({\pi }^k\right)$. According to Fig. 3a in iteration 1, the residual (black bar) contributes more to the proportional AUROC than the expert1 (blue bar) for effusion with both achieving a cumulative proportional AUROC ~ 0.92. All the final experts collectively extract the entire interpretable component from BB $f^0$ in the final iteration, resulting in their more significant contribution to the cumulative performance. In subsequent iterations, the proportional AUROC decreases as the experts are distilled from the BB of the previous iteration. The BB is derived from the residual that performs progressively worse with each iteration. The residual of the final iteration covers the “hardest” samples. Tracing these samples back to the original BB $f^0,f^0$ underperforms on these samples (Fig. 3 (d–f)) as the residual.

Fig. 3.

Performance of experts and residuals across iterations. (a-c): Coverage and proportional AUROC of the experts and residuals. (d-f): Routing the samples covered by MoIE-CXR to the initial $f^0$, we compare the performance of the residuals with $f^0$.

Applying MoIE-CXR to the Unseen Domain.

In this experiment, we utilize Algorithm 1 to transfer MoIE-CXR trained on MIMIC-CXR dataset to Stanford Chexpert [9] dataset for the diseases – effusion, cardiomegaly and edema. Using 2.5%, 5%, 7.5%, 10%, and 15 % of training data from the Stanford Chexpert dataset, we employ two variants of MoIE-CXR where we (1) train only the selectors $\left(\pi \right)$ without finetuning the experts $\left(g\right)$ (“No finetuned” variant of MoIE-CXR in Fig. 4), and (2) finetune $\pi$ and $g$ jointly for only 5 epochs (“Finetuned” variant of MoIE-CXR and MoIE-CXR + R in Fig. 4). Finetuning $\pi$ is essential to route the samples of the target domain to the appropriate expert. As later experts cover the “harder” samples of MIMIC-CXR, we only transfer the experts of the first three iterations (refer to Fig. 3). To ensure a fair comparison, we finetune (both the feature extractor $\Phi$ and classifier $h^0$) BB: $f^0=h^0\circ \Phi$ of MIMIC-CXR with the same training data of Stanford Chexpert for 5 epochs. Throughout this experiment, we fix $\Phi$ while finetuning the final residual in MoIE+R as stated in Eq. 3. Figure 4 displays the performances of different models and the computation costs in terms of Flops. The Flops are calculated as, Flop of (forward propagation + backward propagation) × (total no. of batches) × (no of training epochs). The finetuned MoIE-CXR outperforms the finetuned BB (on average ~ 5% ↑ for effusion and cardiomegaly). As experts are simple models [1] and accept only low dimensional concept vectors compared to BB, the computational cost to train MoIE-CXR is significantly lower than that of BB (Fig. 4 (d–f)). Specifically, BB requires ~ 776T flops to be finetuned on 2.5% of the training data of Stanford CheXpert, whereas MoIE-CXR requires ~ 0.0065T flops. As MoIE-CXR discovers the sample-specific domain-invariant concepts, it achieves such high performance with low computational cost than BB.

Fig. 4.

Transferring the first 3 experts of MoIE-CXR trained on MIMIC-CXR to Stanford-CXR. With varying % of training samples of Stanford CXR, (a-c): reports AUROC of the test sets, (d-g) reports computation costs in terms of log (Flops) (T). We report the coverages in Stanford-CXR on top of the “finetuned” and “No finetuned” variants of MoIE-CXR (red and blue bars) in (d-g).

4. Conclusion

This paper proposes a novel iterative interpretable method that identifies instance-specific concepts without losing the performance of the BB and is effectively fine-tuned in an unseen target domain with no concept annotation, limited labeled data, and minimal computation cost. Also, as in the prior work, MoIE-captured concepts may not showcase a causal effect that can be explored in the future.