ContiStain: Cross-Domain Relation-Preserving Distillation for Continual Multi-Domain Virtual IHC Staining

Abstract: A unified multiplex virtual staining model enables scalable and non-destructive multiplex analysis from H&E slides while promoting parameter efficiency, shared pathological knowledge, and consistent cross-biomarker representations. However, in clinical practice, data for new biomarkers are typically acquired sequentially over time. Fine-tuning on such temporally arriving data leads to severe performance degradation on previously learned biomarkers, as sequential optimization disrupts the structured relationships among biomarker representations in the latent space. To address this issue, we propose ContiStain, an IHC multi-domain relational distillation framework for continual virtual staining. We first (i) construct a domain-aware structured feature space using a mixture-of-experts (MoE) feature extractor to reduce representation interference across biomarker domains. Based on this stabilized feature space, we then (ii) propose a relation-preserving distillation strategy that explicitly enforces the consistency of cross-domain token-level cosine similarity matrices between learned biomarker domains during continual adaptation. By maintaining cross-domain structural coherence, ContiStain mitigates forgetting while retaining adaptability to new domains. Experiments on the MIST dataset under a four-domain sequential virtual IHC staining setting show improved stability, reducing FID and ConchFID by 11.1 and 60.9 compared to sequential fine-tuning, enabling scalable and robust multi-domain virtual staining. Code is released at https://github.com/ccitachi/ContiStain.

• The paper introduces ContiStain, a framework that mitigates catastrophic forgetting in sequential multi-domain virtual IHC staining using cross-domain relation-preserving distillation.
• It employs a Mixture-of-Experts structure for domain-aware latent feature extraction, ensuring both domain specificity and morphology consistency.
• Experimental evaluations on breast cancer biomarkers show significant improvements in FID, CONCH-FID, PSNR, and SSIM compared to baseline continual learning methods.

Cross-Domain Relation-Preserving Continual Learning for Virtual IHC Staining: An Analysis of ContiStain

Introduction

The paper "ContiStain: Cross-Domain Relation-Preserving Distillation for Continual Multi-Domain Virtual IHC Staining" (2607.03851) addresses the clinically motivated challenge of continual multi-domain virtual immunohistochemistry (IHC) staining, where biomarker domains are acquired and incorporated sequentially rather than simultaneously. The core issue in this scenario is catastrophic forgetting—sequential fine-tuning on newly arrived biomarker domains leads to severe degradation of performance on previously learned domains, compounded by disruption of the shared and structured relationships among biomarker representations in the latent space. The authors formalize this continual learning setting for multi-domain virtual staining and propose ContiStain, a framework built on explicit cross-domain relational preservation and domain-aware feature disentanglement, as a highly competitive solution.

Figure 1: Challenges of multi-domain continual virtual staining — sequential domain arrival, catastrophic forgetting with naïve fine-tuning, and structural relationship collapse motivate relationship-aware continual learning.

Methodology

Mixture-of-Experts for Domain-Aware Latent Structuring

ContiStain introduces a Mixture-of-Experts (MoE) feature extractor to construct a domain-aware structured latent space. Each biomarker domain (e.g., HER2, ER, PR, Ki67) is modeled with biomarker-specific experts, while shared experts capture tissue morphology–aligned features. Domain and content-adaptive expert routing via learned gating (conditioned jointly on the domain embedding and image content) ensures expert usage is both domain-aligned and morphology-consistent. The MoE component is regularized to prevent expert collapse and maintain balanced, domain-coherent representations.

Cross-Domain Relation-Preserving Distillation

Upon each sequential adaptation, ContiStain freezes the generator from the previous training stage (teacher), then distills knowledge into the student by enforcing the cross-domain relational geometry among latent representations to remain consistent. This is realized by computing patch-level cosine similarity matrices between biomarker domain features extracted using the MoE structure and minimizing the divergence between these matrices across the teacher and student models for all previously seen domain pairs. Crucially, this operation preserves both within-domain features and the inter-domain correlations that encode cross-biomarker morphological consistency.

