Momentum Memory for Knowledge Distillation in Computational Pathology

Conference: CVPR 2026 arXiv: 2602.21395 Code: Available Area: Medical Imaging Keywords: Knowledge Distillation, Computational Pathology, Momentum Memory, Cross-modal Alignment, Multiple Instance Learning

TL;DR

This paper proposes MoMKD, which replaces conventional batch-local feature alignment with a momentum-updated class-conditional memory bank to enable cross-modal knowledge distillation from genomics to pathology whole-slide images (WSIs), achieving genomics-level predictive capability at inference using only H&E slides.

Background & Motivation

1. State of the Field

Multimodal learning (genomics + pathology) has demonstrated strong performance in cancer diagnosis; however, paired omics–pathology data are scarce in clinical settings. Knowledge distillation (KD) offers a practical solution: genomic supervision is utilized during training, while only pathology slides are required at inference.

2. Limitations of Prior Work

Existing pathology KD methods rely on batch-local alignment—performing feature matching or regression-based distillation within the current mini-batch. This paradigm suffers from three problems: (1) supervision signals are transient and unstable, defined solely by the current batch; (2) negative sample diversity is limited; and (3) in the MIL setting with gigapixel WSIs, abundant background-noise patches overwhelm the distillation signal, resulting in poor generalization.

3. Root Cause

Genomic data are information-dense predictors, whereas WSI features are high-dimensional and sparse. Direct joint training causes genomic gradients to dominate the WSI branch, and batch-local alignment is unstable across heterogeneous modalities.

4. Starting Point

Inspired by dynamic dictionaries in self-supervised learning (MoCo), MoMKD introduces a momentum memory as a distillation intermediary, replacing direct batch-level matching.

Method

Overall Architecture

MoMKD maintains a slowly evolving, class-conditional momentum memory bank (\(C^+\) and \(C^-\)) and drives two encoding branches in parallel: a GATv2-based WSI graph encoder and an MLP-based omics encoder. The two modalities do not interact directly but are aligned indirectly through the memory bank. During training, both modalities are encoded and aligned to the memory; at inference, only the WSI branch and the memory bank are required for prediction.

Key Designs

1. Dual-Branch Encoding with Spherical Projection

A spatial graph (\(k=8\) neighbors) is constructed for each WSI, and patch features \(F_{\mathrm{wsi}} \in \mathbb{R}^{I \times D}\) (\(D=256\)) are encoded by a two-layer GATv2 and projected onto the \(L_2\)-normalized hypersphere as \(\mathbf{F}_{\mathrm{N\text{-}wsi}} \in \mathbb{R}^{D_N}\) (\(D_N=128\)). Omics vectors are similarly encoded by an MLP and projected to \(\mathbf{F}_{\mathrm{N\text{-}omics}}\) on the same sphere.

Design Motivation: After normalization onto the hypersphere, the inner product is equivalent to cosine similarity (angle), eliminating the influence of norm discrepancies on cross-modal alignment.

2. Momentum Memory as Distillation Intermediary

The memory bank \(\mathcal{C}\) contains \(n\) components each for the positive class \(C^+\) and the negative class \(C^-\). Initialization is performed by K-means clustering on 10,000 randomly sampled patches. During training, the memory is updated slowly via alignment and regularization losses, accumulating global semantic information across batches.

Mechanism: The memory is not a simple instance cache but a highly compressed global semantic representation. The model aligns to this stable, slowly evolving intermediary rather than chasing the noisy within-batch distribution.

3. Indirect Distillation: Three-Step Alignment Mechanism

• Semantic Anchoring (Omics Alignment): Omics features are aligned to the memory with a self-supervised reconstruction constraint, injecting genomic semantics into the visually initialized memory.
• Knowledge Transfer (WSI Alignment): WSI features are aligned to the memory already calibrated by omics, forcing the WSI encoder to learn the modality-relevant correlations defined by genomics.
• Memory Evolution (Gradient Decoupling): No direct gradient flows between the omics and WSI branches; interaction occurs only indirectly through the memory. Classification head gradients are not back-propagated into the memory, preventing memory collapse.

