Mix, Align, Distil: Reliable Cross-Domain Atypical Mitosis Classification

TL;DR

This paper introduces a robust domain-adaptive method for classifying atypical mitotic figures in histopathology images, combining style perturbations, feature alignment, and knowledge distillation to handle domain shifts effectively.

Contribution

It proposes a novel training recipe that enhances domain robustness in mitosis classification using style perturbations, attention-based feature alignment, and EMA teacher distillation.

Findings

Achieved balanced accuracy of 0.8762 on leaderboard

Attained sensitivity of 0.8873 and specificity of 0.8651

Reached ROC AUC of 0.9499 with negligible inference overhead

Abstract

Atypical mitotic figures (AMFs) are important histopathological markers yet remain challenging to identify consistently, particularly under domain shift stemming from scanner, stain, and acquisition differences. We present a simple training-time recipe for domain-robust AMF classification in MIDOG 2025 Task 2. The approach (i) increases feature diversity via style perturbations inserted at early and mid backbone stages, (ii) aligns attention-refined features across sites using weak domain labels (Scanner, Origin, Species, Tumor) through an auxiliary alignment loss, and (iii) stabilizes predictions by distilling from an exponential moving average (EMA) teacher with temperature-scaled KL divergence. On the organizer-run preliminary leaderboard for atypical mitosis classification, our submission attains balanced accuracy of 0.8762, sensitivity of 0.8873, specificity of 0.8651, and ROC AUC of 0.9499. The method incurs negligible inference-time overhead, relies only on coarse domain metadata, and delivers strong, balanced performance, positioning it as a competitive submission for the MIDOG 2025 challenge.

Tables

Table 1. Table 1: Preliminary leaderboard results (Track 2)

Method	Balanced Acc.	Sensitivity	Specificity	ROC AUC
Baseline	0.7933	0.9014	0.6851	0.8859
Ours	0.8762	0.8873	0.8651	0.9499

Equations

$$\tilde{\mu }=\lambda \mu (x)+(1-\lambda )\mu (x^{\prime })\tilde{\sigma }=\lambda \sigma (x)+(1-\lambda )\sigma (x^{\prime })$$

$$\tilde{x}=\tilde{\sigma }\odot \frac{x-\mu (x)}{\sigma (x)+\epsilon }+\tilde{\mu }.$$

$$\hat{F}=\mathrm{CBAM}(F)$$

$$f_{i,c}\;:=\;\operatorname{GAP}\!\big{(}F_{i,c,:,:}\big{)}\;=\;\frac{1}{HW}\sum_{h=1}^{H}\sum_{w=1}^{W}F_{i,c,h,w}$$

$${\mu }_{d,c}=\frac{1}{|{\mathcal{I}}_d|}\sum _{i\in {\mathcal{I}}_d}{f}_{i,c}\text{and}{\bar{\mu }}_c=\frac{1}{D}\sum _{d=1}^D{\mu }_{d,c}$$

$$s_{c}^{2}=\frac{1}{D}\sum_{d=1}^{D}\bigl{(}\mu_{d,c}-\bar{\mu}_{c}\bigr{)}^{2}$$

$$\mathcal{L}_{\mathrm{align}}=\frac{1}{C}\sum_{c=1}^{C}\log\!\bigl{(}1+s_{c}^{2}\bigr{)}$$

$${\theta }_T^{(k)}=m{\theta }_T^{(k-1)}+(1-m){\theta }_S^{(k)}m\in [0,1)$$

$$p_S^{(T)}=\mathrm{softmax}\left(\frac{z_S}{T}\right)p_T^{(T)}=\mathrm{softmax}\left(\frac{z_T}{T}\right)$$

$$\mathcal{L}_{\mathrm{KD}}\;=\;T^{2}\,\mathrm{KL}\!\big{(}p_{T}^{(T)}\,\|\,p_{S}^{(T)}\big{)}$$

$${\mathcal{L}}_{\mathrm{total}}={\mathcal{L}}_{\mathrm{cls}}+{\lambda }_{\mathrm{align}}{\mathcal{L}}_{\mathrm{align}}+{\lambda }_{\mathrm{KD}}{\mathcal{L}}_{\mathrm{KD}}$$

Full text

{corrauthor}

kaustubh.atey@iitb.ac.in

1 Introduction

Atypical mitotic figures (AMFs) reflect aberrant cell division with mis-segregation of genetic material and are associated with aggressive disease, higher metastatic risk, and poorer outcomes. Even expert pathologists find it difficult to consistently distinguish AMFs from normal mitoses, which motivates automated analysis. Beyond detection, reliable classification of AMFs would enable reproducible assessment of tumor aggressiveness, yet prior work on this subtyping task remains limited. MIDOG 2025 Task 2 formalizes the problem as binary patch-level classification (atypical vs. normal mitosis) under realistic variability. Algorithms must address severe class imbalance, substantial intra-class morphological diversity, subtle inter-class differences, and pronounced domain shifts across tumors, species, stains, and scanners, while maintaining performance on unseen acquisition conditions.

