Skip to content

Deep Reprogramming Distillation for Medical Foundation Models

🕒 Published (v1): 2026-05-06 03:22 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Deep Reprogramming Distillation (DRD) is a framework for adapting large medical foundation models to resource-constrained downstream tasks by transferring knowledge into lightweight student models while jointly resolving task/domain, structural, and training-strategy inconsistencies. It introduces learnable convolutional reprogramming blocks between teacher and student, co-trained end-to-end on downstream data, plus a Centered Kernel Alignment (CKA) distillation loss for robust feature alignment. DRD outperforms prior KD and PEFT methods across 18 medical datasets covering 2D/3D classification and segmentation.

Problem

Standard knowledge distillation (KD) assumes teacher and student share the same task, model family, and training strategy — a condition violated when a large self- or weakly-supervised medical foundation model is the teacher and a heterogeneous lightweight CNN is the student. PEFT methods reduce training cost but do not reduce inference-time memory or latency. Naïve PEFT+KD composition ignores structural and strategy mismatch, empirically causing negative transfer (e.g., up to −11.8pp on FIVES when distilling ViT→CNN without DRD).

Method

Reprogramming Module. For each of \(N\) coarse stages, a lightweight three-layer convolutional projector \(\phi_i\) maps teacher feature maps \(f_i^T \in \mathbb{R}^{C\times H\times W}\) to student-compatible dimensions: \(f_i^{T\to S} = \phi_i(f_i^T)\).

Deep Co-training Reprogramming. The reprogrammed feature from stage \(i\) is injected into the student at block \(B_{i+1}^S\), propagating through the remaining student blocks to produce hybrid logits \(z_i^{T\to S}\). Loss is computed on both hybrid and pure-student logits against ground truth (CE/Dice), plus KL divergence \(L_{KD} = \sum_i \text{KL}(z_i^{T\to S}, z^S)\), jointly optimizing \(\phi_i\) and the student end-to-end.

CKA Distillation. Feature-level alignment uses the normalized HSIC-based CKA loss: $\(L_i^{CKA} = -\frac{\text{HSIC}(K_i, L_i)}{\sqrt{\text{HSIC}(K_i,K_i)\cdot\text{HSIC}(L_i,L_i)}}\)$ where \(K_i = f_i^{T\to S}(f_i^{T\to S})^\top\) and \(L_i = f_i^S(f_i^S)^\top\). CKA is invariant to orthogonal transforms and robust to random-seed/width variation.

Total loss: \(L_\text{train} = L_\text{sup} + \alpha L_\text{hybrid} + \beta L_{KD} + L_{CKA}\), with \(\alpha,\beta\) decayed linearly to relax reprogramming constraints as training progresses.

Key Contributions

  • DRD framework: single-stage, one-framework solution addressing task/domain gap, teacher–student structural heterogeneity, and lightweight deployment simultaneously.
  • Deep Co-training Reprogramming: propagates reprogrammed features through student branches so hybrid logits jointly supervise both teacher reprogramming and student learning.
  • CKA distillation: geometry-aware, scale-invariant feature alignment that improves training stability under diverse hyperparameter/architecture conditions.
  • Comprehensive evaluation: 8 medical foundation models × 18 public datasets × 2D/3D × classification and segmentation; consistent improvement over 9 KD baselines and 5 PEFT baselines.

Results

  • 2D classification (BUSI): DRD reaches 88.31% accuracy (PMC-CLIP teacher) vs. vanilla 77.92% and best KD baseline 83.12%; ResNet18 student matches teacher performance.
  • 2D segmentation (FIVES): DRD ShuffleNet 73.06 Dice (MedSAM teacher) vs. vanilla 56.25 (+16.8pp); next-best CIRKD 66.97.
  • 2D segmentation (CDPRD, ViT-Tiny): DRD 95.33 vs. vanilla 90.91, teacher 78.34 — student surpasses teacher.
  • 3D segmentation (BTCV-Gallbladder, SAT teacher): DRD 43.80 vs. vanilla 19.15, best KD (Hint) 27.23.
  • 3D segmentation (MSD-Spleen): DRD 89.83 vs. vanilla 80.87, SAT teacher 87.21.
  • 3D classification (LAC, DenseNet3D): DRD 69.71 vs. vanilla 63.43, Merlin teacher 77.92.
  • DRD is the only method that covers efficient deployment + one-stage training + both classification and segmentation + heterogeneous architectures simultaneously (Table III).

Limitations

  • Requires a fixed number of coarse stages to be shared between teacher and student; architecturally very dissimilar models (e.g., ViT vs. lightweight CNN with fewer stages) need manual partitioning.
  • CKA is computed over full batch Gram matrices, scaling \(O(n^2)\) in batch size — limits batch size in memory-constrained 3D settings (batch 4 used).
  • Foundation model backbone is fully frozen during adaptation; no mechanism to selectively unfreeze for severely out-of-domain tasks.
  • Requires downstream labeled data; does not address the zero-shot or few-shot regime.
  • Evaluation is restricted to image-based foundation models; text-vision or report-grounded models (e.g., BioViL, LLaVA-Med) are not covered.

Relevance to Foundation Models in Medicine

DRD directly addresses the deployment gap between large-scale pre-trained medical foundation models and clinical edge-device constraints — a central bottleneck in translating models like MedSAM or SAT into practice. By framing efficient deployment as a structural and task alignment problem rather than a compression problem, it provides a principled alternative to PEFT-only or KD-only pipelines. The breadth of evaluation (8 foundation models, 4 modalities, 2D/3D, classification + segmentation) provides one of the more systematic analyses of how foundation model adaptation behaves under heterogeneous student architectures, which is directly relevant to building lightweight clinical AI systems from general-purpose medical priors.