Uncertainty Aware Knowledge Distillation for Imbalanced Disease Grading

Analysis by the aitrendblend editorial team · Medical imaging AI·

Knowledge Distillation Prostate Cancer Grading Diabetic Retinopathy Class Imbalance Uncertainty Estimation
Knowledge Distillation : Illustration of two expert neural networks transferring grading knowledge to a smaller student network for prostate and retinal disease grading
A pathologist reading a prostate biopsy slide and an ophthalmologist grading a retinal photograph are doing the same basic thing. They are placing a patient somewhere on a severity scale using visual patterns that get rarer and harder to read the more advanced the disease gets. That rarity is exactly what trips up the AI models built to help them, because the training data for stage three prostate cancer or early stage diabetic retinopathy is thin by nature, and thin data teaches a model the wrong lesson about how confident it should be.

Key Points

  • Researchers from Zhejiang University and WeDoctor Cloud built UMKD, a framework that distills knowledge from two biased expert models into one smaller student model for grading prostate cancer and diabetic retinopathy.
  • The method separates structural image features from disease specific features, then uses a per region uncertainty score to decide how much to trust each expert during training.
  • Tested on SICAPv2 prostate histology and APTOS fundus photographs, UMKD beat five established distillation baselines including FitNet, RKD, the original Hinton style KD, DKD, and SDD.
  • The student model, a ResNet18, matched or beat its two ResNet50 teachers on several metrics while using roughly a third of the parameters.
  • An ablation study shows the shallow feature alignment step matters more than either of the other two components, a detail the paper’s abstract does not spell out.
  • This is a preprint posted to arXiv on May 1 2025 and has not yet completed formal peer review at the time of this writing.
This article explains published research and is not medical advice. Nothing here should be used to diagnose or treat prostate cancer, diabetic retinopathy, or any other condition. The performance numbers below come from controlled research datasets, not from deployed clinical systems, and a preprint has not yet been through full peer review. Anyone with questions about a diagnosis or a grading result should talk to a qualified clinician.

Why grading a disease from an image is harder than it looks

Gleason grading, the system pathologists use to score prostate cancer aggressiveness from tissue samples, has a documented interobserver variability of about 40 percent among trained pathologists looking at the same slide. That number alone tells you the task is genuinely hard for humans, not just for machines. Diabetic retinopathy grading has its own version of the problem. Early microaneurysms, the first visible sign of retinal damage, are missed or misclassified in more than a quarter of cases according to the clinical literature the paper cites. Catching that early stage matters a great deal, since timely intervention is associated with roughly a 90 percent reduction in blindness risk for diabetic retinopathy patients.

So there is real pressure to build automated grading tools that catch these early and intermediate stages reliably. The trouble is that the patients who most need early detection are, almost by definition, underrepresented in the data used to train these tools. In the SICAPv2 prostate dataset that the UMKD team used, stage III cases make up only about 8 percent of the full cohort. Everyone else falls into the more common categories. A model trained on that distribution learns to be very good at recognizing the common cases and quietly bad at the rare, clinically important ones. Researchers call this domain shift when it shows up as a mismatch between the distribution a model was trained on and the distribution it is deployed on, and imbalance makes that shift worse because the tail of the distribution barely gets represented at all.

Where knowledge distillation fits in

Knowledge distillation is the general technique of training a smaller student model to mimic the behavior of one or more larger teacher models, first popularized by Hinton and colleagues in 2015. The appeal for clinical AI is obvious. Hospitals often cannot deploy a huge model on limited hardware, and different institutions may already have their own trained expert models they cannot share raw patient data for, due to privacy rules. If you can distill the useful parts of several existing expert models into one compact model without moving the original patient data, you get both efficiency and a form of federated knowledge sharing.

Multi expert knowledge distillation has been applied to natural image problems with class imbalance before, but the UMKD authors point out that it remains underexplored for disease grading specifically. That gap is the reason this paper exists. Prior distillation methods, including relational approaches like RKD and feature hint methods like FitNet, were not designed with the idea that the teacher models themselves might be systematically biased by an imbalanced training set. If a teacher is confidently wrong about a rare class, naive distillation just passes that confident wrongness along to the student.

