Interventional Bag Multi-Instance Learning On Whole-Slide Pathological Images

MIL formulation. Due to the gigapixel resolution and lacking fine-grained labels, performing downstream tasks on WSIs is formulated as an MIL problem, such that each WSI is treated as a labeled bag with corresponding patches as unlabeled instances. Take binary classification as an example, let $X=\{(x_{1},y_{1}),...,(x_{n},y_{n})\}$ be a WSI bag, which contains $n$ instances of $x_{i}$. The instance-level labels $\{y_{i},...,y_{n}\}$ are unavailable. Under the standard MIL assumption, the bag label $Y$ is further given by:

$$Y=\begin{cases}0,&\text{ iff }\sum y_{i}=0\\ 1,&\text{ otherwise }\end{cases}$$

(1)

which can be modelled by max-pooling [15]. A general three-stage approach goes like 1) Instance transformation: a feature extractor $f(\cdot)$ is trained for instance-wise features $b$, 2) Instance combination: the pooling operation $\sigma(\cdot)$ is targeted for bag feature $B$, 3) Bag transformation: a downstream classifier $g(\cdot)$ is used for prediction, which can be formulated as:

$$b_{i}=f(x_{i}),B=\sigma(b_{1},\cdots,b_{n}),\hat{Y}=g(B),$$

(2)

where the pooling $\sigma(\cdot)$ should be a permutation-invariant function [17] for the spatial-invariant MIL method. Some works further absorb the classifier $g(\cdot)$ into the pooling operation $\sigma(\cdot)$, referred to as aggregator/aggregation networks. When applying the MIL methods for WSIs, it should be noted that 1) the diagnosis for WSI analysis can be based on different tissue regions with multiple concepts — the collective MIL assumption, 2) the bag length $n$ for a WSI can be extremely large, $e.g.$, about 8,000 on average[19]. Therefore, the bag MIL methods for WSIs are with learnable aggregators and trained in a two-stage procedure, $i.e.$, training the feature extractor and aggregator stage by stage. Current works mainly follow this formulation and improve the framework from both feature extractor and aggregator, while our proposed method aims to empower existing works from a causal perspective.

Analysis MIL through causal inference. As shown in Fig. 2(a), we formulate the MIL framework as a causal graph (a.k.a, Pearl’s structural causal model or SCM [24]), which contains three nodes: $X$: whole-slide pathological image (bag), $Y$: bag label, $C$: bag contextual information.

Stage 1: Training feature extractors. We learn a feature extractor $f(\cdot)$ on the patchified images of WSIs $\{x_{1},...,x_{n}\}$, aiming at encoding each instance as a discriminative feature vector.

Stage 2: Training aggregators. Given the features of instances $\{b_{1},...,b_{n}\}$, the aggregator employs MIL pooling $\sigma(\cdot)$ to assemble them into a bag feature $B$ sequentially or simultaneously, and a classifier $g(\cdot)$ for discrimination. Formally, the loss for training aggregator is defined as:

$$\mathcal{L}=-\frac{1}{N}\sum_{i=1}^{N}Y_{i}\log\hat{Y}_{i}+\left(1-Y_{i}\right)\log\left(1-\hat{Y}_{i}\right),$$

(4)

where $N$ is number of bags in the training set. Note our IBMIL is no not limited to specific feature extractor or aggregators, including the architectures and training paradigms. Please refer to Sec. 4 for our choices.

Stage 3: Causal intervention via backdoor adjustment. The traditional two-stage bag MIL stops at stage 2 and uses the trained models for inference directly. Instead, we introduce another stage of interventional training, which needs the practical implementation of Eq. 3. Note that backdoor adjustment assumes that we can observe and stratify the confounders of a bag context. Thanks to the powerful ability of deep MIL models, context information is naturally encoded in the higher-level layers [43, 20]. To constitute the confounder set, we use a confounder dictionary $C=\left[c_{1},\ldots,c_{K}\right]$ for approximation, as collecting all confounders is impossible. Given the trained feature extractor and aggregator, we use $K$-means over all the bag features in the training set, partitioning the bags into clusters. We average the bag features of each cluster to represent a confounder stratum $c_{i}$, resulting in a confounder dictionary with the shape of $d\times K$, where $d$ is the dimension of bag features. Note that our approximation is reasonable in that these global clusters are susceptible to the visual biases [29], which is exactly the confounders. Then, we define:

