[Paper Review] Task-specific Fine-tuning via Variational Information Bottleneck for Weakly-supervised Pathology Whole Slide Image Classification

Honglin Li, Chenglu Zhu+6 more|arXiv (Cornell University)|Mar 15, 2023

AI in cancer detection6 citations

TL;DR

This paper proposes a variational information bottleneck-based fine-tuning framework for weakly-supervised whole slide image (WSI) classification that reduces computational cost by distilling over 10,000 patches into fewer than 1,000 key instances per WSI. By enabling end-to-end training of the backbone using only WSI-level labels, the method achieves significant accuracy and generalization gains over ImageNet-pretrained and self-supervised features, with performance approaching fully supervised baselines using less than 1% of patch annotations.

ABSTRACT

While Multiple Instance Learning (MIL) has shown promising results in digital Pathology Whole Slide Image (WSI) classification, such a paradigm still faces performance and generalization problems due to challenges in high computational costs on Gigapixel WSIs and limited sample size for model training. To deal with the computation problem, most MIL methods utilize a frozen pretrained model from ImageNet to obtain representations first. This process may lose essential information owing to the large domain gap and hinder the generalization of model due to the lack of image-level training-time augmentations. Though Self-supervised Learning (SSL) proposes viable representation learning schemes, the improvement of the downstream task still needs to be further explored in the conversion from the task-agnostic features of SSL to the task-specifics under the partial label supervised learning. To alleviate the dilemma of computation cost and performance, we propose an efficient WSI fine-tuning framework motivated by the Information Bottleneck theory. The theory enables the framework to find the minimal sufficient statistics of WSI, thus supporting us to fine-tune the backbone into a task-specific representation only depending on WSI-level weak labels. The WSI-MIL problem is further analyzed to theoretically deduce our fine-tuning method. Our framework is evaluated on five pathology WSI datasets on various WSI heads. The experimental results of our fine-tuned representations show significant improvements in both accuracy and generalization compared with previous works. Source code will be available at https://github.com/invoker-LL/WSI-finetuning.

Motivation & Objective

• To address the high computational cost and performance limitations of existing weakly-supervised WSI classification methods using frozen ImageNet or self-supervised representations.
• To enable end-to-end fine-tuning of the backbone in a Multiple Instance Learning (MIL) framework despite the large number of patches in gigapixel WSIs.
• To improve feature task-specificity and generalization by converting task-agnostic self-supervised features into task-specific representations using WSI-level labels.
• To support versatile training-time augmentations in a weakly-supervised setting to enhance robustness under domain shift.

Proposed method

• Introduces a Variational Information Bottleneck (VIB) module to distill over 10,000 patches per WSI into a sparse set of less than 1,000 most representative instances, reducing training cost by over tenfold.
• Applies the VIB module to learn minimal sufficient statistics of the WSI bag, enabling gradient backpropagation through the backbone despite the large number of patches.
• Combines the VIB-based fine-tuning with self-supervised pretraining to convert task-agnostic SSL features into task-specific representations using only WSI-level labels.
• Employs a Lagrangian optimization framework to balance information compression and prediction accuracy, ensuring the model retains essential diagnostic features.
• Incorporates diverse training-time augmentations during fine-tuning to improve generalization on datasets with domain shift.
• Uses a differentiable top-K selection mechanism to identify the most supported instances within each WSI bag, forming the distilled bag for efficient training.

Experimental results

Research questions

• RQ1 Can a variational information bottleneck framework effectively reduce the computational burden of training on gigapixel WSIs while preserving diagnostic information?
• RQ2 To what extent can WSI-level weak supervision improve the task-specificity of self-supervised or ImageNet-pretrained features in whole slide image classification?
• RQ3 Does end-to-end fine-tuning of the backbone using distilled patch representations lead to better performance than frozen feature extraction?
• RQ4 Can the proposed method generalize effectively under domain shift, especially when combined with diverse data augmentations?

Key findings

• The proposed fine-tuning method achieves a 18.43% improvement in F1 score on Camelyon-16-C and 18.86% on Camelyon-17 compared to freezing ImageNet-pretrained features.
• Max-pooling with fine-tuned features outperforms CLAM-SB and TransMIL by 10.91% and 6.28% in F1 on Camelyon-16-C, respectively.
• The method improves generalization under domain shift, with DTFD-MIL achieving a 11.30% F1 gain on Camelyon-17 when fine-tuned, indicating strong robustness.
• The F1 metric shows significantly larger improvements than AUC, indicating that the method enhances classification performance (sensitivity and specificity) more than ranking ability, which is more clinically relevant.
• T-SNE visualization confirms that fine-tuned features form a more separable and task-specific distribution compared to ImageNet-1K or SSL features, with higher v-scores.
• Attention maps from CLAM-SB show that models using fine-tuned features focus more on tumor regions, demonstrating improved interpretability and diagnostic relevance.