WiSDoM pipeline.

Fig 1 illustrates the WiSDoM general pipeline, applicable to both fully-supervised and weakly-supervised settings. The process begins with initialization, where sample-label pairs are encoded to set the parameters of the joint KDM. This initialization uses encoded patch-label pairs from the training dataset to form the initial KDM .

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. WiSDoM architecture.

A. Initialization: The process begins with an initialization step where pairs of samples and labels are encoded to set the parameters of the joint KDM. B. Encoding: For the fully-supervised patch Gleason classifier, individual patches are fed directly, bypassing attention weighting since each patch has a weight of 1. In the weakly-supervised WiSDoM, patch bags are extracted from the WSI, encoded into a feature space by a CNN, and processed through an attention network that aggregates local and global information. These feature vectors and attention weights are represented as a density matrix, modeling the information as an input probability distribution. C. Inference/Training: From the joint KDM of weighted prototypes and their labels, an output distribution of labels is derived, providing a whole slide-level label posterior distribution with an expected value and variance. During training, all parameters in the encoder, attention network, and joint KDM are updated via gradient descent. D. Prototyping: When a new sample is input into WiSDoM, it is first encoded into latent space by the encoder. This encoded feature vector is then compared with all prototypes learned in the joint KDM using Euclidean distance. The closest prototype is used to retrieve the corresponding example and label from the initialization set.

https://doi.org/10.1371/journal.pone.0335826.g001

WiSDoM employs a deep neural network as a feature extractor, transforming input patches into 128-dimensional feature vector representations . In the fully-supervised setting, such as Gleason pattern grading, individual patches are fed directly into the model. The patch feature vectors are represented as a KDM with m = 1 components and .

For weakly-supervised tasks, WiSDoM incorporates an attention mechanism. Patch bags extracted from the WSI are processed through an attention network that aggregates local and global information, assigning weights to different patches. The feature vectors, along with attention weights in the weakly-supervised case, are represented as a density matrix, modeling the information as an input probability distribution. This forms the joint KDM of weighted prototypes and their labels. From this joint KDM, WiSDoM derives an output distribution of labels, providing a whole slide-level label posterior distribution with an expected value and variance.

The joint KDM is initialized by an arbitrary set of encoded patch-label pairs from the training dataset D. During training, all parameters in the encoder, attention network (for weakly-supervised tasks), and joint KDM are updated via gradient descent.

WiSDoM predicts a probability distribution over ordinal labels, providing expected value and variance for each prediction. The learned internal parameters of the joint KDM allow prototypes to be obtained as examples of the internal learned representations.

The inference procedure involves using the KDM and input KDM , and performing an inference operation (see Eq (8)). The resulting KDM contains a discrete probability distribution of output labels , where each , represents the probability associated with the i-th label. When performing a classification task, we select the most probable label from the distribution. When performing ordinal regression, we slightly modify the inference procedure: First, we convert categorical labels to continuous labels in the range [0,1], the conversion operation from a categorical label to an ordinal label is simply achieved by normalizing the categorical label of each sample by the total number of possible labels as follows (Eq (9)).

(9)

From the probability distribution obtained from we can calculate the expected value and variance.

Given a density matrix , represented by a vector , where each represents the probability associated with the i-th label, the expected value and variance are computed as follows:

(10)

Where are the values associated with each label.

The variance is calculated as the expected value of the squares minus the square of the expected value. This is given by:

(11)

where is calculated similarly to but using the square of the values ():

(12)

The expected value and variance are then output for each input patch. Algorithms 1 and 2 summarize the training and prediction procedure of fully-supervised WiSDoM.

Algorithm 1 Fully-supervised WiSDoM training algorithm.

Input:

: Training dataset.

m: number of components (encoded set of training patches) of KDM

Z: Deep learning backbone

1. KDM is initialized with a sample of size m from dataset D

2. for each : (see Algorithm 2)

3. If task = classification: i.

ii. Minimize

4. If task = ordinal regression:

i. Calculate and , where

ii. Minimize , where α is a penalization parameter for variance.

5. Update all backbone weights w, and KDM parameters using gradient descent.

6. Return

Algorithm 2 Fully-supervised WiSDoM prediction procedure.

Input:

: joint KDM

x: input patch

Z: Deep Learning backbone

1. Encode patch x using Z:

2. Create

3. Calculate probabilities for output KDM using and with Eq (8)

4.

5. Return

Weakly-supervised tasks.

Our prior approach relied on tissue annotations to select relevant patches and to gather Gleason pattern labels for patches across an entire slide. This strategy facilitated an interpretable method for quantifying the extent of each Gleason pattern on the whole slide, similar to how a pathologist would conduct their diagnosis, all centered around WiSDoM classifying patches into specific Gleason patterns.

However, given the high cost and relative unavailability of tissue annotation masks in real-world clinical scenarios, we aim to eliminate the necessity for tissue annotations during training. In this section, we propose a novel method that requires only a whole-slide diagnosis, typically an ISUP grade group for prostate biopsies, which can be easily obtained from pathology reports.