The MIDOG 2022 report (1) emphasized domain shift as a central obstacle: models trained on particular scanners, stains, laboratories, or species showed noticeable degradation on unseen settings. It recommended domain-aware experimental design, including patient or slide disjoint splits, domain-stratified validation, and containerized submissions, to achieve true cross-site robustness. Methodologically, successful entries relied on stain and color augmentation, stain normalization, and explicit domain generalization methods such as feature alignment or adversarial training, often combined with self-supervision, hard-negative mining, test-time augmentation, and ensembling. The report also noted calibration and threshold transfer issues across domains and highlighted annotation ambiguity for borderline cases, reinforcing the need for morphology-centric, domain-invariant representations and principled cross-domain evaluation.

We propose a lightweight training-time recipe for domain-robust atypical mitosis classification. The approach increases feature diversity via style perturbations in the backbone, aligns attention-refined features across sites using weak domain labels, and yields stable predictions by distilling from an EMA teacher. On the preliminary leaderboard, the model shows strong and balanced performance.

2 Methodology

To improve the generalization ability of our model across domains, we employ three components: (i) MixStyle layers inside the feature extractor to promote domain-agnostic, diverse feature extraction; (ii) a feature-alignment loss on backbone features to explicitly align representations from different domains; and (iii) knowledge distillation using an exponential moving average (EMA) teacher to obtain a model that remains stable across domains.

2.1 Simulating Inter Domain Style Shifts

We use MixStyle (2) to simulate inter-domain stain and contrast shifts directly in feature space while preserving class labels. In histopathology, cross-site variation largely stems from staining and imaging differences. By perturbing only the style statistics (channelwise mean and variance) of intermediate features, MixStyle leaves morphology intact and pushes the encoder to rely on structure rather than colorimetric quirks.

Formulation: Given per-instance backbone features $x\in\mathbb{R}^{C\times H\times W}$ with channel-wise statistics $(\mu(x),\sigma(x))$, we draw $\lambda\sim\mathrm{Beta}(\alpha,\alpha)$ and form mixed statistics with another feature $x^{\prime}$:

[TABLE]

and re-standardize

[TABLE]

Here, $\varepsilon>0$ ensures numerical stability and $\odot$ is channelwise multiplication. Here $\tilde{x}$ is supervised with the source label, without any label blending.

By convexly mixing style statistics across instances, MixStyle broadens the training distribution along color and contrast axes that dominate cross-site variability while holding content fixed. This acts as an implicit domain generalizer, improving robustness on unseen centers and scanners with negligible overhead.

2.2 Feature Alignment with CBAM

To complement the primary binary cross-entropy (BCE) classification objective, we introduce an auxiliary feature alignment loss that reduces inter-domain mismatch. Let $F\in\mathbb{R}^{B\times C\times H\times W}$ be encoder features extracted before global average pooling (GAP). We refine them using the Convolutional Block Attention Module (CBAM) (3):

[TABLE]

which applies channel followed by spatial attention to emphasize morphology while suppressing stain and illumination artifacts. We impose alignment on $\hat{F}$ (instead of raw $F$) because attention-refined representations are easier to align across domains. These refined features are then passed to the classifier head.

Domain-aware alignment. Using domain labels within each mini-batch, we reduce feature-distribution discrepancies across domains. Let $\hat{F}_{i}\in\mathbb{R}^{C\times H\times W}$ be the CBAM feature map for sample $i$ ($C$ channels, spatial size $H{\times}W$), and let $\mathcal{I}_{d}$ denote the index set of samples from domain $d\in\{1,\dots,D\}$. For channel $c$, define the per-sample channel descriptor by global average pooling (GAP)

[TABLE]

We then define the per-domain channel mean $\mu_{d,c}$ and the inter-domain mean $\bar{\mu}_{c}$ as follows:

[TABLE]

Using these means, we compute the inter-domain variance as:

[TABLE]

We use a log-stabilized channel average as the alignment loss:

[TABLE]

which penalizes large inter-domain discrepancies. At test time, the encoder followed by CBAM produces a refined, domain-invariant representation that is passed to the classifier head to make predictions.

2.3 Knowledge Distillation with an EMA Teacher

To promote prediction stability across domains, we employ knowledge distillation (KD) (4) with a teacher maintained as an exponential moving average (EMA) of the student (5). After each student update $\theta_{S}^{(k)}$, the teacher is updated as

[TABLE]

Given student logits $z_{S}$ and teacher logits $z_{T}$, we form temperature-softened distributions

