Med-LEGO: Editing and Adapting toward Generalist Medical Image Diagnosis

🕒 Published (v1): 2025-03-03 04:27 UTC · Source: Arxiv · Venue: MICCAI · link

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Med-LEGO is a training-free framework that merges multiple specialist medical imaging models into a generalist CAD model by combining SVD-enhanced LoRA adapters (SVD-LoRA) through simple averaging and re-decomposition. It eliminates the need for original training data during fusion, addressing privacy barriers to building generalist medical AI. Using only 0.18% of full model parameters, it outperforms existing model-merging baselines across cross-domain and in-domain benchmarks.

Problem

Privacy constraints prevent centralizing multi-domain medical data for pre-training or continually updating a single foundation model, forcing practitioners to maintain isolated specialist models. Existing model-merging methods (Task Arithmetic, Ties-Merging, PEM Composition, etc.) operate on full model parameters—expensive to compute and store—and suffer from task-interference noise when directly adding or subtracting weights across heterogeneous medical domains with large distribution gaps.

Method

Each specialist model is fine-tuned with SVD-LoRA: the weight update is parameterized as ∆ = BEA, where B ∈ ℝ^{d_m×r} and A ∈ ℝ^{r×d_n} are orthonormal singular-vector matrices (enforced by a Frobenius-norm regularization term) and E ∈ ℝ^{r×r} is a diagonal singular-value matrix initialized to zero. This structure is applied to Q and V projection matrices of ViT-base-patch16 with r = 4.

To merge N specialist adapters, the full-rank average ∆_full = (1/N)∑ B_i E_i A_i is computed, then re-decomposed via SVD. The top-k singular components whose cumulative singular values exceed threshold v = 99.7% are retained to form the merged SVD-LoRA. This re-decomposition recovers the low-rank structure, avoids the pre-merge vs. post-merge inequality inherent to standard LoRA averaging, and suppresses noise from less informative singular directions. Cross-domain merging yields a "General Doctor"; in-domain merging (e.g., multiple CXR datasets) yields a domain-specialized doctor. The merged weights then serve as initialization for fine-tuning on new tasks.

Key Contributions

• SVD-LoRA adapter: augments standard LoRA with an explicit diagonal singular-value matrix E, enabling mathematically sound post-merge SVD re-decomposition and orthogonality-regularized training.
• Training-free merging: specialist models are fused without access to any original training data, directly addressing medical data privacy constraints.
• Noise suppression via rank truncation: retaining only top-k singular components filters task-interference noise automatically, eliminating manual sign/scale tuning required by prior methods.
• Better initialization for new tasks: merged weights provide faster convergence on unseen medical tasks than ImageNet pre-trained ViT weights, effectively bridging natural-to-medical domain gap.

Results

Cross-domain merging (7 MedMNIST datasets, single merged model): - Med-LEGO achieves best or second-best accuracy across most tasks; representative scores: Blood 0.814, Breast 0.872, Pneumonia 0.849, Retina 0.593. - Competing methods collapse on several tasks (e.g., Task Arithmetic: Blood 0.078, Derma 0.201; MagMax: Derma 0.052, Pathology 0.101).

In-domain merging (3 chest X-ray datasets): - Med-LEGO: Pneumonia ACC 0.845 / AUC 0.966; Tuberculosis ACC 0.803 / AUC 0.897; NIH-CXR14 ACC 0.382 / AUC 0.710. - Baseline collapse on Pneumonia: ViT Averaging ACC 0.354, Task Arithmetic 0.372, Ties-Merging 0.330. - PEM Composition is the closest competitor but degrades on NIH-CXR14 (AUC 0.579 vs. Med-LEGO 0.710).

Generalization to new tasks (OAI, Blood-Cell, Tuberculosis): - Med-LEGO merged weights converge faster and reach higher test accuracy than both ImageNet ViT and ViT Averaging across all three new datasets (exact values shown in Fig. 2 convergence curves; no single-number tabulation provided).

Limitations

• Backbone restricted to ViT-base-patch16 with ImageNet pre-training; generalization to larger or domain-pre-trained ViTs (e.g., medical foundation models) is not evaluated.
• NIH-CXR14 multi-label performance remains low (AUC 0.710 merged), only marginally below specialist AUC 0.783, suggesting multi-label heterogeneity is not fully resolved by merging.
• Only classification tasks are tested; segmentation and detection are not addressed.
• Threshold v = 99.7% for rank selection is fixed; sensitivity analysis is not reported.
• Orthogonality regularization adds a hyperparameter and training overhead not required by plain LoRA.
• Small dataset sizes (e.g., Tuberculosis: ~660 images) limit conclusions about scalability.

Relevance to Foundation Models in Medicine

Med-LEGO directly targets the core bottleneck in medical foundation model deployment: the inability to aggregate multi-institutional, multi-modality data due to privacy regulations. By enabling training-free, data-free composition of specialist adapters into a generalist model, it provides a practical federated-style alternative to centralized pre-training. The SVD-LoRA representation is parameter-minimal (0.18% of model weights), making it compatible with communication-constrained or storage-constrained clinical settings. This work connects to a growing line of research on model merging and parameter-efficient adaptation for medical AI, and the resulting merged weights' superior initialization for downstream fine-tuning suggests a pathway toward iteratively building richer medical foundation models without data sharing.