Risk Stratification of Lung Nodules Using 3D CNN-Based Multi-task Learning

TL;DR

This paper introduces a 3D CNN-based multi-task learning framework for lung nodule risk stratification, leveraging volumetric CT data, transfer learning, and radiologist disagreement modeling to improve malignancy prediction accuracy.

Contribution

It presents a novel 3D CNN approach with multi-task learning and graph regularization to enhance lung nodule characterization from CT scans.

Findings

Achieved state-of-the-art malignancy score regression results.

Utilized transfer learning for discriminative feature extraction.

Incorporated radiologist disagreement modeling to improve attribute scoring.

Abstract

Risk stratification of lung nodules is a task of primary importance in lung cancer diagnosis. Any improvement in robust and accurate nodule characterization can assist in identifying cancer stage, prognosis, and improving treatment planning. In this study, we propose a 3D Convolutional Neural Network (CNN) based nodule characterization strategy. With a completely 3D approach, we utilize the volumetric information from a CT scan which would be otherwise lost in the conventional 2D CNN based approaches. In order to address the need for a large amount for training data for CNN, we resort to transfer learning to obtain highly discriminative features. Moreover, we also acquire the task dependent feature representation for six high-level nodule attributes and fuse this complementary information via a Multi-task learning (MTL) framework. Finally, we propose to incorporate potential disagreement among radiologists while scoring different nodule attributes in a graph regularized sparse multi-task learning. We evaluated our proposed approach on one of the largest publicly available lung nodule datasets comprising 1018 scans and obtained state-of-the-art results in regressing the malignancy scores.

Figures

Tables

Table 1. Table 1: Classification accuracy and mean absolute score difference of the proposed multi-task learning method in comparison with the other methods.

Methods	Accuracy	Mean score diff.
GIST features with LASSO	76.83%	0.6753
3D CNN MTL with Trace norm	80.08%	0.6259
Proposed method (Equation 7)	91.26%	0.4593

Equations

$$\underset{W}{\min }{\Vert XW-Y\Vert }_2^2+\lambda {\Vert W\Vert }_1,$$

$$\begin{array}{r}\underset{W}{\min }{\Vert XW-Y\Vert }_2^2,\\ s.t.{\Vert W\Vert }_1\le t,\end{array}$$

$$\underset{\mathbf{W}}{\min }\sum _{i=1}^M{\Vert X_i{W}_i-Y_i\Vert }_2^2+\rho {\Vert \mathbf{W}\Vert }_{\ast },$$

$${\Vert \mathbf{W}S\Vert }_F^2=\sum _{i=1}^{\Vert \mathcal{E}\Vert }{\Vert \mathbf{W}{e}^i\Vert }_2^2=\sum _{i=1}^{\Vert \mathcal{E}\Vert }{\Vert {\mathbf{W}}_{e_a^i}-{\mathbf{W}}_{e_b^i}\Vert }_2^2.$$

$${\Vert \mathbf{W}S\Vert }_F^2=\text{tr}((\mathbf{W}S{)}^T(\mathbf{W}S))=\text{tr}(\mathbf{W}S{S}^T{\mathbf{W}}^T)=\text{tr}(\mathbf{W}\mathcal{L}{\mathbf{W}}^T),$$

$$\mbox{\large$\Psi(j)=(e^{\frac{-\sum_{i}(x_{i}^{j}-\mu^{j})^{2}}{2\sigma^{p}}})^{-1}$}.$$

$$\underset{\mathbf{W}}{\min }\sum _{i=1}^M\stackrel{①}{\overbrace{{\Vert (X_i+{\Psi }_i)W_i-Y_i\Vert }_2^2}}+\stackrel{②}{\overbrace{{\rho }_1{\Vert \mathbf{W}S\Vert }_F^2}}+\stackrel{③}{\overbrace{{\rho }_2{\Vert \mathbf{W}\Vert }_1}},$$

Full text

1 Introduction

Cancer is the number-one cause of deaths in the world. Out of 8.2 million deaths due to cancer worldwide, lung cancer accounts for the highest number of mortalities i.e. 1.59 million [1]. Risk stratification of lung nodules can aid in identifying cancer stage leading to improved treatment and higher chances of survival. In addition, any significant development to accurately and automatically characterize lung nodules can save significant manual exertion as well as valuable time.

Early diagnosis is one of the ways to reduce deaths related to lung cancer [2]. In this regard, lung screening programs are especially beneficial. Low Dose Computed Tomography (CT) scans are usually used to perform lung nodule diagnosis, including both detection and risk stratification. Although CT imaging remains the gold standard for lung cancer detection and diagnosis, Computer-Aided Diagnosis (CAD) and quantification tools are often necessary. Moreover, research in developing CAD algorithms can help explore the domain of imaging features and biomarkers which can be then studied by radiologists to further improve clinical decision making.