What UMKD actually does differently

The framework has three moving parts. Two of them operate in feature space, before the final classification layer, and one operates on the output logits. The authors call these shallow feature alignment, compact feature alignment, and uncertainty aware decoupled distillation.

Shallow feature alignment, or teaching structure before disease specifics

Early layers in a convolutional network tend to encode general structural information, edges, textures, and coarse shapes, rather than disease specific patterns. UMKD treats this structural information as something the student should absorb cleanly, without the noise that comes later. It does this with a frequency domain trick. For each expert model, the method applies average pooling at several different kernel sizes, effectively a set of low pass filters that keep the coarse structure and discard high frequency detail. The student gets a learnable version of the same idea, built from a downsampling convolution and a depthwise separable convolution, and its filtered output is aligned to each expert’s filtered output.

The alignment itself uses two loss terms combined. A Maximum Mean Discrepancy term measures how far apart the statistical distributions of student and expert features are, and a reconstruction loss keeps the expert model’s own features unchanged through the alignment process, which matters because the experts are treated as frozen and untouchable to respect the privacy constraint that keeps their original training data separate.

\[ \mathcal{L}_{\text{MMD}} = \frac{1}{B} \sum_{t=1}^{N} \left\| \sum_{i=1}^{B} \phi\left(\hat{F}_{T_t}^i\right) – \sum_{j=1}^{B} \phi\left(\hat{F}_S^j\right) \right\|_2^2 \]

Here \(B\) is the batch size, \(N\) is the number of expert models, and \(\phi\) is a mapping function applied to the projected features of expert \(T_t\) and student \(S\).

Compact feature alignment, or teaching the disease specific signature

Where shallow alignment handles coarse structure, compact feature alignment works on the deepest features, the ones right before the final fully connected classification layer. These features carry the disease specific signature the network has learned. Because the expert and student networks can have different architectures and therefore different feature dimensions, UMKD appends a small 1 by 1 convolutional adapter to each encoder so everything lands in a shared spherical feature space before alignment. The same MMD plus reconstruction loss combination from shallow alignment is reused here.

Uncertainty aware decoupled distillation, or knowing when to trust an expert

This is the component that gives UMKD its name, and it operates on the model outputs rather than internal features. The authors build on a 2022 technique called decoupled knowledge distillation, which splits the standard distillation loss into two separate terms, one for the target class prediction and one for the non target class predictions. UMKD adds a spatial twist and an uncertainty twist on top of that foundation.

First, instead of comparing whole logit maps at once, the method partitions them into regions at several scales, from the full image down to finer grids, and compares expert and student predictions region by region. Second, for each region it computes an uncertainty coefficient from the expert’s own prediction confidence. If the expert’s softmax output is close to a one hot distribution, meaning it is very sure of its answer, the uncertainty coefficient is close to zero. If the expert is unsure, spreading probability across several classes, the coefficient climbs toward one.

\[ \mathcal{L}_{\text{UDD}}(w,n) = (2 + U_{T_t}) \cdot \mathcal{L}_{\text{TCKD}} + (1 – U_{T_t}) \cdot \mathcal{L}_{\text{NCKD}} \]

where \(U_{T_t} = 1 – \max(\sigma(\psi_{T_t}(w,n)))\) measures how far the expert’s prediction at scale \(w\) and region \(n\) is from a confident one hot answer.

The effect is that regions where the expert is confidently wrong get more weight on precise logit matching, since the (1 minus uncertainty) term dominates there and enforces tight alignment on the parts the expert is sure about. Regions where the expert is uncertain get amplified target class supervision through the (2 plus uncertainty) term, which pushes the student to pay closer attention exactly where an imbalance biased expert is most likely to be shaky. That is a reasonable way to encode the intuition that an expert trained on a skewed dataset should be trusted less on the classes it rarely saw, without requiring anyone to manually specify which classes those are.

