Explainable Feature Embeddings from Histopathology Foundation Models: A Case Study for End Stage Kidney Disease Risk Analysis in Diabetic Nephropathy Patients

Abstract

Foundational models (FMs) based on advanced neural network architectures have demonstrated improved performance in pathology image analysis across various organs due to their increased generalizability. However, their clinical adoption requires explainability, as their black box nature limits transparency. Understanding the specific features these models learn for a given downstream task is crucial for explainability and integrating FMs into clinical workflows more effectively.

We propose a computational pipeline that enhances explainability by correlating domain-specific handcrafted features (HFs), with hidden features i.e., feature embeddings (FEs) from FMs. We correlate and combine HFs from Detectron 2 DeepLabv3+ segmentation with FEs from Prov-Gigapath (PG) and UNI FMs for improved explainability and performance. In this work, HFs are extracted from segmented functional tissue units, including arteries, tubules, globally sclerotic glomeruli, and non-globally sclerotic glomeruli. FEs are extracted at the tile and slide levels for PG and at the tile level for UNI. We use the Pearson correlation coefficient to identify significant correspondences between these feature sets.

To evaluate our proposed methodology, we use 56 diabetic nephropathy kidney biopsy whole slide images (WSIs) from Seoul National University Hospital. The task is to predict end-stage kidney disease (ESKD) two years post-biopsy using leave-one-out cross-validation on 56 WSIs, with 16 from ESKD patients and 40 from non-ESKD patients. We combine top correlated features from FEs of FMs with HFs and train logistic regression (LR) and $k$ nearest neighbor ($k\text{NN}$) classifiers. LR model trained on combined feature set improved accuracy, balanced accuracy, Matthew’s correlation coefficient, F1-score, precision, and recall to 0.8393, 0.7938, 0.5993, 0.8377, 0.8367, 0.8393 respectively, when compared to LR and $k\text{NN}$ models trained on individual feature sets. PG excelled in specificity (1.000) and AUROC (0.8281), while UNI showed superior AUPRC (0.7813) performance. We also present feature explainability maps corresponding to each feature in FE.

1. INTRODUCTION

Foundation models (FMs) are large-scale artificial intelligence (AI) models trained using self-supervised learning (SSL) on a diverse range of data. They represent a significant advancement in AI development, with applications spanning various fields such as radiology,¹ genomics,² coding,³ and mathematics.⁴ In the domain of digital pathology, several FMs have been developed, including Prov-Gigapath (PG),⁵ UNI,⁶ Virchow,⁷ RudolfV,⁸ and Hibou.⁹ These models have been trained on hundreds of thousands of slides from various data sources and organs, enabling them to produce hidden features i.e., feature embeddings (FEs) which are generalizable.

Previous studies have used handcrafted features (HFs),¹⁰ neural networks,¹¹ and a combination of neural networks and HFs¹² to perform predictions. Providing a physical interpretation to FEs is crucial for helping clinicians to use them in a clinical setting. Popular techniques for visualizing and understanding the decisions made by deep learning models, particularly convolutional neural networks, include class activation mapping^13–15 and attention-based methods.^{16, 17} These techniques enhance clinicians trust in FMs and facilitate their use in clinical settings. However, interpreting the FEs from FMs in digital pathology remains largely unexplored. Our contributions in this paper are summarized as follows:

• Propose a computational pipeline for explainability of FEs of FMs by correlating them with HFs.

• Propose a methodology to combine FEs and HFs to enhance accuracy.

• Apply the proposed methodology to build a computer-assisted diagnosis model that can predict ESKD.

Our research work enables better integration of FMs into clinical settings.

2. METHODS

We predict end-stage kidney disease (ESKD) two years post- biopsy using leave-one-out cross-validation. We used data from 56 participants with type 2 diabetes who were recruited at the Seoul National University (SNU) Hospital to evaluate our methodology. The biopsy WSIs were stained histologically using periodic acid-Schiff (PAS). Of the 56 WSIs, 16 were from patients who developed ESKD, and 40 from those who did not. Data collection followed protocols approved by the Institutional Review Board at the SNU College of Medicine (H-1812-159-998), Seoul, Korea, and all experiments complied with federal guidelines and regulations.

In this work, we develop a methodology for predicting the pathological meaning of the FEs produced by FMs and also propose a computational module to combine useful features from FEs of various FMs with handcrafted features to improve performance on the downstream tasks. Fig. 1 shows the computational pipeline developed in this research work. Our methodology comprises of the following steps:

Fig. 1:

A schematic diagram that illustrates the proposed computational pipeline for obtaining explainable features from histopathology FMs. HFs and FEs are extracted from the segmentation pipeline and the FMs respectively. Thereafter, a feature ranking module is used to rank the FEs using HFs via correlation analysis. Then we fuse the top features in FEs with HFs for the downstream analysis.

Preparation and Preprocessing:

First, we utilize Otsu thresholding technique¹⁸ to segment the WSIs into foreground and background, divide them into non-overlapping tiles, and filter out those with more than 90% background. These tiles are passed through FMs. UNI and PG FMs take a tile of size 224 × 224 and 256 × 256 respectively.

Foundation Model Embeddings:

Thereafter, FEs are derived from the pre-trained FMs PG and UNI, which provide robust representations of the WSIs. For the PG model, tile-level embeddings (TLE) are obtained from the last_layer_embed layer, with each tile embedding having dimensions of 1536 × 1. These embeddings are averaged to obtain a single representation of a WSI at the tile level. Additionally, we obtain slide-level embeddings (SLE), using slide-level encoder based on the LongNet architecture, which generates contextualized 768 × 1 embeddings from TLE. For the UNI model, we can only obtain TLE as there is no slide level encoder. Each 1024 × 1 TLE from UNI is averaged to obtain a single representation of a WSI.

Tissue Segmentation and Histomorphometry Features:

We utilize a multi-compartment segmentation algorithm, as detailed in our previous work,¹⁹ to segment different functional tissue units (FTUs) including arteries, renal tubules, globally sclerotic glomeruli, and non-globally sclerotic glomeruli in the PAS stained WSIs. Our approach employs the DeepLabv3 model,²⁰ pre-trained on a large dataset of renal tissue biopsy WSIs, to achieve precise segmentation of FTUs. The multi-compartment segmentation algorithm is based on a panoptic model which provides a robust framework for segmenting various FTUs providing generally high performance, enabling effective feature extraction.²¹ From each segmented FTU, we extract a set of 72 HFs, encompassing various morphological, textural, and intensity-based characteristics as discussed in our prior work.²² A total of 288 features derived from the four different FTUs provide a comprehensive dataset for downstream analysis. Fig. 2 shows the physical interpretation of the two of the handcrafted features.

Fig. 2:

Illustrative figure for histomorphometry of selected handcrafted features

(a) Max. distance transform represents the maximum of the distances between a PAS + ve pixel and the closest non PAS + ve pixel

(b) Luminal Space Area is the space occupied by the lumen inside an FTU, and PAS + ve area is the area occupied by the purple region inside an FTU

Feature Selection:

Let ${\mathbf{E}}_{\text{PG}}^{\text{slide}},{\mathbf{E}}_{\text{PG}}^{\text{tile}}$, and ${\mathbf{E}}_{\text{UNI}}$ denote the set of embeddings obtained from the PG FM at the slide level, tile level, and from the UNI FM, respectively. Let $\mathbf{H}$ represent the set of HFs.

Correlation Calculation:

FEs obtained from the FMs and the HFs are informative for predicting the target outcome. For each feature in the FEs, we calculate the Pearson correlation with the HFs using the scipy.stats library. We calculate the correlation matrices as follows:

$${\mathbf{C}}_{\text{PG}}^{\text{slide}}=\text{corr}\left({\mathbf{E}}_{\text{PG}}^{\text{slide}},\mathbf{H}\right),$$ (1)

$${\mathbf{C}}_{\text{PG}}^{\text{tile}}=\text{corr}\left({\mathbf{E}}_{\text{PG}}^{\text{tile}},\mathbf{H}\right),$$ (2)

$${\mathbf{C}}_{\text{UNI}}=\text{corr}\left({\mathbf{E}}_{\text{UNI}},\mathbf{H}\right)$$ (3)

where corr(⋅,⋅) denotes the correlation.

Sparsification of Correlation Matrices:

To ensure that only the most relevant features are selected, we apply a threshold of $\left|\rho \right|\ge 0.5$ and $p\le 0.05$ for selecting the most effective correlation coefficients statistically in the correlation matrix, thus sparsifying the correlation matrices and enhancing the robustness of the feature selection process. Let ${\mathbf{C}}_{\text{PG}}^{\text{slide,sparse}},{\mathbf{C}}_{\text{PG}}^{\text{tile,sparse}}$, and ${\mathbf{C}}_{\text{UNI}}^{\text{sparse}}$ denote the sparsified correlation matrices obtained from ${\mathbf{C}}_{\text{PG}}^{\text{slide}},{\mathbf{C}}_{\text{PG}}^{\text{tile}}$, and ${\mathbf{C}}_{\text{UNI}}$ correlation matrices respectively.

Top $N$ Feature Extraction:

We assume that the top $N$ features selected from the correlation matrices capture the most significant relationships between the embeddings and the HFs. Extract the top $N$ features from each sparsified correlation matrix. We took a value of $N=288$, since these are the number of HFs present.

$${\mathbf{F}}_{\text{PG}}^{\text{slide}}={\text{top}}_N\left({\mathbf{C}}_{\text{PG}}^{\text{slide},\text{sparse}}\right),$$ (4)

$${\mathbf{F}}_{\text{PG}}^{\text{tile}}={\text{top}}_N\left({\mathbf{C}}_{\text{PG}}^{\text{tile},\text{sparse}}\right),$$ (5)