The development of a fast, robust and accurate system to perform risk stratification of lung nodules is therefore of significant importance. Specially the availability of large publicly available datasets such as LIDC-IDRI from Lung Image Database Consortium [3] has helped accelerate the research in this regard. However, the variability in nodule characteristics, including shape, size, intensity, location and uncertainty among radiologists’ interpretation have made this problem particularly challenging. The advancement in machine learning methods, including the development of novel classification and feature learning techniques, has increased the efficacy of this task. However, there remains a substantial progress to be done in order to develop a CAD system attractive enough to be used in routine clinical evaluations of lung nodules.

In this work, we address the challenge of risk-stratification of lung nodules in low-dose CT scans. Capitalizing on the significant progress of deep learning technologies for image classification and their potential applications in radiology [4], we propose a 3D Convolutional Neural Network (CNN) based approach for rich feature representation of lung nodules. We argue that the use of 3D CNN is paramount in the classification of lung nodules in low-dose CT scans which are 3D by nature. By using the conventional CNN methods, however, we implicitly lose the important volumetric information which can be very significant for accurate risk stratification. The superior performance of 3D CNN over 2D networks is well studied in [5]. We also avoid hand-crafted feature extraction, painstaking feature engineering, and parameter tuning. Moreover, any information about six high-level nodule attributes such as calcification, sphericity, margin, lobulation, spiculation and texture (Figure 1) can help in improving the benign-malignant risk assessment of the nodules. Taking forward this idea, we identify features corresponding to these high-level nodule attributes and fuse them in a multi-task learning framework to obtain the final risk assessment scores. An overview of the proposed approach is presented in Figure 2. Overall, our main contributions in this work can be summarized as follows:

• •

We propose a 3D CNN based method to utilize the volumetric information from a CT scan which would be otherwise lost in the conventional 2D CNN based approaches. Moreover, we also circumvent the need for a large amount of volumetric training data to train the 3D network by transfer learning. We use the CT data to fine-tune a network which is trained on 1 million videos. To the best of our knowledge, our work is the first to empirically validate the success of transfer learning of a 3D network for lung nodules.

• •

We perform experimental evaluations on one of the largest publicly available datasets comprising lung nodules from more than 1000 low-dose CT scans.

• •

We employ graph regularized sparse multi-task learning to fuse the complementary feature information from high-level nodule attributes for malignancy determination. We also propose a scoring function to measure the inconsistency in risk assessment among different experts (radiologists).

2 Related Work

Conventionally, the characterization of lung nodules comprised nodule segmentation, extraction of hand-crafted imaging features, followed by the application of an off-the-shelf classifier/regressor. The method by Uchiyama et al. [6] was based on the extraction of various physical measures, including intensity statistics and then classification using Artificial Neural Networks. El-Baz et al. [7] first segmented the lung nodules using appearance-based models and used spherical harmonic analysis to perform shape analysis. The final step was the classification using $k$-nearest neighbor. Proposing a study based on texture analysis, Han et al. [8] extracted 2D texture features such as Haralick, Gabor and Local Binary Patterns (LBP) and extended them to 3D. Support Vector Machine (SVM) was employed to perform the classification. In another classical work by Way et al. [9], segmentation is performed using 3D active contours followed by the extraction of texture features from the rubber band straightening transform of the surrounding voxels. The classification was performed using Linear Discriminant Analysis (LDA) classifier. In another study, Lee et al. [10] proposed a feature selection based approach using both imaging and clinical data. An ensemble classifier, combining genetic algorithm (GA) and random subspace method (RSM) was then used to gauge feature relevance and information content. Finally, LDA was employed to perform classification on the reduced feature set.

Following up on the success of deep learning, the medical imaging community has moved from feature engineering to feature learning. In those frameworks, CNN had been used for feature extraction and an off-the-shelf classifier such as Random Forest (RF) was employed for classification [11, 12]. Recently, Buty et al. [12] combined spherical harmonics along with deep CNN features and then classified them using RF. However, the use of CNN for lung nodule classification has been confined to 2D image analysis [13], thus falling short of utilizing the important volumetric and contextual information.

Moreover, the use of high-level image attributes had been found to be instrumental in the risk assessment and classification of lung nodules. In an effort to study the relationship between nodules attributes and malignancy, Furuya et al. [14] found that in a particular dataset, 82% of the lobulated, 97% of the densely spiculated, 93% of the ragged and 100% of the halo nodules were malignant. Moreover, 66% of the round nodules were found to be benign. Inspired by this study, in this work we utilize 3D CNN to learn discriminative feature set corresponding to each of the 6 attributes. We then fuse these feature representations via MTL to determine the malignancy likelihood.

3 Method

3.1 Problem Formulation

Let $X=[x_{1},x_{2}\dots x_{n}]\in\mathbb{R}^{n\times d}$ be the data matrix comprising features from $n$ data points in $\mathbb{R}^{d}$. Each sample corresponds with a regression score given by $Y=[y_{1},y_{2}\dots y_{n}]$ where $Y\in\mathbb{R}^{n\times 1}$. Here the objective is to learn the coefficient vector or the regression estimator $W$ from the training data. In this case, the $\ell_{1}$ regularized least square regressor is defined as:

[TABLE]

where $\lambda$ controls the sparsity level for coefficient vector $W=[w_{1},w_{2}\dots w_{d}]$. The problem in Eq. 1 is an unconstrained convex optimization problem, and it remains non-differentiable when $w_{i}=0$. Hence, the closed form solution corresponding to the global minimum for Eq. 1 is not possible. Thus, the above equation is represented in the following way as a constrained optimization function:

[TABLE]

where $t$ is inversely proportional to $\lambda$. In the representation given in Eq. 2, both optimization function and the constraint are convex.

3.2 Network Architecture and Transfer Learning

We use the lung nodules dataset to fine-tune a 3D CNN trained on Sports-1M dataset [15]. The sports dataset comprises 1 million videos with 487 classes. In the absence of a large number of training examples from lung nodules, we use transfer learning strategy to obtain rich feature representation from a larger dataset (Sports-1M) for lung nodule characterization. The Sports-1M dataset is used to train a 3D CNN [5]. The network comprises 5 convolution, 5 max-pooling, 2 fully-connected and 1 soft-max classification layers. The input to the network is 3$\times$16$\times$128$\times$171 where there are 16 non-overlapping slices in the input volume. The first 2 convolution layers have 64 and 128 filters respectively, whereas there are 256 filters in the last 3 layers. The outputs of the fully connected layers are of 4096 dimensions.

3.3 Multi-task learning

Consider a problem with $M$ tasks representing different attributes corresponding to a given dataset $D$. These tasks may be related and share some feature representation, both of which are unknown. The goal in Multi-task learning (MTL) is to perform joint learning of these tasks while exploiting dependencies in feature space so as to improve regressing one task using the others. In contrast to multi-label learning, tasks may have different features in MTL. Each task has model parameters denoted by $W_{m}$, used to regress the corresponding task $m$. Moreover, when W $=[W_{1},W_{2}\dots W_{M}]\in\mathbb{R}^{M\times d}$ represents a rectangular matrix, rank is considered as an extension to the cardinality. In that case, trace norm, which is the sum of singular values is a replacement to the $\ell_{1}$-norm. Trace norm, also known as nuclear norm is the convex envelope of the rank of a matrix (which is non-convex), where the matrices are considered on a unit ball. By replacing $\ell_{1}$-norm with trace norm in Eq. 1, the trace norm regularized least square loss function is given by:

[TABLE]

where $\rho$ tunes the rank of the matrix W, and trace-norm is defined as: $\left\|\mathbf{W}\right\|_{*}=\sum_{i=1}\sigma_{i}(\mathbf{W})$ with $\sigma$ representing singular values.

Another regularizer, pertinent to MTL, is the regularization on the graph representing the relationship between the tasks [16, 17]. Consider a complete graph $G=(V,\mathcal{E})$, such that nodes $V$ represent the tasks and the edges $\mathcal{E}$ encode any relativity between the tasks. The complete graph can be represented as a structure matrix $S=[e^{1},e^{2}\dots e^{\left\|\mathcal{E}\right\|}]$ and the difference between all the pairs connected in the graph is penalized by the following regularizer:

[TABLE]

Herein, $e^{i}_{a}$, $e^{i}_{b}$ are the edges between the nodes $a$ and $b$. The above regularizer can also be written as:

[TABLE]

where ‘tr’ represents trace of a matrix and $\mathcal{L}=SS^{T}$ is the Laplacian matrix. Since there may exist disagreements between the scores from different experts (radiologists), we propose a scoring function to measure potential inconsistencies:

[TABLE]

The inconsistency measure corresponding to a particular example $j$ is represented by $\Psi(j)$. $x_{i}^{j}$ is the score given by the expert (radiologist) $i$ and $\mu^{j}$ and $\sigma^{j}$ denote mean and standard deviation of the scores, respectively. Here, for simplicity, we have dropped the index for the task; however, note that the inconsistency measure is computed for all the tasks. The final proposed graph regularized sparse least square optimization function with the inconsistency measure can then be written as:

[TABLE]

where $\rho_{1}$ controls the level of penalty for graph structure and $\rho_{2}$ controls the sparsity. In the above optimization, the least square loss function ① considers tasks to be decoupled whereas ② and ③ consider the interdependencies between different tasks.

3.4 Optimization

The optimization function in Eq. 7 cannot be solved through standard gradient descent because the $\ell_{1}-$norm is not differentiable at $\mathbf{W}=0$. Since the optimization function in Eq. 7 has both smooth and non-smooth convex parts, estimating the non-smooth part can help solve the optimization function. Therefore, accelerated proximal gradient method [18, 19] is employed to solve the Eq. 7. The accelerated proximal method is the first order gradient method with a complexity of $O(1/k^{2})$, where $k$ is the iteration counter. Note that in Eq. 7, the $\ell_{1}$-norm comprises the non-smooth part and the proximal operator is used for its estimation. The steps in the proposed approach are summarized in Algorithm 1.