Competitive performance can be achieved in line with current state-of-the-art methodologies for whole-slide grading, which only needs a weak label for training while being constrained by providing interpretability.

We extend WiSDoM probabilistic deep learning framework for weakly supervised, interpretable ordinal regression and classification. It operates on the principle that each WSI in the training set is an individual data point with an established slide-level diagnosis yet lacks specific pixel or region-level annotations. The framework adopts a similar approach to MIL, viewing each WSI as a collection of numerous smaller segments or patches (see Fig 1).

Traditionally, MIL focuses on binary classification, discerning positive from negative classes under the assumption that the presence of one positive patch classifies the entire slide as positive. This approach typically employs a max-pooling aggregation function, choosing the patch with the highest probability of the positive class for slide-level classification. However, this method is unsuitable for multiclass or binary classifications without explicit positive/negative annotations.

WiSDoM differentiates itself by not using the standard max-pooling or other conventional aggregation functions like average pooling, generalized mean, or log-sum-exp, which are limited in terms of problem-specific adaptability and interpretability. Instead, WiSDoM integrates an attention-guided KDM for aggregating information from patches. This method allows for a more nuanced integration of patch-level data into a unified WSI prediction or representation, offering enhanced interpretability and adaptability for various classification problems.

Following tissue detection and patch extraction, WiSDoM involves encoding the N patches constituting a WSI into a feature vector representation utilizing a deep learning backbone.

In our study on ISUP grading, we adopt a novel approach by using a collection of sample instances, known as ‘bags,’ instead of labeling each sample individually. This method is particularly suited to scenarios where patches from a whole slide collectively form a specific ISUP grade group, but individual Gleason patterns at the patch level remain unknown, a common occurrence in real-world settings.

We interpret a specific collection of patches from a WSI as a ‘bag.’ The challenge for the model is to learn to assign accurate labels to each patch within these bags and then synthesize this information to make a comprehensive prediction at the whole-slide level. This approach inherently involves uncertainties, especially regarding the individual characteristics of each patch within a bag. Our objective is to model these uncertainties effectively. This integration allows for a more accurate and reliable prediction process, closely mirroring the complexities encountered in actual pathological assessments.

During the training process, WiSDoM takes bags of training samples . The training dataset corresponds to a set of pairs , where each is a vector expressing the label proportions of the i-th sample. Each input sample is represented by a KDM with m_i components. For our specific problem, where the goal is to obtain a whole-slide ISUP grade group from a patch bag, we model a variant of the original implementation of KDM [5]. It receives training sample bags as a set of pairs , where each is the ISUP grade group of the bag, perceived as the whole-slide ’weak’ label.

Furthermore, considering the density matrix representation inherent to the KDM, which ascribes a probability to each possible label, we can model the significance or contribution of each instance within a bag towards the overall bag’s label. To accomplish this, we employ a local-global attention method, as shown in [28]. This method assigns a weight, or a contribution factor, to each instance within the bag. Its application to natural images has proven to be useful, as it not only enhances performance but is also able to pinpoint regions of interest (ROIs), providing an additional layer of interpretability. This contribution of each patch to the bag class can be modeled into the probability p_i of each KDM component . By incorporating this additional information, we enhance the weakly-supervised learning process by compelling the model to assign greater importance to certain instances within the bags over others.

This attention module receives patches from a bag and processes it using two multi-layer perceptrons (MLPs), which form the means to extract attention weights from these patch bags. The initial MLP is charged with computing the local context, which essentially encapsulates the local information available in each patch . This is accomplished by passing the input through the first MLP, defined as MLP₁, which yields .

(13)

Subsequently, a global context is obtained by aggregating the local context across all patches.

(14)

Where k is the number of patches in the bag. This step provides understanding of the information present in the input data and forms the basis for the subsequent attention distribution. The local () and global () information are then combined, and this representation of both local and global information, are local-global embeddings that are fed to the second MLP MLP₂, yielding another set of weights, z which are the importance of each patch in the bag. The raw attention weights, z, are then passed through a Softmax operation.

(15)

The final attention weights are for each patch in the bag. This Softmax operation normalizes these weights such that they all lie between 0 and 1 and their total sum equals 1. The application of this operation allows the model to weigh each patch based on both the unique contribution of each patch and the global context, enhancing the model’s performance by considering both individual and collective factors.

The primary differentiation in this approach, in comparison to the previous experiment, resides in the KDM creation process. Instead of uniformly distributing weights by assigning to every patch in the bag, where m_i is the total number of patches, this novel approach utilizes the attention mechanism to determine these weights. This inclusion allows for a more informative weight assignment that takes into account both individual patch contributions and their collective influence: we assign in . Each MLP is configured with 64 neurons and uses a ReLU activation function. This configuration, with the number of neurons being half of the input’s dimension, is chosen based on the feature vector size of 128 neurons, effectively reducing the input dimensionality by half, balancing between model complexity and computational efficiency, ensuring that the model is capable of learning a rich set of features without being prohibitively expensive to train on top of the encoder and KDM parameters. The training is conducted in an end-to-end manner, optimizing the parameters across all components of the model. This includes the patch encoder, the global-local attention mechanism, and the KDM .

