Momentum Memory for Knowledge Distillation in Computational Pathology

TL;DR

This paper introduces MoMKD, a novel knowledge distillation framework using momentum memory to improve multimodal learning in computational pathology, enabling better histology-only inference by aggregating cross-modal information across batches.

Contribution

We propose a momentum memory-based distillation method that enhances supervision in multimodal models and decouples gradient updates to improve stability and generalization.

Findings

MoMKD outperforms state-of-the-art methods on TCGA-BRCA benchmarks.

The framework improves histology-only inference accuracy.

It demonstrates strong generalization on independent datasets.

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.