4. Soft Angle-based Alignment Loss

Feature-to-memory similarities are aggregated via LogSumExp:

\[\phi(F, C) = \frac{1}{\tau_{\text{agg}}} \ln \sum_{j=1}^{n} \exp(\tau_{\text{agg}} F^T c_j)\]

The memory differential \(\Delta(F; C^+, C^-) = \phi(F, C^+) - \phi(F, C^-)\) is computed, and a softplus loss with margin=0.3 enforces proximity to \(C^+\) and distance from \(C^-\):

\[L_{\text{align}}(F, y=1) = \text{softplus}(\beta(\text{margin} - \Delta(F; C^+, C^-)))\]

Design Motivation: LogSumExp smoothly approximates the max similarity, ensuring gradient flow to all memory components; the margin prevents overfitting caused by perfect alignment.

5. Memory-Guided Unimodal Inference

At inference, each patch feature computes a differential affinity score \(\text{Score}_i\) with respect to \(C^+\) and \(C^-\). A softmax with temperature (\(\tau=0.2\)) generates attention weights, which are used to aggregate a slide-level representation. The memory acts as a global genomic anchor, directing attention toward patches consistent with omics-defined patterns.

Loss & Training

Total loss: \(L_{\text{total}} = \lambda_{\text{ce}} L_{\text{ce}} + \lambda_{\text{mse}} L_{\text{mse}} + \alpha_{\text{wsi}} L_{\text{align}}^{\text{wsi}} + \alpha_{\text{omics}} L_{\text{align}}^{\text{omics}} + \lambda_{\text{mem}} L_{\text{mem}}\)

• \(L_{\text{ce}}\): Classification cross-entropy (\(\lambda_{\text{ce}}=0.5\)), applied to the WSI branch only.
• \(L_{\text{mse}}\): Omics self-supervised reconstruction (\(\lambda_{\text{mse}}=0.01\)), preserving biological fidelity of omics encodings.
• \(L_{\text{align}}\): Cross-modal alignment loss (\(\alpha_{\text{wsi}}=0.2\), \(\alpha_{\text{omics}}=0.05\)).
• \(L_{\text{mem}}\): Memory regularization (\(\lambda_{\text{mem}}=0.1\)), comprising a VQ loss (patch-to-nearest-memory MSE) and an orthogonality constraint among memory components.

Feature backbone: UNI v2 (frozen); five-fold cross-validation; TCGA-BRCA dataset.

Key Experimental Results

Main Results

Table 1: Internal Comparison on TCGA-BRCA (AUC%)

Method	HER2 AUC	PR AUC	ODX AUC	Type
ABMIL	72.9±3.1	84.5±2.3	79.3±2.5	WSI-only
WIKG	75.5±5.0	84.9±3.0	78.3±3.7	WSI-only
TDC	76.2±2.1	84.7±5.3	81.0±2.2	Multimodal KD
MKD	77.1±2.3	85.1±1.2	80.1±1.5	Multimodal KD
G-HANet	76.1±5.6	85.0±2.3	80.5±1.3	Multimodal KD
MoMKD	79.6±0.7	87.9±0.9	82.3±2.3	Multimodal KD

MoMKD achieves top performance across all three tasks, with gains of +4.1%, +3.0%, and +4.0% over the best WSI-only baseline (WIKG), respectively.

Table 2: External Validation (In-house ODX)

Method	AUC	ACC	F1
DTFDMIL	76.2±2.2	86.5±1.5	63.5±3.9
TDC	76.5±2.1	86.2±3.0	63.5±3.2
MoMKD	79.4±0.8	87.1±1.7	68.0±3.0

MoMKD demonstrates strong cross-domain generalization, with AUC +2.9% and F1 +4.5%.

Ablation Study

