Momentum Memory for Knowledge Distillation in Computational Pathology

Yongxin Guo; Hao Lu; Onur C. Koyun; Zhengjie Zhu; Muhammet Fatih Demir; Metin Nafi Gurcan

cs.CVFeb 24, 2026

Momentum Memory for Knowledge Distillation in Computational Pathology

Authors: Yongxin Guo, Hao Lu, Onur C. Koyun, Zhengjie Zhu, Muhammet Fatih Demir, Metin Nafi Gurcan

Organizations: Wake Forest University School of Medicine, Winston-Salem, NC, USA

Abstract

Multimodal learning that integrates genomics and histopathology has shown strong potential in cancer diagnosis, yet its clinical translation is hindered by the limited availability of paired histology-genomics data. Knowledge distillation (KD) offers a practical solution by transferring genomic supervision into histopathology models, enabling accurate inference using histology alone. However, existing KD methods rely on batch-local alignment, which introduces instability due to limited within-batch comparisons and ultimately degrades performance. To address these limitations, we propose Momentum Memory Knowledge Distillation (MoMKD), a cross-modal distillation framework driven by a momentum-updated memory. This memory aggregates genomic and histopathology information across batches, effectively enlarging the supervisory context available to each mini-batch. Furthermore, we decouple the gradients of the genomics and histology branches, preventing genomic signals from dominating histology feature learning during training and eliminating the modality-gap issue at inference time. Extensive experiments on the TCGA-BRCA benchmark (HER2, PR, and ODX classification tasks) and an independent in-house testing dataset demonstrate that MoMKD consistently outperforms state-of-the-art MIL and multimodal KD baselines, delivering strong performance and generalization under histology-only inference. Overall, MoMKD establishes a robust and generalizable knowledge distillation paradigm for computational pathology.

Figures & tables

Figure 1: a. The classical teacher-student knowledge distillation (KD) method. b. The correlation-based KD method. c. The proposed momentum memory KD method. Compared with a and b, the proposed method uses momentum memory for knowledge distillation to solve the batch-local problem.

Figure 2: The overall framework of the proposed momentum memory knowledge distillation framework. a) presents the multi-modal KD training process; b) indicates the proposed cross-modal alignment. Here we assume the input is a positive case so that the alignment loss is aiming to push it closer to the positive memory set (blue triangle) and pull it away to the negative memory set (yellow star); c) presents the uni-modal inference stage, and the visual interpretation on the memory.

Methods	HER2(%)			PR(%)			ODX(%)
Methods	AUC	ACC	F1	AUC	ACC	F1	AUC	ACC	F1
ABMIL [ 22 ]	72.9 $\pm$ 3.1	77.1 $\pm$ 1.8	64.8 $\pm$ 2.9	84.5 $\pm$ 2.3	78.8 $\pm$ 1.6	75.0 $\pm$ 2.1	79.3 $\pm$ 2.5	83.8 $\pm$ 2.6	68.8 $\pm$ 1.6
DSMIL [ 25 ]	71.3 $\pm$ 4.3	76.2 $\pm$ 2.8	63.5 $\pm$ 4.1	81.6 $\pm$ 2.8	76.4 $\pm$ 2.0	72.7 $\pm$ 2.2	78.7 $\pm$ 2.0	83.9 $\pm$ 1.9	70.8 $\pm$ 1.4
TransMIL [ 34 ]	69.8 $\pm$ 2.8	75.6 $\pm$ 1.9	61.1 $\pm$ 5.2	85.2 $\pm$ 1.0	79.0 $\pm$ 1.2	76.1 $\pm$ 1.6	77.9 $\pm$ 3.5	83.5 $\pm$ 3.7	66.8 $\pm$ 5.0
DTFDMIL [ 46 ]	74.4 $\pm$ 1.9	76.2 $\pm$ 1.0	61.7 $\pm$ 5.4	83.9 $\pm$ 2.4	79.6 $\pm$ 3.1	76.3 $\pm$ 2.9	79.1 $\pm$ 2.1	82.2 $\pm$ 3.1	68.9 $\pm$ 2.0
WIKG [ 26 ]	75.5 $\pm$ 5.0	77.1 $\pm$ 3.1	53.7 $\pm$ 4.4	84.9 $\pm$ 3.0	79.2 $\pm$ 3.3	76.1 $\pm$ 4.1	78.3 $\pm$ 3.7	84.1 $\pm$ 1.2	71.1 $\pm$ 2.5

Table 1: Internal comparison on the TCGA-BRCA dataset. The best performance is in bold, and the second-best performance is underlined.

Methods	In-house dataset(%)
Methods	AUC	ACC	F1
ABMIL	75.1 $\pm$ 1.7	86.1 $\pm$ 0.2	60.9 $\pm$ 3.7
DSMIL	74.3 $\pm$ 2.8	86.1 $\pm$ 0.7	61.2 $\pm$ 3.2
TransMIL	71.7 $\pm$ 2.4	85.0 $\pm$ 1.7	60.5 $\pm$ 4.1
DTFDMIL	76.2 $\pm$ 2.2	86.5 $\pm$ 1.5	63.5 $\pm$ 3.9
WIKG	75.9 $\pm$ 3.5	86.7 $\pm$ 1.4	58.3 $\pm$ 5.3