All three components feed into one combined loss with two weighting coefficients that balance feature alignment against output distillation, alongside the standard classification cross entropy loss on the true labels.

\[ \mathcal{L}_{\text{Total}} = \mathcal{L}_{\text{cls}} + \alpha \cdot (\mathcal{L}_{\text{SFA}} + \mathcal{L}_{\text{CFA}}) + \beta \cdot \sum_{w \in W} \sum_{n \in N_w} \mathcal{L}_{\text{UDD}}(w,n) \]
The interesting design choice here is not any single loss term, it is treating expert uncertainty as a per region signal rather than a single number per model. A confident expert can still be wrong in one corner of an image, and a generally shaky expert can still nail the obvious cases. Scoring uncertainty at the patch level rather than the whole image level is what lets UMKD react to that.

How the experiments were set up

The team tested two distinct and genuinely different imbalance scenarios, which is worth understanding because they measure different failure modes. In the source imbalanced setting, the two expert models are each trained on an imbalanced dataset, then distillation happens onto a class balanced target dataset built by random sampling. In the target imbalanced setting, the experts are trained on balanced data, and distillation happens onto an imbalanced target dataset that mirrors the real world class distribution. The balanced subsets used per grading category were 2500, 2222, 2500, and 948 samples for SICAPv2, and 600, 370, 300, 193, and 295 for APTOS.

Both datasets were split 8 to 1 to 1 for training, validation, and testing, with random cropping and flipping for augmentation. One detail worth calling out because it shows domain awareness, the team deliberately excluded color jittering, a very common augmentation for natural images, because pathology images are sensitive to color shifts that could disrupt the actual diagnostic signal in the tissue stain. That is a small methodological choice, but it is the kind of thing that separates a paper written by people who understand medical imaging from one that treats it as just another image classification benchmark.

The two expert models were ImageNet pretrained ResNet50 networks, and the student was an ImageNet pretrained ResNet18, giving roughly a threefold reduction in parameter count from teacher to student. Baselines compared against UMKD included FitNet and RKD as feature based methods, and Hinton style KD, DKD, and SDD as logit based methods, plus the raw performance of each individually trained ResNet as a reference point.

What the results actually show

On SICAPv2 prostate grading, UMKD posted the best overall accuracy, mean accuracy, weighted F1, and lowest mean absolute error among all distillation methods in both imbalance settings.

SICAPv2 prostate grading, source imbalanced setting, percent except MAE
MethodOAMean AccF1MAE
ResNet50 Expert 191.5389.4891.470.1098
ResNet50 Expert 292.0589.8891.930.1100
ResNet18 Student baseline89.5889.0689.490.1463
FitNet78.7878.4178.420.3136
RKD88.5488.2488.440.1690
KD (Hinton)89.0688.3988.970.1624
DKD86.9185.0186.680.1739
SDD87.8286.6687.670.1594
UMKD91.0290.2390.940.1294

That mean accuracy figure is the one worth sitting with. UMKD reaches 90.23 percent mean accuracy, which is a class balanced metric, while the raw ResNet18 student baseline sits at 89.06 and the strongest prior distillation baseline, SDD, sits at 86.66. A gain of 3.57 points in mean accuracy specifically, more than the gain in overall accuracy, tells you the improvement is concentrated in the harder, less common grading categories rather than just riding on the easy majority class.

The target imbalanced setting, which more closely resembles what a hospital would face in practice since the deployment population is naturally skewed, tells a similar story. UMKD reached 91.75 percent overall accuracy and 90.72 percent mean accuracy, ahead of every baseline including the individually trained expert models themselves.

On APTOS fundus grading, a dataset the authors describe as more imbalanced and harder to annotate consistently than SICAPv2, UMKD again led on most metrics but the picture gets more textured, and this is where it pays to look past the headline numbers.

