Source-linked AI summary
Extending the Horizon of Early Diagnosis: Lung Cancer Prediction with Vision Transformers
Olivera Kotevska, Ian Goethert, Michael McGee, Maria Mahbub, Sean R. Wilkinson, Rowena Yip, Myvizhi Esai Selvan, Zeynep H. Gumus, Claudia Henschke, Robert J. Klein, Providencia Morales, Samuel M Aguayo, Ioana Danciu, Mayanka Chandrashekar
TL;DR
Early lung cancer signs can be difficult to detect on routine chest X-rays, motivating prediction before clinical diagnosis. This study evaluates ViTs using different pretraining and imbalance-mitigation strategies, finding improved performance with transfer learning but results below clinical thresholds.
Problem
Early-stage lung cancer can be difficult to detect on chest X-rays, while existing methods have limited ability to forecast diagnosis one or two years in advance.
Method
The study fine-tuned scratch, ImageNet-pretrained, and Corona-pretrained ViTs on routine-care chest X-rays labeled by registry diagnosis within one and two years.
Results
Pretrained models surpassed the scratch baseline by 6–10 percentage points in AUC and approximately 10–12% in balanced accuracy, while ImageNet models were more stable.
Takeaways & Limitations
The findings support further development of ViT-based triage systems that flag high-risk patients for earlier evaluation.
Takeaways & Limitations
Performance remained below clinical acceptability thresholds, and the study lacked external validation from independent cohorts.
Abstract
from arXiv · showhide
Lung cancer remains a leading cause of cancer-related mortality worldwide, and early diagnosis is critical for improving survival. However, early-stage malignancies can be subtle on chest X-rays, creating challenges for radiologists. This study evaluates Vision Transformers (ViTs) for predicting lung cancer one to two years before clinical diagnosis. We analyzed 259,361 chest X-rays from 91,020 imaging studies at the Jamaica Plains VA Hospital in Boston, MA. The dataset showed extreme class imbalance, approximately 1:150 cancer to non-cancer, which was addressed using hybrid under- and over-sampling and class-weighted loss optimization. Three ViT configurations were evaluated: a model trained from scratch, an ImageNet-pretrained model, and a Corona-pretrained model fine-tuned on the lung cancer dataset. Transfer learning improved performance, with pretrained models exceeding the scratch baseline by 6-10 percentage points in AUC and about 10-12 percent in balanced accuracy. ImageNet-pretrained models showed the most stable overall performance, while Corona-pretrained models achieved higher sensitivity in some settings but greater variability. Moderate resampling ratios, including 1:1 undersampling and 1.5:2 oversampling, provided favorable trade-offs between sensitivity, precision, and computational efficiency, reducing runtime by up to 70 percent without major performance loss. These findings demonstrate the potential of ViTs for early lung cancer risk prediction from routine chest X-rays. Although performance remains below clinical deployment thresholds, the results support further development of ViT-based triage systems to flag high-risk patients for earlier evaluation.
1 Introduction
Lung cancer is often diagnosed late, while early radiographic signs can be subtle and difficult to detect. This study evaluates ViT-based prediction of lung cancer diagnosis one to two years after routine chest X-ray imaging.
- Motivation: Lung cancer causes millions of deaths annually, and fewer than 16% of eligible U.S. individuals undergo annual low-dose CT screening.Limited screening uptake contributes to continued advanced-stage diagnosis.
- Research gap: Existing methods have limited ability to forecast cancer development years in advance, motivating earlier risk identification.The study targets a defined one- or two-year interval before diagnosis.
- Clinical context: Routine chest radiographs contain subtle pulmonary and mediastinal changes that may precede cancer diagnosis but are often imperceptible to human observers.AI-based interpretation could support opportunistic identification of high-risk individuals.
- Prior work: ViTs use high-resolution image processing and pretrained models for medical vision tasks, including cancer detection.Prior work reports applications in chest X-rays, retinal images, and brain MRIs.
- Study objective: The study fine-tunes pretrained ViTs on more than 150,000 VHA chest radiographs to predict cancer diagnosis within one or two years of imaging.It assesses pretraining, class-imbalance mitigation, and sampling strategies.
2 Methods
The study uses routine-care chest X-rays labeled by subsequent registry diagnosis, with patient-level separation and standardized preprocessing for ViT inputs.
- Data: Chest X-rays from Boston VHA routine care between 2010 and 2022 were labeled by registry lung cancer diagnosis within one and two years after imaging.Binary indicators were computed for each image.
- Data: Patient-level splits created non-overlapping training and testing cohorts, with the same test set used across experiments.This design supports unbiased comparison between modeling conditions.
- Preprocessing: The preprocessing pipeline converted images, resized and normalized them, and transformed them into tensors for ViT processing.The workflow standardized image format and scale before batch processing.
- Preprocessing: DICOM images were converted to JPEG to simplify processing and maintain compatibility with preprocessing tools.The conversion was intended to retain image quality and diagnostic information.
- Preprocessing: Images were resized to 224 × 224 pixels to match the ViT patch-based architecture.Pixel intensities were then normalized to the pretrained model’s expected input distribution.
- Preprocessing: Rescaled intensities were standardized using ImageNet mean µ and standard deviation σ values across three replicated channels.The applied values were µ = [0.485, 0.456, 0.406] and σ = [0.229, 0.224, 0.225].
2.2 Experimental setup
Experiments were run on a single NVIDIA A100 GPU using Python, PyTorch, and SLURM under a fixed computational configuration.
- Computational setup: Training used one NVIDIA A100 GPU with 80 GB memory on a DGX node.Jobs used one node, one GPU, 60 GB memory, and a 20-hour wall-time limit.
2.3 Vision Transformer (ViT)
ViTs represent chest X-rays as patch sequences and use self-attention to model global spatial context. A classification token is mapped to cancer and no-cancer probabilities.
- Architecture: Unlike CNNs, ViTs treat images as sequences of fixed-size patches, enabling long-range dependency and global spatial-context modeling.This architecture adapts transformers originally developed for natural language processing to image data.
- Image tokenization and embedding: Images are divided into non-overlapping patches that are flattened and linearly projected into embeddings with positional encodings.A learnable classification token aggregates global information.
- Transformer encoder: Transformer encoder layers combine multi-head self-attention, feed-forward networks, residual connections, and layer normalization.These components process the sequence of patch embeddings.
- Classification head: The final classification-token representation enters an MLP that outputs cancer-versus-no-cancer probabilities.The model therefore performs binary medical image classification.
- Architecture: The ViT formulation is intended to learn localized and global radiographic patterns relevant to subtle early-stage lung cancer detection.Figure 1 summarizes the patch-tokenization, transformer-encoder, and binary-output workflow.
2.4 Workflow methodology
The workflow compares non-pretrained and pretrained ViT configurations for lung cancer imaging, including ImageNet and Corona-based transfer learning. All configurations use ViT models fine-tuned or trained for binary classification on the lung cancer dataset.
- Experimental variants: The from-scratch model used a lightweight ViT architecture with 4×4 patches, two transformer blocks, two attention heads, hidden dimension 8, and 10 output classes.It was trained for 20 epochs with Adam, batch size 32, CrossEntropyLoss, and weighted random sampling.
- Pretrained models: The pretrained models used ViT-Base-Patch16-224 architectures configured for binary classification and fine-tuned with AdamW and class-weighted CrossEntropyLoss.Inputs were resized to 224 × 224 RGB, normalized, sampled by class frequency, and trained for 20 epochs with three repetitions.
- Pretrained models: The ImageNet model used google/vit-base-patch16-224, while the Corona model used identical architecture with weights further fine-tuned on chest radiographs.The CoronaHack–ChestXRay data included healthy, pneumonia, and COVID-19 cases, providing test domain-specific transfer learning.
- Experimental variants: Three ViT workflows were evaluated: training from scratch, ImageNet pretraining, and two-stage Corona-based pretraining followed by lung cancer fine-tuning.The Corona configuration was first adapted on the CoronaHack–ChestXRay dataset before final fine-tuning on sensitive lung cancer data.
2.5 Handling the unbalanced classes
The study addressed severe cancer-to-non-cancer imbalance through controlled sampling and class-weighted loss reweighting. Undersampling reduced majority-class representation, while oversampling increased minority-class representation within predefined ratios.
- Imbalance mitigation: > 1 : 100 cancer-to-non-cancer imbalance was addressed using controlled sampling and loss reweighting.Both strategies were designed to reduce imbalance effects during model training.
- Controlled sampling: Undersampling reduced majority-class samples to predefined ratios from 1:1 to 1:4 while limiting excessive data loss.The approach ensured that both classes contributed comparably during training.
- Controlled sampling: Oversampling replicated or augmented minority-class samples to reach ratios between 1:1 and 1.5:2 while limiting overfitting risk from excessive duplication.Experiments sought trade-offs between robustness and computational cost.
- Loss reweighting: Class-weighted loss assigned inverse-frequency weights so errors from underrepresented classes contributed more strongly during optimization.This encouraged balanced attention across classes and reduced majority-class bias.
2.6 Experiments and Analysis
Experiments used patient-level disjoint data splits, a fixed hold-out test set, repeated random seeds, and validation balanced accuracy for checkpoint selection. Performance was reported using mean and standard deviation unless otherwise stated.
- Experimental design: Patient-level disjoint splits prevented any patient from appearing in more than one data split.A fixed hold-out test set was identical across experiments, while model development used training and validation portions only.
- Experimental design: Each experiment was repeated three times with different random seeds, and the best checkpoint was selected by validation balanced accuracy.Reported results used the mean and standard deviation unless otherwise stated.
- Experimental design: The evaluation covered ViT models fine-tuned for lung cancer detection across different class ratios and pretraining sources.
2.7 Metrics
AUC and balanced accuracy were selected to evaluate classification performance under severe class imbalance. AUC summarizes class separability across thresholds, while balanced accuracy averages recall across classes; execution time measures elapsed training-to-testing duration.
- Metric selection: AUC and balanced accuracy were chosen because they address imbalanced medical-imaging datasets.
- AUC: AUC summarizes the model’s ability to distinguish classes across all classification thresholds using a single scalar value.It captures the trade-off between true-positive rate and true-negative rate and is less biased by class distribution.
- Balanced accuracy: Balanced accuracy averages recall, or sensitivity, for each class to prevent minority-sample performance from being masked by class imbalance.
- Runtime: Execution time was defined as the elapsed wall-clock duration between training start and testing end.
3 Results
Results show that pretrained Vision Transformers generally outperform the non-pretrained baseline, with ImageNet initialization providing stronger overall performance and Corona initialization showing competitive sensitivity. Across sampling configurations, moderate ratios improved efficiency while preserving performance.
- Dataset: 259,361 chest X-rays from 91,020 imaging studies formed the Boston VHA dataset, with substantially more no-cancer than cancer cases.The dataset included multiple clinical views, primarily PA and LL, with fewer AP images.
- Pretrained model evaluation across sampling techniques: ImageNet-initialized models mostly outperformed Corona-initialized models in Balanced Accuracy and AUC, although Corona models achieved 59.02–62.27% balanced accuracy across the whole dataset.The comparison supports stronger overall performance from ImageNet pretraining while retaining competitive Corona results.
- Pretrained model evaluation across sampling techniques: 6–10 percentage points: ImageNet-based fine-tuning surpassed the non-pretrained model in AUC.ImageNet-based fine-tuning achieved 52.23–61.30% balanced accuracy, while the non-pretrained baseline achieved 49.77–51.75%.
- Visual Performance Comparison: ImageNet models improved Sensitivity from ≈0.08 to ≈0.49 and F1-score from ≈0.15 to ≈0.56 between Year 1 and Year 2, while Specificity declined from ≈0.92 to ≈0.51.Balanced Accuracy remained stable, indicating improved recall without severe distortion of global performance.
- Pretrained model evaluation across sampling techniques: 25–70% reduction in runtime: under-sampling and over-sampling reduced training time to 14–18 hours without noticeable performance degradation.The reported favorable ratios included 1:1 under-sampling and 1.5:2 over-sampling.
4 Discussion
The study finds that pretrained ViTs and carefully controlled resampling improve early lung-cancer risk prediction from routine chest radiographs, but performance remains below diagnostic-use thresholds. The results support triage-oriented risk stratification while highlighting limitations in signal strength, generalizability, and clinical integration.
- Clinical implications: Chest-radiograph signals were detectable one to two years before diagnosis but insufficient for prediction alone and below diagnostic-use thresholds.The authors therefore frame contextual integration and triage as more appropriate than isolated image-based diagnosis.
- Transfer learning: 6–10 percentage points in AUC and approximately 10–12% in balanced accuracy separated pretrained models from the untrained baseline.Pretrained models likely benefit from transferable visual features captured during large-scale pretraining.
- Class imbalance: A ≈1:150 cancer-to-non-cancer ratio caused majority-class overfitting, while moderate 1:1 or 1.5:2 resampling improved sensitivity and F1-score without major precision loss.Heavy 1:3 resampling instead produced instability and degraded generalization.
- Pretraining comparisons: ImageNet-pretrained models were more stable overall, whereas Corona-pretrained models sometimes achieved higher sensitivity but showed greater variability across metrics.The comparison indicates a trade-off between generalization and domain proximity.
- Task difficulty: COVID pneumonia exceeded 90% accuracy after a few training steps, whereas lung-cancer performance plateaued around 60–65% balanced accuracy under the same configuration.Diffuse pneumonia opacities are easier to distinguish than subtle, localized cancer lesions confounded by anatomical variation.
- Computational efficiency: Up to 70% training-time reductions from balanced subsampling occurred without significant accuracy loss.The full dataset required two days per experiment, making sampling relevant to resource-constrained retraining and domain adaptation.
- Limitations: The study remains an initial step because it used one VHA site, lacked external validation, and did not model longitudinal imaging, calibration, or explainability.These constraints motivate multi-institutional, longitudinal, and clinically oriented validation.
5 Conclusion
This work evaluates ViT models for early lung-cancer prediction from chest X-rays, examining pretrained initialization, class-imbalance mitigation, and domain-specific transfer learning. Pretraining and moderate sampling improved performance, but results remained below clinical thresholds, supporting further development of multimodal and federated approaches.
- Model performance: Both ImageNet- and Corona-pretrained ViTs outperformed untrained baselines on the lung-cancer imaging task.ImageNet pretraining produced more stable and balanced performance, while Corona pretraining showed higher sensitivity but greater variability.
- Data balance: Moderate sampling strategies improved recall without major accuracy loss despite severe class imbalance.The conclusion emphasizes sampling as an effective mitigation strategy rather than unrestricted resampling.
- Clinical scope: Overall performance remained below clinical thresholds because subtle cancer lesions are difficult to identify from X-rays.The authors identify multimodal and federated approaches as directions for improving reliability and efficiency.