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.
Vision transformers (ViT) have been shown to allow for more flexible feature detection and can outperform convolutional neural network (CNN) when pre-trained on sufficient data. Due to their promising feature detection capabilities, we deployed ViTs for morphological classification of anaplastic large cell lymphoma (ALCL) versus classic Hodgkin lymphoma (cHL). We had previously designed a ViT model which was trained on a small dataset of 1,200 image patches in fully supervised training. That model achieved a diagnostic accuracy of 100% and an F1 score of 1.0 on the independent test set. Since fully supervised training is not a practical method due to lack of expertise resources in both the training and testing phases, we conducted a recent study on a modified approach to training data (weakly supervised training) and show that labeling training image patch automatically at the slide level of each whole-slide-image is a more practical solution for clinical use of Vision Transformer. Our ViT model, trained on a larger dataset of 100,000 image patches, yields evaluation metrics with significant accuracy, F1 score, and area under the curve (AUC) at 91.85%, 0.92, and 0.98, respectively. These are respectable values that qualify this ViT model, with weakly supervised training, as a suitable tool for a deep learning module in clinical model development using automated image patch extraction.
Nghia, Nguyen, Amer Wahed +9
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
Vision Transformers (ViT) have become the architecture of choice for many computer vision tasks, yet their performance in computer-aided diagnostics remains limited. Focusing on breast cancer detection from mammograms, we identify two main causes for this shortfall. First, medical images are high-resolution with small abnormalities, leading to an excessive number of tokens and making it difficult for the softmax-based attention to localize and attend to relevant regions. Second, medical image classification is inherently fine-grained, with low inter-class and high intra-class variability, where standard cross-entropy training is insufficient. To overcome these challenges, we propose a framework with three key components: (1) Region of interest (RoI) based token reduction using an object detection model to guide attention; (2) contrastive learning between selected RoI to enhance fine-grained discrimination through hard-negative based training; and (3) a DINOv2 pretrained ViT that captures localization-aware, fine-grained features instead of global CLIP representations. Experiments on public mammography datasets demonstrate that our method achieves superior performance over existing baselines, establishing its effectiveness and potential clinical utility for large-scale breast cancer screening. Our code is available for reproducibility here: https://aih-iitd.github.io/publications/attend-what-matters
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
Three-dimensional models are widely assumed preferable for volumetric medical imaging, yet their practical value depends on whether performance gains justify added computational cost and complexity. Rather than proposing a new architecture, we study how input dimensionality (2D, 2.5D, 3D) affects model behavior across convolutional neural networks (CNNs) and Vision Transformers (ViTs) under a fixed training protocol. Using a leakage-free NLST cohort (n = 1,977) with supporting LIDC-IDRI data, we find that the 2.5D CNN offers the most favorable discrimination-stability trade-off in our comparison (ROC-AUC 0.682, 95% CI [0.546, 0.799]) with a stable operating point. In contrast, 3D CNNs show threshold instability, and transformers exhibit degenerate predictions, such as all-positive predictions. Confidence intervals are wide and overlapping, so we present these results as a controlled resource-performance frontier and a failure-mode taxonomy rather than as definitive superiority claims. For class-imbalanced lung cancer screening classification, 2D and 2.5D inputs provide a more reliable trade-off between performance, stability, and computational efficiency than full 3D representations.
Md Enamul Hoq, Sharafat Hossain, Imraul Emmaka +4
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