APTOS fundus grading, target imbalanced setting, percent except MAE
MethodOAMean AccF1MAE
ResNet18 Student baseline82.3470.5681.040.2521
RKD85.0069.7584.380.2482
KD (Hinton)83.4371.5683.500.2578
DKD81.8773.9582.430.2743
SDD83.1273.5582.830.2662
UMKD83.9174.3884.030.2476

Look closely and you will notice RKD actually beats UMKD on both overall accuracy, 85.00 against 83.91, and weighted F1, 84.38 against 84.03, in this specific target imbalanced APTOS setting. The paper is upfront about this rather than hiding it. Their explanation is that RKD clusters same class samples together within each training batch using angular relationships between sample triplets, which tends to favor whichever classes dominate a given batch. That inflates overall accuracy on an imbalanced test set precisely because the majority classes carry more weight in that metric. UMKD trades a little of that majority class accuracy for a meaningfully better mean accuracy, 74.38 against 69.75, which is the metric that actually reflects fair performance across all five DR severity grades rather than being dominated by the common ones. Whether that tradeoff is the right one depends entirely on what a deployment actually needs, and the paper does not pretend UMKD wins on every axis.

The absence of the SFA and CFA components leads to significant performance degradation in UMKD, because task agnostic structural features and task specific semantic features are not effectively decoupled. Section 3.3, Ablation Study

What the ablation study reveals that the abstract does not

The authors ran an ablation on SICAPv2, systematically removing one of the three components at a time. The results in the paper’s Table 3 rank the components by how much each one contributes.

Ablation on SICAPv2, source imbalanced setting, percent except MAE
ConfigurationOAMean AccF1MAE
No SFA, no CFA, no UDD (base)87.8286.6687.670.1594
SFA and CFA, no UDD90.6989.9290.580.1314
SFA and UDD, no CFA90.3689.6890.250.1355
CFA and UDD, no SFA88.1586.8487.970.1566
Full UMKD91.0290.2390.940.1294

Removing SFA hurts far more than removing either of the other two pieces. The CFA and UDD combination without SFA barely edges past the no distillation baseline at all, 88.15 against 87.82 overall accuracy, while dropping CFA still leaves the model at 90.36. That asymmetry is a genuinely useful finding for anyone thinking about adapting this approach, since it suggests the frequency domain structural alignment step is doing most of the heavy lifting, not the more elaborate uncertainty weighted output distillation that gives the method its name. The paper also states plainly that the same ablation was run on APTOS but not reported due to space limitations, so readers only get to see this component breakdown validated on one of the two datasets, which is a gap worth flagging rather than glossing over.

The clinical translation gap

Everything above describes performance on held out test splits of SICAPv2 and APTOS, both established public research datasets used for benchmarking, not on data collected from a live clinical workflow. That distinction matters more than it might seem. Public benchmark datasets are typically curated, meaning cases with ambiguous scans or conflicting expert labels are often excluded or resolved during dataset construction, and image acquisition conditions tend to be more uniform than what a hospital sees across different scanners, staining protocols, and patient populations over time.

A model that reaches 91 percent overall accuracy on SICAPv2 has not been shown to reach 91 percent accuracy on a new hospital’s scanner with its own staining chemistry, on a population with different demographics than the original cohort, or on the messier edge cases that get excluded from benchmark curation. The paper itself does not claim clinical deployment readiness, and the balanced subsets used for some experiments were built through random sampling specifically to create controlled test conditions rather than to represent a real patient intake stream. Any organization considering this kind of method for an actual grading tool would need prospective validation on data from the specific clinical setting it is meant to serve, ideally with comparison against expert clinician performance on the exact same cases, before any claim about real world accuracy would be appropriate.

Key takeaway. A benchmark result, even a strong one across multiple metrics and baselines, describes how a model performs on a specific curated dataset. It does not by itself describe how the model would perform inside a hospital’s actual diagnostic pipeline, and the paper does not claim otherwise.

Honest limitations

A few limitations are worth naming clearly, drawn from what the paper itself reports rather than speculation.