Configuration	HER2 AUC (%)	Note
WSI baseline	73.9±3.1	No distillation
WSI + WSI Alignment only	75.2±2.4	Memory shaped by WSI only
WSI + Omics Alignment only	75.7±2.5	Memory calibrated by omics only
w/o Omics Recon	78.0±3.6	Unstable omics encoding
MoMKD (Full)	79.6±0.7	All components combined

Fixed vs. Momentum Memory: Momentum memory yields +4.4% on HER2 and +5.9% on the in-house set. Fixed memory collapses under domain shift (81.9→73.5%), while momentum memory remains robust (82.3→79.4%).

Key Findings

1. Momentum update is critical: Fixed memory performs reasonably within the source domain but degrades severely under domain shift, demonstrating that momentum updates are indispensable for resisting distributional drift.
2. Dual-modal alignment is complementary: Omics alignment injects semantics while WSI alignment transfers knowledge; neither alone is sufficient.
3. Adaptive memory capacity: More active memory components are retained for the harder HER2 task, while PR/ODX converge to fewer—indicating automatic adaptation to task complexity.
4. Visualization confirms biological validity: Positive memory components activate tumor-enriched and stromal interaction regions, while negative components activate adipose tissue and normal ducts, confirming that the memory captures biologically meaningful patterns.

Highlights & Insights

1. Transferring the MoCo dictionary concept to cross-modal KD: This elegantly resolves the instability of batch-local alignment; the memory simultaneously serves as an information bottleneck and a stable intermediary.
2. Elegant gradient decoupling design: The omics and WSI branches interact only indirectly through the memory, and classification head gradients do not affect the memory—triple isolation ensures slow, stable memory evolution.
3. Substantially reduced variance: MoMKD's standard deviations (0.7–2.3%) are far lower than those of competing methods (2–5%), indicating that the momentum mechanism enhances training stability.
4. Strong interpretability: Memory component-to-patch mappings can be visualized on WSIs, facilitating review by pathology experts.

Limitations & Future Work

1. Only binary classification evaluated: HER2/PR/ODX are all binary tasks; multi-class scenarios (e.g., fine-grained molecular subtyping) remain unexplored.
2. Manual memory size selection: The choice of \(n\) lacks an adaptive mechanism.
3. H&E staining only: IHC-stained WSIs may provide richer information, particularly for HER2.
4. Limited dataset scale: Each task in TCGA-BRCA contains only 800–1,000 cases; large-scale external validation is absent.
5. Single backbone: Only UNI v2 features are used; the impact of different patch encoders is not explored.

Related Work & Insights

• MoCo → MoMKD transfer: The insight from self-supervised learning—that "a large, stable dictionary is key to stable learning"—is successfully transferred to cross-modal KD.
• Evolution of pathology KD: TDC (gradient distillation) → MKD (online multi-teacher) → G-HANet (reconstruction distillation) → MoMKD (memory alignment), representing a progression from batch-local to global paradigms.
• Inspiration from VQ mechanisms: The VQ loss in memory regularization (patch-to-nearest-memory) is consistent with the VQ-VAE paradigm; replacing stop-gradient with EMA could be a promising direction.

Rating

⭐⭐⭐⭐ (4/5)

MoMKD innovatively transplants the MoCo dictionary concept into cross-modal KD, with elegant gradient decoupling and indirect distillation designs. The evaluation is thorough, encompassing three tasks, external validation, ablation studies, and visualizations, though the dataset scale is limited. The work establishes a new paradigm for cross-modal distillation in computational pathology.

Related Papers

• [CVPR 2026] A protocol for evaluating robustness to H&E staining variation in computational pathology models
• [ICLR 2026] Fusing Pixels and Genes: Spatially-Aware Learning in Computational Pathology
• [NeurIPS 2025] Revisiting End-to-End Learning with Slide-level Supervision in Computational Pathology
• [CVPR 2026] Cell-Type Prototype-Informed Neural Network for Gene Expression Estimation from Pathology Images
• [CVPR 2026] Prototype-Based Knowledge Guidance for Fine-Grained Structured Radiology Reporting