Extending the Horizon of Early Diagnosis: Lung Cancer Prediction with Vision Transformers
Authors: Olivera Kotevska, Ian Goethert, Michael McGee, Maria Mahbub, Sean R. Wilkinson, Rowena Yip, Myvizhi Esai Selvan, Zeynep H. Gumus, +6 more
Organizations: Mathematics in Computation, Computer Science and Mathematics Division, Oak Ridge National Laboratory · Research Computing Support Division, Oak Ridge National Laboratory · Advanced Intelligent Systems, Cyber Resilience and Intelligence Division, Oak Ridge National Laboratory · National Center for Computational Sciences, Oak Ridge National Laboratory · Icahn School of Medicine at Mount Sinai · Phoenix VA Medical Center · Department of Biomedical Informatics, Vanderbilt University Medical Center · Advanced Computing for Health Sciences, Computational Sciences and Engineering Division, Oak Ridge National Laboratory
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.
Figures & tables
Figure 1: Vision Transformer (ViT) workflow for binary medical image classification. The model tokenizes the input image into patches, processes them through transformer encoders with self-attention, and outputs two classes: Cancer and No-Cancer .
Figure 2: Overview of the experimental workflows for ViT fine-tuning on the lung cancer dataset. The top pathway represents the baseline configuration trained directly on the sensitive dataset without pretraining. The middle pathway illustrates the ImageNet-pretrained model fine-tuned on the lung cancer dataset. The bottom pathway shows the two-stage fine-tuning strategy, where the model is first adapted using an open-source COVID pneumonia dataset before final fine-tuning on the sensitive lung cancer data. All configurations are evaluated using the same testing dataset to ensure consistent performance comparison.
Year 1 Dataset
Year 2 Dataset
# Number of Patients
28,637
28,281
# Number of Studies
63,818
63,322
# Chest X-ray Images
169,001
167,684
Category-wise # Images:
Class 0: No-cancer diagnosed
167,936
166,124
Viewposition-wise
Table 1: Summary statistics for Year 1 and Year 2 datasets. It shows the number of samples per class, study and patient ID, and different view positions such as LL, PA, and AP.
Sampling
Class Ratio
Year 1
Year 2
VIT-ImageNet
VIT-Corona
VIT-ImageNet
VIT-Corona
-
Original Dataset
59.24 ± 0.85
59.70 ± 1.43
57.78 ± 1.12
58.83 ± 1.05
Under sampling
U: 1:1
64.07 ± 1.67
61.51 ± 1.10
60.79 ± 0.94
57.15 ± 1.42
U: 1:2
58.18 ± 1.32
58.23 ± 0.93
58.29 ± 1.25
58.18 ± 1.06
U: 1:3
56.65 ± 1.21
55.98 ± 0.87
58.26 ± 1.08
55.84 ± 1.10
Over sampling
O: 1.5:1.5
60.04 ± 1.54
61.63 ± 1.11
58.25 ± 0.95
57.27 ± 1.26
Table 2: AUC Scores (est. mean ± std) across different class ratios and pretrained models. U = undersampling and O = oversampling.
Sampling
Class Ratio
Year 1
Year 2
VIT-ImageNet
VIT-Corona
VIT-ImageNet
VIT-Corona
-
Original Dataset
52.23 ± 1.17
62.27 ± 3.44
61.30 ± 1.38
59.02 ± 1.55
Under sampling
U: 1:1
64.7 ± 1.52
59.73 ± 2.80
62.47 ± 1.15
59.15 ± 1.67
U: 1:2
61.43 ± 1.57
60.29 ± 1.92
63.39 ± 0.60
61.90 ± 1.32
U: 1:3
64.43 ± 1.26
59.07 ± 2.55
62.13 ± 1.19
59.63 ± 2.19
Over sampling
O: 1.5:1.5
65.00 ± 2.00
63.60 ± 1.17
63.13 ± 0.23
61.37 ± 0.47
Table 3: Balanced Accuracy (mean ± std) across different class ratios and pretrained models. U = undersampling and O = oversampling.
Figure 3: Side-by-side radar plots comparing the evolution of model performance from Year 1 to Year 2 using two pretraining strategies. Both models are evaluated on the same five metrics: Balanced Accuracy, Sensitivity, Specificity, Precision, and F1-score.
Figure 4: Side-by-side comparison of ImageNet-pretrained Vision Transformer models across Year 1 (left) and Year 2 (right). Each polygon represents a fine-tuning configuration (1×1, 1×2, 1×3), evaluated on five metrics: Balanced Accuracy, Sensitivity, Specificity, Precision, and F1-score.
Figure 5: Side-by-side comparison of Corona-pretrained Vision Transformer models across Year 1 (left) and Year 2 (right). Each polygon represents a fine-tuning configuration (1×1, 1×2, 1×3), evaluated across five metrics: Balanced Accuracy, Sensitivity, Specificity, Precision, and F1-score.
Andy · Department of Pathology and Laboratory Medicine, The University of Texas McGovern Medical School, Houston, Texas · Department of Hematopathology, The University of Texas MD Anderson Cancer, Houston, Texas
Department of Computer Science and Engineering, IIT Delhi, New Delhi, India · Yardi School of AI, IIT Delhi, New Delhi, India · Department of Radiodiagnosis, PGIMER Chandigarh, Chandigarh, India
Department of Biomedical Informatics, University of Arkansas for Medical Sciences · Department of Information Science, University of Arkansas at Little Rock · Department of Neuroscience, University of Arkansas for Medical Sciences