The sample sizes, while reasonable for this kind of research, are not enormous. SICAPv2’s balanced subset uses category counts as low as 948 samples for the rarest grade, and APTOS’s balanced subset goes as low as 193 for its smallest category. Deep learning models trained on datasets this size, even with a smaller ResNet18 student, can still overfit to dataset specific artifacts that would not generalize to new imaging equipment or new patient populations.

Both datasets carry their own known selection characteristics. SICAPv2 draws from a specific institutional pipeline for histology slide preparation, and APTOS was collected as part of a Kaggle competition with its own acquisition protocol. Neither is a multi institution, multi scanner dataset designed explicitly to test generalization across sites, which is generally considered the harder and more clinically relevant test.

The paper reports ablation results for only one of its two datasets, as noted above, which limits how confidently readers can generalize the finding that SFA matters most to the diabetic retinopathy task specifically. It is entirely possible the balance between components shifts for a different imaging modality.

Finally, and this applies to essentially every method paper of this kind, the work is a preprint that has not completed formal peer review at the time of writing. The reported numbers have not yet been independently verified through the peer review process, and readers should treat them as promising research findings rather than settled clinical facts.

Where this could matter beyond these two diseases

The core idea, using per region expert uncertainty to weight knowledge transfer rather than trusting every expert equally everywhere, is not tied to prostate cancer or diabetic retinopathy specifically. Any diagnostic imaging task with a similarly long tailed severity distribution, skin lesion grading, certain cardiac imaging classifications, or tumor staging from radiology scans, faces the same basic problem of experts trained on imbalanced data being unreliable exactly where it matters most. The privacy motivated design choice, keeping the original expert models frozen and never requiring access to their source training data, also fits naturally with how hospitals and health systems tend to guard patient data, which makes the underlying architecture worth watching even for teams working on entirely different organs or modalities.

Complete PyTorch implementation

Below is a full, runnable implementation of the three UMKD components described above, built with PyTorch and torchvision ResNet backbones. It includes the multi scale low pass filter for shallow feature alignment, the spherical projection for compact feature alignment, the uncertainty weighted decoupled distillation loss, the combined training loss, a training loop skeleton, an evaluation function, and a smoke test on random dummy data so you can confirm the pieces fit together before touching real patient data.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models

# ----------------------------------------------------------------------
# Multi scale low pass filter used by Shallow Feature Alignment (SFA)
# ----------------------------------------------------------------------
class MultiScaleLowPassFilter(nn.Module):
    def __init__(self, kernel_sizes=(2, 4, 8)):
        super().__init__()
        self.kernel_sizes = kernel_sizes

    def forward(self, feat):
        # feat shape N,C,H,W. Apply average pooling at several kernel sizes
        # then upsample back to the original resolution and average.
        _, _, h, w = feat.shape
        outputs = []
        for k in self.kernel_sizes:
            k = min(k, h, w)
            pooled = F.avg_pool2d(feat, kernel_size=k, stride=k)
            upsampled = F.interpolate(pooled, size=(h, w), mode='bilinear', align_corners=False)
            outputs.append(upsampled)
        return torch.stack(outputs, dim=0).mean(dim=0)