Table 2: External validation results for the ODX prediction task on the independent in-house dataset. The best performance is in bold, and the second-best performance is underlined.

Model Variant	Active Components	AUC (%)
Baseline	WSI solely	73.9 $\pm$ 3.1
MoMKD ( $\alpha_{omics}=0$ )	WSI + OMICS Recon + WSI Alignment	75.2 $\pm$ 2.4
MoMKD ( $\alpha_{wsi}=0$ )	WSI + OMICS Recon + OMICS Alignment	75.7 $\pm$ 2.5
MoMKD (w/o Recon)	WSI + Joint Alignment (WSI & OMICS)	78.0 $\pm$ 3.6
MoMKD (Proposed)	WSI + OMICS Recon + Joint Alignment	79.6 $\pm$ 0.7

Table 3: Ablation study on the proposed method. All settings are trained on the TCGA-BRCA dataset for the HER2 classification.

Figure 3: Memory dynamics across biomarker prediction tasks. Each column corresponds to one task ( HER2 , PR , and ODX ), while each row depicts one statistic of the memory storage: ( upper ) Active memory Ratio , and ( bottom ) Perplexity . The active memory ratio quantifies the proportion of memory components utilized during training, ensuring that no dead memory components emerge. The corresponding perplexity represents the effective number of active memory components ( $c_{j}^{+}$ and $c_{j}^{-}$ ). Across all tasks, the momentum memory maintains high active ratios ( $>0.75$ ) and perplexity, indicating stable and interpretable utilization of the memory storage.

Task	AUC (%)
	Fixed	Momentum
HER2	75.2 $\pm$ 3.0	79.6 $\pm$ 0.7
PR	84.7 $\pm$ 2.1	87.9 $\pm$ 0.9
ODX	81.9 $\pm$ 2.3	82.3 $\pm$ 2.3
In-house dataset	73.5 $\pm$ 3.7	79.4 $\pm$ 0.8

Table 4: Fixed memory vs. momentum memory.

Figure 4: Visualization result based on memory, with mapped patches in the WSI. The upper figure is the positive case with the lower one being negative based on the Oncotype DX test. For each memory component, we present the top-3 mapped patches here.

Knowledge Distillation

Knowledge distillation transfers useful behavior or representations from larger or specialized models into more compact or adaptable learners.

See all 819 papers

Pathway-Structured Privileged Distillation for Deployable Computational Pathology
Jun 1, 2026Yongxin Guo, Hao Lu, Onur Koyun +2Computational Pathology Oncology
School of Medicine, Wake Forest University, Winston-Salem, NC, USA.
A Generative Foundation Model for Multimodal Histopathology
Apr 4, 2026Jinxi Xiang, Mingjie Li, Siyu Hou +9Computational Pathology Pathology Foundation Models
Department of Radiation Oncology, Stanford University School of Medicine, Stanford, CA · Department of Psychiatry and Behavioral Sciences, Stanford University School of Medicine, Stanford, CA · Department of Statistics and Data Science, Yale University, New Haven, CT +5
Hierarchical Prototype-based Domain Priors for Multiple Instance Learning in Multimodal Histopathology Analysis
Apr 27, 2026Xuemei Qiu, Dawei Fan, Yebin Huang +2Cross-Modal Alignment Prototype Learning
College of Computer and Information Science, Fujian Agriculture and Forestry University, Fuzhou, 350002, China · Fujian Agriculture and Forestry University, Fuzhou, 350002, China · Department of Pathology, Clinical Oncology School of Fujian Medical University, and Fujian Cancer Hospital, Fuzhou, 350014, China +2

Jun 1, 2026J/K move · Enter open · S save

Pathway-Structured Privileged Distillation for Deployable Computational Pathology

Yongxin Guo, Hao Lu, Onur Koyun, Muhammet Demir+1

School of Medicine, Wake Forest University, Winston-Salem, NC, USA.

Computational Pathology Oncology Mixture-of-Experts Models Learning Using Privileged Information Knowledge Distillation Whole-Slide Image Analysis

Open paper AskPDF arXiv

Integrating transcriptomics and histopathology can improve cancer risk modelling, yet practical use is constrained by the limited availability of RNA profiling in routine settings. Here we introduce Mixture of Pathway Experts (MoPE), a knowledge-distillation framework that reframes multimodal learning as privileged distillation for histology-only inference. MoPE is motivated by the partial observability between RNA profiles and whole-slide images: histology can capture morphology-linked consequences of certain molecular programmes, but cannot be expected to reconstruct the full transcriptomic state. MoPE encodes RNA-derived pathways and transfers the molecular supervision to pathway-indexed pathology experts through memory-usage alignment. Across diverse public benchmarks and two independent breast cancer cohorts, MoPE consistently improved WSI-only inference performance relative to baseline methods. Pathway-usage analyses and human-audited visual inspection provide bounded inspection of model behaviour and candidate morphology-linked readouts. These results support pathway-structured privileged distillation as a promising route to using molecular information during training while preserving RNA-free inference.

Abstract

Knowledge Distillation

Explore similar work

Pathway-Structured Privileged Distillation for Deployable Computational Pathology

A Generative Foundation Model for Multimodal Histopathology

Hierarchical Prototype-based Domain Priors for Multiple Instance Learning in Multimodal Histopathology Analysis