$$\begin{array}[]{l}h(X,c_{i})=\alpha_{i}c_{i},\\ \left[\alpha_{1},\cdots,\alpha_{K}\right]=\operatorname{softmax}\left(\frac{\left({W}_{1}B\right)^{T}\left({W}_{2}C\right)}{\sqrt{l}}\right),\end{array}$$

(5)

where $B=\sigma(f(X))$ is the bag feature, ${W}_{1},{W}_{2}\in\mathbb{R}^{l\times d}$ are two learnable projection matrices to project bag feature ${B}$ and confounder ${C}$ into a joint space, and $\sqrt{l}$ is used for feature normalization [33]. Since the prediction comes from both bag $X$ and confounder $C$ (see Fig. 2(a)), we further define

$$P(Y\mid X,h(X,c_{i}))=P(Y\mid B\oplus h(X,c_{i})),$$

(6)

where $\oplus$ denotes vector concatenation, and other implementations can be found in ablation studies. We assume $P({c}_{i})$ is a uniform prior of $1/K$ for a safe estimation, and a more reasonable assumption, $e.g.$, incorporating expert knowledge, will be our future work. Plugging Eq. 5, Eq. 6 and defined $P({c}_{i})$ into Eq. 3, we are ready to calculate $P(Y|do(X))$ via passing the network multiple times. In practice, to avoid the expensive cost, we further apply Normalized Weighted Geometric Mean [38] to move the outer sum into the Softmax:

$\displaystyle P(Y\mid do(X))\approx P\left(Y\mid B\oplus\sum\nolimits_{i=1}^{K}\alpha_{i}{c}_{i}P(c_{i})\right).$

(7)

Thus, backdoor adjustment can be achieved by one feed-forward of the network.

Compatible with large-scale unlabelled datasets. We constitute the confounder set in an unsupervised fashion. One alternative implementation is to use the available bag labels for guidance, preserving the intra-class variation and capturing the class-relevant characteristics . There are two main reasons for our choice. 1) The unsupervised fashion makes our scheme compatible with large-scale unlabelled datasets, $e.g.$, The Cancer Genome Atlas (TCGA), for better approximation of confounders. 2) The confounder could be irrelevant to the class identity, $e.g.$, the stain color of positive and negative instances can be the same. We explore the other implementations in Sec. 4.3.

One possible more elegant scheme. As we need the trained aggregator to generate the bag features (the stage 2), one more stage is needed to retrain the aggregator (the stage 3). We are thus motivated to further simplify our scheme to avoid extra computational cost. Specifically, we can achieve the bag features by applying the traditional non-parametric aggregators, $e.g.$, max/mean-pooling, to the instances in a bag. It is inspired by the fact that these non-parametric aggregators serve as strong baselines, and we conjecture that statistic bag information they provide can be used for a reasonable approximation of confounders. Therefore, we can omit the stage 2. The experiment results in Sec. 4.3 support that our scheme can be more elegant.

Connection to other methods. Embedding-based MIL: As we approximate the confounder set based on bag features, these confounders can be seen as a denoised abstraction of bag features. From this perspective, we share the same spirit with the embedding-based MIL [35], $i.e.$, exploring the relations between bags. That means our IBMIL also explains the effectiveness of embedding-based MIL. Color Normalization: Some works [47] propose color normalization methods for H&E stained WSIs. However, color is just one of the confounders, and some confounders are even unobserved. Our method does not focus on color only, and thus is the more reasonable partially observed children of the unobserved confounder [11]; Instance augmentation: IMIL [20] uses strong instance augmentation to train the feature extractor for instance prediction. However, the augmentation may affect the statistical information in the bag. Therefore, our method is more suitable for bag MIL. Remix [39] proposes data augmentations for MIL by exploring the relations of instances, but our method explores the bag-level relations based on the causal theory.

Size of confounder dictionary. We ablate size $K$ of the confounder dictionary on three feature extractors, including ResNet-18, CTransPath and ViT. From Fig. 4, the performance of IBMIL is relatively robust to the size of confounder dictionary. Therefore, we need not elaborately tune this hyper-parameter and an arbitrary size within a wide range is able to boost the performance.

Dimension of joint feature space. As mentioned above, the confounders and bag features are projected into a joint feature space with a dimension of $l$ and attention scores are calculated subsequently. We ablate dimension $l$ on three feature extractors. The results in Fig. 4 reveal that performance does not improve monotonically with increased dimension and is saturated at $l=256$. We choose a dimension of 128 as the default configuration.