class StudentLowPassBranch(nn.Module):
    # Learnable counterpart used on the student side of SFA
    def __init__(self, channels, downsample_stride=2):
        super().__init__()
        self.msLF = MultiScaleLowPassFilter()
        self.downsample = nn.Conv2d(channels, channels, kernel_size=3,
                                     stride=downsample_stride, padding=1)
        self.upsample_match = nn.Upsample(scale_factor=downsample_stride, mode='bilinear', align_corners=False)
        self.dsconv = nn.Sequential(
            nn.Conv2d(channels * 2, channels, kernel_size=3, padding=1, groups=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, feat_s):
        low = self.msLF(feat_s)
        down = self.downsample(feat_s)
        down = self.upsample_match(down)
        down = F.interpolate(down, size=feat_s.shape[2:], mode='bilinear', align_corners=False)
        fused = torch.cat([down, low], dim=1)
        return self.dsconv(fused)


# ----------------------------------------------------------------------
# Feature alignment loss shared by SFA and CFA (MMD plus reconstruction)
# ----------------------------------------------------------------------
def mmd_loss(student_feat, expert_feats):
    # student_feat: N,C,H,W flattened to N,D
    # expert_feats: list of N,C,H,W tensors, one per expert
    s_flat = student_feat.flatten(1)
    total = 0.0
    for e_feat in expert_feats:
        e_flat = e_feat.flatten(1)
        diff = e_flat.sum(dim=0) - s_flat.sum(dim=0)
        total = total + torch.norm(diff, p=2) ** 2
    return total / s_flat.shape[0]


def reconstruction_loss(original_feats, aligned_feats):
    total = 0.0
    for orig, aligned in zip(original_feats, aligned_feats):
        total = total + F.mse_loss(aligned, orig, reduction='sum')
    return total


# ----------------------------------------------------------------------
# Compact Feature Alignment: project penultimate layer features into a
# shared spherical space with a 1x1 adapter per encoder
# ----------------------------------------------------------------------
class CompactFeatureAdapter(nn.Module):
    def __init__(self, in_channels, shared_dim=256):
        super().__init__()
        self.adapter = nn.Conv2d(in_channels, shared_dim, kernel_size=1)

    def forward(self, feat):
        projected = self.adapter(feat)
        pooled = F.adaptive_avg_pool2d(projected, 1).flatten(1)
        # project onto the unit sphere
        return F.normalize(pooled, p=2, dim=1)


# ----------------------------------------------------------------------
# Uncertainty aware Decoupled Distillation (UDD)
# ----------------------------------------------------------------------
def spatial_partition_logits(logits, w):
    # logits: N,C,H,W. Pool into a w by w grid of region logits.
    return F.adaptive_avg_pool2d(logits, output_size=(w, w))


def udd_loss(student_logits, expert_logits, scales=(1, 2, 4)):
    total_loss = 0.0
    for w in scales:
        psi_t = spatial_partition_logits(expert_logits, w)
        psi_s = spatial_partition_logits(student_logits, w)

        prob_t = F.softmax(psi_t, dim=1)
        uncertainty = 1.0 - prob_t.max(dim=1, keepdim=True).values

        # target class channel knowledge distillation term
        l_tckd = F.mse_loss(F.softmax(psi_s, dim=1), prob_t, reduction='none')
        l_tckd = l_tckd.mean(dim=1, keepdim=True)

        # non target class knowledge distillation term, raw logit alignment
        l_nckd = F.mse_loss(psi_s, psi_t, reduction='none').mean(dim=1, keepdim=True)

        weighted = (2.0 + uncertainty) * l_tckd + (1.0 - uncertainty) * l_nckd
        total_loss = total_loss + weighted.mean()
    return total_loss


# ----------------------------------------------------------------------
# Full UMKD model wrapper
# ----------------------------------------------------------------------
class UMKDModel(nn.Module):
    def __init__(self, num_classes=4, shared_dim=256):
        super().__init__()
        # student backbone, pretrained weights would be loaded in practice
        student = models.resnet18(weights=None)
        self.student_backbone = nn.Sequential(*list(student.children())[:-2])
        self.student_head = nn.Linear(512, num_classes)
        self.student_logit_conv = nn.Conv2d(512, num_classes, kernel_size=1)

        self.sfa_branch = StudentLowPassBranch(channels=512)
        self.cfa_adapter_student = CompactFeatureAdapter(512, shared_dim)

    def forward(self, x):
        feat_map = self.student_backbone(x)
        pooled = F.adaptive_avg_pool2d(feat_map, 1).flatten(1)
        logits = self.student_head(pooled)
        logit_map = self.student_logit_conv(feat_map)
        return {
            'logits': logits,
            'logit_map': logit_map,
            'feat_map': feat_map,
        }


# ----------------------------------------------------------------------
# Combined training loss
# ----------------------------------------------------------------------
def umkd_total_loss(student_out, expert_outs, labels,
                     cfa_adapter_student, cfa_adapters_experts,
                     sfa_branch, msLF,
                     alpha=0.5, beta=0.5):
    cls_loss = F.cross_entropy(student_out['logits'], labels)

    # shallow feature alignment
    student_low = sfa_branch(student_out['feat_map'])
    expert_lows = [msLF(e['feat_map']) for e in expert_outs]
    sfa_loss = mmd_loss(student_low, expert_lows) + reconstruction_loss(
        [e['feat_map'] for e in expert_outs], expert_lows)

    # compact feature alignment
    student_z = cfa_adapter_student(student_out['feat_map'])
    expert_zs = [adapter(e['feat_map']) for adapter, e in zip(cfa_adapters_experts, expert_outs)]
    cfa_loss = mmd_loss(student_z.unsqueeze(-1).unsqueeze(-1),
                         [z.unsqueeze(-1).unsqueeze(-1) for z in expert_zs])

    # uncertainty aware decoupled distillation, averaged across experts
    udd_total = 0.0
    for e in expert_outs:
        udd_total = udd_total + udd_loss(student_out['logit_map'], e['logit_map'])
    udd_total = udd_total / len(expert_outs)

    total = cls_loss + alpha * (sfa_loss + cfa_loss) + beta * udd_total
    return total, {
        'cls': cls_loss.item(), 'sfa': sfa_loss.item(),
        'cfa': cfa_loss.item(), 'udd': udd_total.item()
    }


# ----------------------------------------------------------------------
# Evaluation function
# ----------------------------------------------------------------------
def evaluate(model, dataloader, device='cpu'):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)
            out = model(images)
            preds = out['logits'].argmax(dim=1)
            correct = correct + (preds == labels).sum().item()
            total = total + labels.size(0)
    return correct / max(total, 1)


