A 3D Probabilistic Deep Learning System for Detection and Diagnosis of Lung Cancer Using Low-Dose CT Scans

TL;DR

This paper presents a 3D deep learning system for lung cancer detection and diagnosis from low-dose CT scans, achieving state-of-the-art results and providing calibrated probabilities for better clinical decision-making.

Contribution

It introduces an end-to-end 3D CNN system that couples detection and diagnosis, incorporates model uncertainty, and improves robustness without false positive reduction stages.

Findings

State-of-the-art detection and classification performance on LUNA16 and Kaggle datasets.

Coupled detection and diagnosis improves robustness.

Model uncertainty enables well-calibrated probabilities for clinical use.

Abstract

We introduce a new computer aided detection and diagnosis system for lung cancer screening with low-dose CT scans that produces meaningful probability assessments. Our system is based entirely on 3D convolutional neural networks and achieves state-of-the-art performance for both lung nodule detection and malignancy classification tasks on the publicly available LUNA16 and Kaggle Data Science Bowl challenges. While nodule detection systems are typically designed and optimized on their own, we find that it is important to consider the coupling between detection and diagnosis components. Exploiting this coupling allows us to develop an end-to-end system that has higher and more robust performance and eliminates the need for a nodule detection false positive reduction stage. Furthermore, we characterize model uncertainty in our deep learning systems, a first for lung CT analysis, and show that we can use this to provide well-calibrated classification probabilities for both nodule detection and patient malignancy diagnosis. These calibrated probabilities informed by model uncertainty can be used for subsequent risk-based decision making towards diagnostic interventions or disease treatments, as we demonstrate using a probability-based patient referral strategy to further improve our results.

Figures

Full text

Our CADe system was developed to have as high of a sensitivity/recall as possible for identifying and localizing pulmonary nodules while still keeping the number of false positives low enough to be manageable for our CADx system. It is able to analyze any size CT scan and uses only local visual information to make its decisions. Our CADe system is made up of 1) a 3D segmentation network, which labels every voxel of the CT scan with a nodule probability, and 2) a 3D scoring network, which computes refined nodule probability estimates for full nodule candidates generated from the segmentation and allows for both greater interpretability and false positive reduction. The following subsections explain these networks in detail.

0.1 Segmentation for Candidate Extraction

Our primary nodule segmentation network is a 3D fully-convolutional neural network based on the V-Net architecture, which has been demonstrated to be effective for 3D medical image segmentation [vnet_3dv16]. Our architecture111We began with the V-Net architecture as described in the paper [vnet_3dv16] and tuned the number of layers, number of features, non-linearity functions, and normalization layers based on performance on the LUNA16 dataset. uses three encoder-decoder block pairs, with corresponding skip connections, in addition to the input and output blocks. Our encoder blocks are made up of a $2\times$ downsampling convolution, two layers of kernel size 3 convolutions, and residual connection to the output. Decoder blocks are the same, but with a $2\times$ upsampling deconvolution. The innermost encoder-decoder block pair include channel-wise dropout between the sampling convolution and the two main convolutional layers. All blocks use instance normalization [ulyanov16] instead of batch normalization as well as ReLU nonlinearities.

For training, we use the LUNA16 pulmonary nodule annotations, which provide nodule centers and diameters, to annotate individual voxels of CT scans with spherical masks. Since the number of voxels corresponding to the nodule annotations is very small compared to the total number of voxels, and each CT scan is too large to realistically use to train our network, we sample $64^{3}$-shaped blocks from the original images. For each patient, we sample blocks near known nodules with a probability of $1-0.7^{N}$ and sample randomly from the image the with a probability of $0.7^{N}$, where $N$ is the number of nodules for that patient. We employed this random sampling strategy in order to ensure that the network is exposed to true nodules a sufficient number of times as well as a diverse selection of background. We train the network with a cross entropy loss function that weights voxels within a nodule twice as much as background voxels. Each patient scan is seen once per epoch, regardless of how many nodules are contained within. Our network is trained with 16 block samples per batch using the Adam optimizer [adam_iclr15] with learning rate of $10^{-3}$ for 2500 epochs.

As full preprocessed CT scan images are too large to fit into GPU memory222We used an NVIDIA Quadro GP100 or equivalent for all model training and evaluation described in this paper. while being passed through the trained nodule segmentation model at test time, we split images into eight $256^{3}$-shaped overlapping blocks and stitch the output segmentations together, weighting voxels in the overlapping edge regions by how much of their field-of-view was contained within the borders of their respective input blocks. This results in a full image segmentation that is as good as if the entire image were evaluated by the network in one pass on a CPU, but requires an order of magnitude less time.

CADe candidates are generated from the full segmented image by thresholding the output voxel scores, applying a nearest-neighbors binary opening filter, and labeling all separate regions based on a voxel connectivity of one. The center of each nodule candidate is taken to be the center of mass of the voxel scores of each constituent voxel after thresholding. Sample segmentations of generated candidates on LUNA16 test data are shown in Figure 1 alongside the original spherical annotation.

0.2 False Positive Reduction for LUNA16 Evaluation

All generated candidates are supplied to the CADx system. However, for the purposes of evaluating our CADe system on the LUNA16 benchmark, we developed a scoring network that operates on $32^{3}$-shaped blocks around the candidate center. The architecture of this network is shown in Figure LABEL:fig:cnn_arch. The network was trained with the SGD optimizer (learning rate 0.01, momentum 0.9, and $0.8\times$ learning rate decay every 100 epochs) with batches of 16 candidates for 2500 epochs, with true nodule candidates weighted twice as heavily as false positive candidates in the cross-entropy loss. As we show in Section LABEL:sec:results, our scoring network can be used to dramatically reduce the number of false positives found by the segmentation network.

Our candidate scoring network outputs include model uncertainty as quantified by Monte Carlo (MC) dropout [gal_iclr16, gal_icml16]. We show in Section LABEL:sec:referral that the addition of this model uncertainty makes these scores well-calibrated to the LUNA16 data distribution and thus the probabilities output by the network for each candidate are directly interpretable as real probabilities of being a nodule.