Source-linked AI summary
Domain Generalization for Medical Imaging Classification with Linear-Dependency Regularization
Haoliang Li, YuFei Wang, Renjie Wan, Shiqi Wang, Tie-Qiang Li, Alex C. Kot
TL;DR
Medical imaging classifiers need large, representative datasets but may encounter limited data and shifts across vendors, protocols, and patient populations. The paper learns a representative latent feature space using variational encoding and linear-dependency regularization, and reports better cross-domain generalization on two medical imaging tasks.
Problem
Deep networks trained on limited, distribution-specific medical imaging data may not generalize to other domains, while representative training data are often difficult to access.
Method
The method combines variational encoding with rank-based linear-dependency regularization to capture shareable information in a representative latent feature space.
Results
Experiments on skin lesion classification and spinal cord gray matter segmentation show the proposed method's effectiveness and better cross-domain generalization than state-of-the-art baselines.
Takeaways & Limitations
The formulation provides a medical-imaging domain-generalization approach that jointly models latent linear dependency and domain alignment, with an empirical target-domain risk upper bound established theoretically.
Takeaways & Limitations
The reported clinical potential is framed for settings with limited training samples and diverse vendors or environments, while the method's broader diagnostic-system implications remain expected future possibilities.
Abstract
from arXiv · showhide
Recently, we have witnessed great progress in the field of medical imaging classification by adopting deep neural networks. However, the recent advanced models still require accessing sufficiently large and representative datasets for training, which is often unfeasible in clinically realistic environments. When trained on limited datasets, the deep neural network is lack of generalization capability, as the trained deep neural network on data within a certain distribution (e.g. the data captured by a certain device vendor or patient population) may not be able to generalize to the data with another distribution. In this paper, we introduce a simple but effective approach to improve the generalization capability of deep neural networks in the field of medical imaging classification. Motivated by the observation that the domain variability of the medical images is to some extent compact, we propose to learn a representative feature space through variational encoding with a novel linear-dependency regularization term to capture the shareable information among medical data collected from different domains. As a result, the trained neural network is expected to equip with better generalization capability to the "unseen" medical data. Experimental results on two challenging medical imaging classification tasks indicate that our method can achieve better cross-domain generalization capability compared with state-of-the-art baselines.
1 Introduction
Medical imaging models can fail under domain shifts and limited training data, while target-domain access is often infeasible. The paper therefore proposes learning shareable latent information across source domains for generalization to unseen data.
- Limited training data can weaken deep-learning generalization when testing images differ from training images in protocols, vendors, or patient populations.
- Target-domain data may be unavailable because real-time clinical settings limit collection and privacy regulations restrict access.
- Domain generalization trains models without target-domain information by using multiple source domains to generalize to an unseen testing domain.
- Existing medical-imaging domain-generalization work includes data augmentation, but selecting suitable augmentation types and magnitudes can be difficult for clinical deployment.
- The proposed method combines latent-space linear-dependency modeling with variational encoding and rank regularization to capture shareable information across domains.
2 Related Works
Related work addresses domain shift through adaptation, augmentation, synthesis, and domain-generalization methods. Domain generalization is more challenging because it must predict on an unseen target without accessing target-domain data.
- Domain adaptation methods address source–target shifts through subspace learning, instance re-weighting, distribution alignment, or adversarial learning.
- Domain generalization is more challenging than domain adaptation because the target domain is unavailable during training.
- Prior domain-generalization methods use multiple source domains to learn shareable or invariant information for unseen-target prediction.
- Medical-imaging approaches to cross-domain classification include image synthesis with generative adversarial networks to mitigate distribution shifts.
3 Methodology
LDDG learns a latent representation that captures cross-domain, within-category linear dependency while aligning latent features to a Gaussian prior. It combines rank and distribution regularization with classification to support generalization to unseen but related domains.
- Linear-Dependency Modeling: LDDG models shareable information across domains by constraining latent features to rank C, the number of task categories.The method assumes linear dependency among same-category latent features from multiple source domains.
- Linear-Dependency Modeling: Rank(Z) = C is optimized by minimizing the (C + 1)th singular value of the latent-feature matrix.This provides the rank loss and its sub-gradient for optimization.
- Distribution Alignment: The method adds KL regularization to match source-domain latent features to the Gaussian prior N(0, 1).The combined linear-dependency and distribution regularization is described as bounding empirical risk on unseen but related domains.
- Theoretical Analysis: Under the stated assumptions, unseen related-domain latent features lie on the prior manifold, and target expected loss is bounded when source empirical risks are small.Theorem 2 gives the bound Mϵ + log C for a linear softmax classifier trained with cross-entropy.
- Model Training: The network maps samples through feature extraction, variational encoding, and a linear classification module to produce predictions.The architecture comprises Qθ, Fω, and Tφ; latent features are resampled before classification.
- Model Training: The training objective combines cross-entropy, rank, and distribution regularization to penalize prediction error, latent-distribution mismatch, and latent rank.The classifier uses a linear module so the stated theoretical assumptions apply.
4 Experiments
Experiments evaluate the method on skin lesion classification and spinal cord gray matter segmentation across domains. The method generally outperforms domain-generalization baselines, with ablations supporting category-specific rank regularization.
- Experimental setup: The evaluation covers skin lesion classification and spinal cord gray matter segmentation across multiple medical domains.Skin lesion experiments use seven public datasets; segmentation data come from four medical centers with different MRI systems.
- Skin lesion classification: Domain-generalization methods outperform DeepAll on skin lesion classification, while the proposed method achieves a clear margin over the baselines.The comparison includes MASF, MLDG, and CCSA, with DeepAll trained directly using classification loss.
- Spinal cord gray matter segmentation: The proposed method achieves the best segmentation results for target sites 1, 2, and 3 across all metrics, and for site 4 under ASD.MASF performs better than CCSA and MLDG among the other domain-generalization methods.
- Spinal cord gray matter segmentation: Qualitative results show competing methods failing to segment or over-segment gray matter on unseen sites, whereas the proposed method generally performs better.Examples include failure on site2 and extension into white matter on site3.
- Ablation study: Both rank regularization and KL-based distribution alignment benefit generalization, while category-specific rank regularization outperforms nuclear-norm minimization.The ablation study defines Rank, KL, and LR as the corresponding components or alternatives.
- Ablation study: Accuracy peaks at rank(Z) = 7, matching the task’s seven categories, and decreases when rank(Z) becomes larger.The authors associate larger ranks with noise information and report better performance than nuclear-norm regularization.
5 Conclusion
The paper addresses medical-imaging generalization by combining variational encoding with linear-dependency modeling and domain alignment, reporting effectiveness on two tasks while noting fairness risks from protected-data constraints.
- The method learns a representative feature space through variational encoding with a novel rank regularization term for linear-dependency modeling.Its formulation is intended to capture shareable information across domains.
- The approach combines linear-dependency modeling and domain alignment to address generalization in medical imaging classification.
- Experimental results on skin lesion classification and spinal cord gray matter segmentation show the effectiveness of the proposed method.
- Privacy restrictions on protected attributes may leave trained models vulnerable to unfairness or discrimination.The paper identifies gender and ethnicity as examples of protected attributes that may not be publicly released for training.
A.1 Experimental Setting for Skin Lesion Classification Task
For skin lesion classification, the experiments use a pretrained ResNet18 feature extractor, variational encoding, and a seven-output fully connected classifier, trained with specified regularization and optimization settings.
- The model uses an ImageNet-pretrained ResNet18 without its fully connected layer as the feature extractor, with input size 224 × 224.For the proposed method, the network before average pooling serves as the feature extractor.
- A variational encoding network is inserted between the feature extractor and the final fully connected classification layer.
- The classification network is a fully connected layer with output size 7.
- Training uses λ1 = 0.001, λ2 = 0.4, focal loss, Adam, learning rate 0.0001, weight decay 0.001, batch size 32, and 200 epochs.The learning rate decreases by a factor of 10 after every 80 epochs, and evaluation uses the best validation model.
A.2 Experimental Setting for Spinal Cord Gray Matter Segmentation Task
For spinal cord gray matter segmentation, the method uses a 2D-UNet backbone with a variational encoder and weighted binary cross-entropy, trained separately across domains with stated hyperparameters.
- The backbone is a 2D-UNet without its final 1 × 1 convolutional layer, using MRI axial slices as input.
- The variational encoder uses two identical networks, a 64-channel latent layer, ReLU activation, and 8-channel mean and standard-deviation outputs.
- Training uses λ1 = 0.001, λ2 = 0.01, weighted binary cross-entropy, Adam, learning rate 1e-4, weight decay 1e-8, and batch size 8 per domain.Models are trained for 200 epochs with tenfold learning-rate reductions every 80 epochs.
B BigAug [42] for Segmentation
BigAug applies stacked image transformations for segmentation, but default parameters can damage image quality and boundaries; tuning helps, while the proposed method still performs better.
- BigAug stacks transformations including sharpness, blurriness, noise, brightness, contrast, rotation, and scaling for spinal cord gray matter segmentation.
- Default BigAug parameters produce undesired segmentation performance, worse than the DeepAll baseline in DSC and JI.
- Default augmentation can deteriorate input quality and oversmooth boundaries, increasing discrepancy between source and target domains.
- Parameter tuning improves BigAug over DeepAll, but the proposed method still outperforms tuned BigAug.
- Table 5 reports domain-generalization results for gray matter segmentation using BigAug.