# ----------------------------------------------------------------------
# Smoke test on dummy data, confirms shapes and a backward pass work
# ----------------------------------------------------------------------
if __name__ == '__main__':
    torch.manual_seed(0)
    device = 'cpu'
    num_classes = 4
    batch_size = 2

    student = UMKDModel(num_classes=num_classes).to(device)
    cfa_student = CompactFeatureAdapter(512, 256).to(device)
    cfa_experts = [CompactFeatureAdapter(512, 256).to(device) for _ in range(2)]
    sfa_branch = StudentLowPassBranch(channels=512).to(device)
    msLF = MultiScaleLowPassFilter().to(device)

    dummy_images = torch.randn(batch_size, 3, 224, 224).to(device)
    dummy_labels = torch.randint(0, num_classes, (batch_size,)).to(device)

    student_out = student(dummy_images)

    # fake two frozen experts using the same architecture for the smoke test
    expert_outs = []
    for _ in range(2):
        expert = UMKDModel(num_classes=num_classes).to(device)
        for p in expert.parameters():
            p.requires_grad = False
        expert_outs.append(expert(dummy_images))

    loss, breakdown = umkd_total_loss(
        student_out, expert_outs, dummy_labels,
        cfa_student, cfa_experts, sfa_branch, msLF,
        alpha=0.5, beta=0.5
    )

    loss.backward()

    print('Smoke test passed')
    print('Total loss', loss.item())
    print('Loss breakdown', breakdown)
    assert student_out['logits'].shape == (batch_size, num_classes)
    print('Output shape check passed', student_out['logits'].shape)

Conclusion

UMKD tackles a specific and clinically grounded version of a familiar machine learning problem, what happens when you distill knowledge from teacher models that are themselves biased by imbalanced training data. Rather than treating each expert as an equally trustworthy source everywhere in the image, the method scores expert confidence at the regional level and lets that score control how strongly the student leans on each expert’s guidance. Combined with a two stage feature alignment process that separates general structure from disease specific signal, the approach delivers consistent gains across two genuinely different imaging modalities, histology slides for prostate cancer and fundus photographs for diabetic retinopathy.