Learnable vs. unlearnable confounders. We explore the effect of learnable and unlearnable confounders. For the former, we update them in an end-to-end manner via backpropagation.

As can be seen, both of them outperform the baseline accuracy of 81.43%. However, freezing confounders during interventional training beats learnable confounders by 1.43% on accuracy. The reason may be that it is challenging to learn both confounders and interventional training with only bag-level labels, and introducing context-level supervision could be a solution [33]. We set confounders unlearnable as the default configuration.

Implementation of backdoor adjustment. We study the effect of different implementations of backdoor adjustment. Given a bag feature $B\in\mathbb{R}^{d}$ and the combination of confounders $\sum\nolimits_{i=1}^{K}\alpha_{i}{c}_{i}P(c_{i})\in\mathbb{R}^{d}$, we explore three variants to combine them, $i.e.$ $B\star\sum\nolimits_{i=1}^{K}\alpha_{i}{c}_{i}P(c_{i})$ and $\star\in\{\oplus,+,-\}$, where $+/-$ is element-wise addition/subtraction.

We observe all these implementations can lead to performance improvements, which demonstrates the stability and effectiveness of the proposed intervention.

Do improvements come from more epochs? Note that our proposed method requires an extra stage to train the aggregator. A natural question is whether we can improve baseline performance by training the baseline methods for as many epochs as the extra stage. Fig. 4 displays that more epochs do not bring about performance improvement, showing that our proposed method could empower baseline methods by backdoor adjustment rather than more training epochs. In most cases, training longer even brings performance degradation, which can be caused by the over-fitting problem in MIL [44].

Is stage 2 necessary? Recent MIL methods aggregate the instance features into a bag feature via the weighted average operator. The weights, also referred to as attention scores, are generated by parametric networks, which need an extra stage of training aggregators. Alternatively, we turn to three non-parametric settings to skip this stage and efficiently achieve the bag features. We consider:

• •

  “Mean” / “Max” denotes a bag feature is obtained through a mean-pooling / max-pooling layer among a bag of instance features, which is inspired by the strong baseline of non-parametric MIL method [35].

• •

  “Instance” denotes that $K$-means is directly performed over all the instance features in training set, since each instance can be regarded as a bag with length of one.

Then, interventional training is applied to baseline methods (including ABMIL and DSMIL) and we report the results in Fig. 4. Notably, even with such simple aggregation strategy, IBMIL still outperforms the baseline, and remains competitive or even better compared to “Default” setting, which indicates stage 2 in our scheme is unnecessary. By omitting stage 2, our scheme can be more elegant without performance degradation in most cases.

Can IBMIL improve non-parametric baselines? Besides using non-parametric aggregators to generate bag features for confounder set, we further take them as baselines and verify whether IBMIL is also able to improve them ($i.e.$, max/mean-pooling). Surprisingly, in Tab. 2, IBMIL brings significant improvements under all settings, where the best performance is even comparable to these attention-based aggregators. It indicates that IBMIL is indeed compatible with all compared MIL methods, including the non-parametric ones.

Constituting confounder set w/ or w/o bag labels? Given bag labels, we explore the class-specific $K$-means. In particular, we apply $K$-means to each class respectively, preserving the intra-class variation and class-relevant characteristics. From Fig. 5, we observe no obvious performance gap between class-specific and class-agnostic $K$-means. We conjecture that 1) the confounders could be independent of the class identity, and 2) bag features are already separable by classes. We will explore the way of incorporating bag labels in future work. On the other hand, the unsupervised fashion makes our scheme compatible with large-scale unlabelled datasets. We explore more unlabelled bags via combining the bags of TCGA and Camelyon16, and constituting the confounder set via the non-parametric aggregators. From Fig. 5, we observe a clear improvement on AUC under both max- and mean-poolings. That indicates, with more bags, our implementation can achieve better approximation of confounders.

Is IBMIL just post-processing? Since our proposed IBMIL shares some commonalities with the embedding-based MIL [35], one may ask: Do the improvements only come from exploring the bag relations? To answer this question, we make minor modifications on ABMIL. Instead of interventional training with confounders, we obtain the confounder dictionary via the class-specific $K$-means and treat it as a KNN classifier for evaluation. As can be seen, it brings limited improvements or even degrades the performance, verifying the improvements comes from interventional training, which is not just post-processing.