[TABLE]

and minimize the temperature-scaled KL divergence

[TABLE]

with gradients flowing only through the student. This EMA teacher acts as a temporal ensemble that smooths spurious domain-specific variations, encouraging the student to match stable targets rather than noisy, site-dependent signals. In practice, this reduces sensitivity to stain and illumination shifts while improving cross-site generalization.

Training objective: Our final objective combines the primary binary classification loss with the auxiliary alignment and distillation terms:

[TABLE]

Here, $\mathcal{L}_{\mathrm{cls}}$ is the BCE-with-logits loss on the classifier output, $\mathcal{L}_{\mathrm{align}}$ is the domain-aware alignment loss defined in Section 2.2, and $\mathcal{L}_{\mathrm{KD}}$ is the temperature-scaled distillation loss from the EMA teacher (Section 2.3). The weights $\lambda_{\mathrm{align}},\lambda_{\mathrm{KD}}>0$ balance the auxiliary terms.

3 Implementation Details

This section details on the datasets and splits used for training, validation, and testing, followed by the augmentation pipeline applied to all inputs. Next, we summarize the feature extractor configuration and its intermediate refinement prior to aggregation. Finally, we provide training details, including objectives, optimization settings and regularization.

3.1 Datasets and Splits

We use the MIDOG 2025 Atypical Training Set (6) together with the MIDOG 2021 subset of AMi-Br (7) for model development. The combined pool is partitioned into training and validation sets with a 75/25 split, ensuring no data leakage between splits. For held-out evaluation, we use two unseen test domains: the TUPAC16 subset of AMi-Br and the OMG-Octo Atypical dataset (8). These sets serve solely to assess generalization to unseen domains and are never used for tuning or selection. A model is considered for submission only if it improves on these held-out sets.

For domain supervision, domain labels are defined as follows: within MIDOG 2025 Atypical Training Set, each domain corresponds to a unique combination of Scanner, Origin, Species, and Tumor. The MIDOG 2021 subset of AMi-Br constitutes a separate domain. These domain labels are used exclusively for the domain alignment loss.

3.2 Data Augmentation

To improve robustness while preserving mitotic morphology, we apply a compact augmentation suite covering staining, illumination, and mild geometric variability. Scale/geometry: random resized cropping, horizontal/vertical flips, rotations, and mild perspective warp. Photometric: color jitter, Gaussian blur, automatic contrast adjustment, and sharpness adjustment. Synthetic artifacts: gamma adjustment, posterization, and solarization to mimic acquisition/compression effects. Coarse occlusion: random erasing with multiple window sizes.

Collectively, these transforms broaden the training distribution to approximate cross-site shifts, promoting invariance to nuisance factors and improving generalization to unseen domains.

3.3 Model Architecture and Training

Model architecture.

We adopt a DenseNet-121 backbone (9) for feature extraction, using its convolutional trunk to produce pre-pooling feature maps. To expose intermediate representations to style perturbations, we insert MixStyle blocks at early and mid stages (Beta mixing with $\alpha=0.1$). The resulting features are first refined with a CBAM module and then, after global average pooling, fed to the classification head.

Training setup.

We optimize with AdamW (initial learning rate $10^{-3}$, weight decay $10^{-2}$), batch size $512$, and a ReduceLROnPlateau scheduler that lowers the learning rate when the validation metric stalls. Gradients are clipped at a max norm of $1.0$. A teacher network is maintained as an EMA of the student ($m{=}0.999$). Knowledge distillation uses temperature $T{=}2.0$ with base weight $0.5$, linearly warmed up over the first 10 epochs and applied every batch. Validation and early stopping monitor the EMA model’s metric. The feature-alignment weight $\lambda_{\mathrm{align}}$ follows the DANN annealing schedule (10), ramping from [math] to $1$ across training.

4 Results

Leaderboard performance (Track 2: Atypical Mitosis Classification).

On the Preliminary Evaluation Phase leaderboard, our submission achieves the following organizer-reported metrics (see Table 1). These results indicate strong discrimination (ROC AUC = $0.9499$) with balanced operating characteristics (sensitivity = $0.8873$, specificity = $0.8651$), reflecting robust cross-domain generalization for atypical mitosis classification.

5 Discussion

Our method addresses domain shift through three complementary components: (i) MixStyle, which perturbs feature-space style statistics to reduce sensitivity to stain and illumination variation; (ii) CBAM-refined feature alignment, which explicitly reconciles representations across domains; and (iii) EMA-teacher distillation, which supplies a temporally stable target without adding inference-time cost.

On the organizer-run preliminary leaderboard for atypical mitosis classification, our approach achieves strong and balanced results (Table 1), indicating robust generalization to unseen domains. These findings position our submission as a strong contender for the MIDOG 2025 challenge.