$${\mathbf{F}}_{\text{UNI}}={\text{top}}_N\left({\mathbf{C}}_{\text{UNI}}^{\text{sparse}}\right)$$ (6)

Concatenation of Top Features:

The selected top $N$ features from each FE are concatenated with handcrafted features to form a feature set of size $4N$.

$${\mathbf{F}}_{\text{combined}}={\mathbf{F}}_{\text{PG}}^{\text{slide}}\cup {\mathbf{F}}_{\text{PG}}^{\text{tile}}\cup {\mathbf{F}}_{\text{UNI}}\cup {\mathbf{F}}_{\text{handcrafted}}$$ (7)

${\mathbf{F}}_{\text{combined}}$ is used for the downstream tasks.

3. EXPERIMENTS AND RESULTS:

Our model aims to predict the development of ESKD two years post-biopsy, relying on features from FMs and HFs. Linear and $k\text{NN}$ probing was done separately on an individual set of FE. On the combined set of features and HFs we train logistic regression and $k\text{NN}$ classifers. The models were evaluated using leave-one-out cross-validation (LOO-CV). For training the $k\text{NN}$ classifiers, a grid search was conducted to determine the optimal number of neighbors. PG’s TLE performed better compared to it’s SLE, UNI TLE, and HFs if we consider an individual set of features. Overall, results show that combining features from different FEs (PG, UNI) and HFs improves performance across various metrics such as accuracy, balanced accuracy, Matthew’s correlation coefficient (MCC), F1 score, precision, and recall, indicating the benefit of combining multiple feature representations for the prediction task. Table 1 shows the results obtained on various FEs, HFs and combined set of features. Fig. 3a shows the interpretation of the top FEs indices clustered based on the feature type, showing that major features types learned are based on color and texture. Fig. 3b shows the interpretation of the top FEs indices clustered based on the FTUs, suggesting that features are mainly learned form arteries and tubules.

Table 1:

Classification metrics for ESKD prediction in DN

	Prov-Gigapath				UNI		H. F.		C. F.
Metrics	SLE-LP	SLE-KP	TLE-LP	TLE-KP	TLE-LP	TLE-KP	LR	$k\text{NN}$	LR	$k\text{NN}$
Accuracy	0.8035	0.7321	0.8214	0.7679	0.8035	0.7500	0.6428	0.7321	0.8393	0.7321
B.A.	0.7322	0.5313	0.7438	0.5938	0.7313	0.6375	0.5438	0.6063	0.7938	0.5500
MCC	0.4948	0.2132	0.5367	0.3762	0.4948	0.3244	0.0912	0.2614	0.5993	0.2006
F1 score	0.7969	0.6351	0.8129	0.7047	0.7969	0.7299	0.6353	0.7054	0.8377	0.6593
Precision	0.7958	0.8052	0.8149	0.8248	0.7958	0.7304	0.6295	0.7058	0.8367	0.7161
Recall	0.8035	0.7321	0.8214	0.7679	0.8035	0.7500	0.6428	0.7321	0.8393	0.7321
Specificity	0.9000	1.0000	0.9250	1.0000	0.9000	0.9000	0.7750	0.9000	0.9000	0.9750
AUROC	0.8219	0.5625	0.8281	0.7008	0.8313	0.6375	0.6781	0.6063	0.8266	0.6844
AUPRC	0.6136	0.4405	0.7654	0.5987	0.7813	0.5768	0.3793	0.5322	0.6913	0.5327

Fig. 3:

Top 20 features from FM correlated with the top 20 handcrafted features

(a) Maximum of absolute correlations of top 20 features clustered based on feature type for PG TLE

(b) Maximum of absolute correlation of top 20 features clustered based on functional tissue unit for PG TLE

Table 1 shows the classification results obtained after doing linear and $k\text{NN}$ probing on the FEs obtained from the PG and UNI FMs. B.A. refers to Balanced Accuracy. SLE, TLE, represent the slide level embeddings, tile level embeddings respectively. LP and KP refers to linear probing and $k\text{NN}$ probing respectively. H.F., C.F. represent the handcrafted features and combined features respectively.

4. CONCLUSIONS

We demonstrate that we can interpret the FEs from FMs with HFs using correlation analysis. We also present the potential of combining FEs from FMs, specifically PG and UNI, with HFs to predict end stage kidney disease (ESKD) in subjects with DN. First, we perform linear and $k\text{NN}$ probing on the FEs from FMs and also train a logistic regression and $k\text{NN}$ classifier on the HFs. Next, we combine the most relevant features from the FEs with HFs and train logistic regression and $k\text{NN}$ model using these combined feature set. Furthermore, we present a set of FEs ranked by their level of explainability and significance in ESKD risk prediction. Our findings suggest that the integration of HFs with FEs can enhance the interpretability and performance of predictive models in computational pathology. The proposed method provides a framework for utilizing FMs in clinical settings, offering a balance between predictive accuracy and explainability. Future work will focus on validating our approach on larger datasets and exploring the application of this method to other pathological conditions.