Additionally, the trained local-global attention layer of our model can be extended to provide qualitative interpretability of the decisions, not only providing a way to visualize the most important patches in the patch bag but effectively showing the most significant patches in the slide for accurately prediction its ISUP grade group.

The core of the slide-level classification task is the inference process using the KDM and input KDM in the same fashion as the fully-supervised case using Eq (8).

The density matrix is then translated into a discrete probability distribution over the classes. A vector of probabilities is computed from the components of , where the weights and vectors are denoted by and , respectively. Both are normalized, and , and the probability distribution is obtained as . This probability vector represents the likelihood of the WSI belonging to each class, forming the basis for the slide-level classification or ordinal regression task. For the ordinal regression task, we add a final regression layer that takes the probability distribution of labels as input. This layer computes the expected value and variance for predictions. Algorithms 3 and 4 show a summary of the training and prediction procedure of WiSDoM in a weakly-supervised setting.

Algorithm 3 Weakly-supervised WiSDoM training algorithm.

Input:

: Training dataset, with a WSI with k patches

m: number of components of KDM

Z: Deep learning backbone

1. KDM is initialized with a sample of size m from dataset D

2. for each : (see Algorithm 4)

3. If task = classification:

i.

ii. Minimize

4. If task = ordinal regression:

i. Calculate and , where

ii. Minimize , where α is a penalization parameter for variance.

5. Update backbone, MLP₁ and MLP₂ weights w, and KDM parameters using gradient descent.

6. Return

Algorithm 4 Weakly-supervised WiSDoM prediction procedure.

Input:

: input WSI with k patches

: joint KDM

Z: Deep Learning backbone

1.

2.

3.

4. Encode patches using Z:

5. Create

6. Calculate probabilities from output KDM using and with Eq (8)

7.

8. Return

Prototype initialization and selection.

For KDM initialization, the joint KDM requires m = 216 components, with 36 prototypes for each of the 6 ISUP grade groups (0-5). We randomly select 36 samples from each class through stratified sampling from the training dataset.

During initialization, each selected prototype is processed through the pre-warmed ConvNeXT encoder to obtain 128-dimensional feature representations. These encoded representations initialize the component set of the joint KDM, where represents the encoded prototype. The corresponding labels are one-hot encoded representations for classification tasks or normalized continuous values for ordinal regression tasks. The prototype mixture weights are initialized uniformly as for all components.

During training, the components of undergo refinement through gradient-based optimization. The prototype positions in the 128-dimensional latent space evolve to maximize class separability according to the RBF kernel . The mixture weights are learned simultaneously, with the optimization process determining which prototypes contribute most to the inference operations defined in Eq 8. Post-training analysis shows that the model utilizes a subset of the initialized prototypes. Components with learned weights below p_i<0.01 contribute minimally to the final predictions, indicating that the KDM framework selects the most discriminative prototypes for each diagnostic category.

Computational complexity analysis.

We select m = 216 prototypes to balance representational capacity with computational efficiency. With 36 prototypes per ISUP grade group (6 groups total), this configuration provides diversity within each diagnostic category while maintaining tractable inference complexity. Table 1 shows the computational impact of different prototype configurations.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Table 1. Computational analysis of prototype count impact.

https://doi.org/10.1371/journal.pone.0335826.t001

The computational complexity of KDM inference scales as where n is the number of input patches and m is the prototype count, following Eq (8). The KDM parameters ( for prototype positions plus for label encodings) constitute 0.10% of the total model parameters with m = 216 prototypes, while the encoder backbone represents 99.90% of model complexity. Prototype scaling has minimal impact on model size but linear impact on inference operations through the kernel evaluations required for the inference procedure in Eq (8).

Optimization and training procedure.

After initializing the KDM , the deep learning backbone, attention module, and KDM parameters are trained end-to-end. We use the Adam [30] optimizer with a learning rate of , , and . A gradual warm-up scheduler with a factor of 10 is applied for 1 epoch, followed by cosine annealing for the remaining epochs. The mini-batch size is set to 4 bags. For the loss function, we use categorical cross-entropy for the classification task. The model is trained for 50 epochs with an early-stopping callback to prevent overfitting, stopping training after 5 epochs without validation loss improvement.

For the ordinal regression task, the warm-up and KDM initialization follow the same procedure. However, we use real-valued labels normalized to the range [0,1] during training instead of one-hot encoded labels, following Eq (9). The loss function is modified to Mean Squared Error with an additional penalization α for high variance predictions as shown in Algorithm 3. The same Adam optimizer settings are employed, along with a gradual warm-up scheduler and a mini-batch size of 4 bags.