The conceptual shift worth remembering is the move from treating distillation as a single global trust relationship between one teacher and one student, to treating it as a collection of local trust relationships that shift depending on where in the image and how confident the expert happens to be. That is a more honest way to model how imbalanced training actually damages a network, since the damage is never uniform. A model trained on scarce stage three prostate samples is not uniformly worse everywhere, it is specifically worse in the regions and patterns associated with that underrepresented class, and a distillation method that can localize its own uncertainty has a real shot at correcting for that rather than just averaging over it.

Whether this generalizes cleanly to other diagnostic imaging tasks with similar long tailed class distributions, skin lesion staging or certain radiology grading tasks come to mind, is an open question the authors do not claim to have answered. The architecture does not depend on anything specific to prostate tissue or retinal photographs, which is a reasonable basis for optimism, but that transferability has not itself been tested in this paper and would need its own validation.

The honest limitations matter as much as the strong numbers here. This is one preprint, tested on two public benchmark datasets with sample sizes in the low thousands, without prospective clinical validation, and without a full multi site generalization study. None of that erases the value of the underlying idea, but it does mean the responsible reading of this paper is as a promising research direction rather than a deployment ready clinical tool.

What happens next matters more than what has already been shown. If the uncertainty weighted distillation idea holds up under peer review and gets tested against a wider range of imaging tasks and a genuinely multi site dataset, it could become a useful building block for hospitals trying to combine several existing diagnostic models without either sharing raw patient data or accepting the biases each individual model picked up from its own training set. That is a real and recurring problem in clinical AI deployment, and this paper offers one concrete, tested piece of a possible answer.

Frequently asked questions

What does UMKD stand for and what problem does it solve

UMKD stands for Uncertainty aware Multi expert Knowledge Distillation. It trains one compact student model using guidance from two larger expert models, while specifically correcting for the fact that those experts were trained on imbalanced disease data and are therefore less reliable on rare grading categories.

What datasets and diseases were used to test this method

The authors tested on SICAPv2, a histology dataset for prostate cancer Gleason grading, and APTOS, a fundus photograph dataset for diabetic retinopathy severity grading.

Does UMKD beat every other distillation method on every metric

No, and the paper says so directly. In the APTOS target imbalanced setting, RKD achieves higher overall accuracy and weighted F1 than UMKD, though UMKD achieves a notably higher mean accuracy, the metric that better reflects fair performance across all severity grades rather than being dominated by common classes.

Has this method been used in an actual hospital

Not according to this paper. All results come from held out test splits of two public research datasets. The authors do not report prospective clinical validation, and the work is a preprint that had not completed formal peer review at the time this article was written.

Which component of UMKD matters most

The ablation study on SICAPv2 shows that removing shallow feature alignment causes the largest performance drop of the three components, more than removing compact feature alignment or the uncertainty weighted distillation step. The same ablation was not reported for APTOS.

Is this article medical advice

No. This article explains a published research paper for an AI and machine learning audience. It is not medical advice, and anyone with questions about a diagnosis or treatment should speak with a qualified clinician.

Read the original research

Read the paper on arXiv View the APTOS dataset

The full method, including the shallow and compact feature alignment mechanisms and the uncertainty aware decoupled distillation loss, is described in Tong, Gao, Liu, Huang, Xu, Ying, and Wu, Uncertainty Aware Multi Expert Knowledge Distillation for Imbalanced Disease Grading, arXiv:2505.00592, posted May 1 2025.

Academic citation. Tong, S., Gao, S., Liu, K., Huang, Z., Xu, H., Ying, H., Wu, J. Uncertainty Aware Multi Expert Knowledge Distillation for Imbalanced Disease Grading. arXiv preprint arXiv:2505.00592, 2025.

This analysis is based on the published paper and an independent evaluation of its claims.

Related reading

1 thought on “Uncertainty Aware Knowledge Distillation for Imbalanced Disease Grading”

Leave a Comment

Your email address will not be published. Required fields are marked *