Figure 2: ContiStain: continual virtual staining pipeline with MoE-structured features and relational distillation for domain-wise adaptation.

Combined Objective

The total loss consists of the original base objective (adversarial and contrastive patch-level supervision), the relational distillation loss enforcing structural consistency, and the MoE gating regularization. The framework does not require replay of prior domain data, operating under the strict continual learning constraint that only the current domain’s data are available.

Experimental Evaluation

Dataset and Setting

Experiments are conducted on the MIST dataset, which spans four key breast cancer biomarkers (HER2, Ki67, ER, PR), processed under the realistic scenario of sequential domain introduction. Models are trained without access to past domain samples, assessing retention and generalization as new domains are incorporated.

Comparative Performance and Robustness

ContiStain is benchmarked against state-of-the-art continual learning baselines—LwF, EWC, iCaRL, and naïve sequential fine-tuning. Across the canonical and reversed domain orderings, ContiStain consistently achieves the lowest FID and CONCH-FID (pathology-semantic FID) scores across most domains, reflecting strong fidelity and clinical realism. Notably, after continual learning, FID/ConchFID for the HER2 domain is reduced by 14.7/153.9 (compared to EWC) and 10.5/86.3 (compared to iCaRL); on average, the reduction is 11.1/60.9 over naïve sequential learning. PSNR and SSIM values remain stable or improve, with minimal degradation due to forgetting, in contrast to significant performance drops observed with baseline methods.

Qualitative and Structural Validation

Qualitative sample visualizations across all four biomarkers indicate that ContiStain preserves domain-specific staining phenotypes (membranous HER2, nuclear ER/PR/Ki67) and morphology, even in the absence of direct replay.

Figure 3: Qualitative visualization of virtual staining outputs for HER2, ER, PR, and Ki67 using ContiStain.

Visualization of feature distributions (t-SNE) reveals that latent features derived using ContiStain remain tightly aligned with the initial, pre-adaptation distributions, whereas sequential fine-tuning results in pronounced feature drift and loss of cross-domain alignment.

Figure 4: Comparison of domain feature distributions; ContiStain maintains proximity to the initial structure, mitigating drift seen in naïve sequential fine-tuning.

Ablation Analyses

Ablation confirms the efficacy of both key innovations. Adding relation-preserving distillation alone substantially improves both distributional and perceptual metrics; further inclusion of domain-aware MoE delivers additional gains. The optimal number of experts is empirically found to be five, with further increases yielding diminishing or negative returns. Excessively high relation loss weights adversely affect model stability, indicating the need for balanced multi-objective optimization.

Implications and Future Prospects

ContiStain’s focus on explicit cross-domain relational maintenance marks a significant shift from classic output- or parameter-regularization continual learning regimes. By structuring the latent space to be both domain-specific in appearance and domain-invariant in tissue morphology, and directly enforcing preservation of the inter-biomarker relationships, the method delivers both practical retention and theoretical consistency with the underlying tissue structure. This paradigm is particularly well suited to medical imaging contexts where domain relationships are semantically meaningful but data arrive in a non-i.i.d., evolving manner.

Practically, ContiStain enables robust deployment of scalable, unified virtual staining models in clinical workflows, accommodating progressive adoption of new biomarkers without retraining or severe catastrophic forgetting. Theoretically, the work opens directions for: (1) relational geometric continual learning in general multi-domain visual tasks beyond pathology, (2) advanced expert routing schemes, possibly with sparse or adaptive expert selection (Li et al., 2022), and (3) development of feature-level continual learning metrics tailored to medical domain constraints.

Conclusion

ContiStain provides a novel framework for continual multi-domain virtual IHC staining, mitigating catastrophic forgetting by combining domain-aware Mixture-of-Experts feature modeling and relation-preserving distillation. By explicitly maintaining the structured inter-domain relationships in the latent space, ContiStain yields superior cross-biomarker retention and fidelity in sequential training regimes. This approach sets the stage for robust, clinically scalable virtual staining systems capable of adapting to evolving diagnostic requirements and provides a methodological foundation for structured continual learning across a broad spectrum of multi-domain vision applications.