A generative flow-based model for volumetric data augmentation in 3D deep learning for computed tomographic colonography

Colorectal cancer (CRC) is the third most common cancer in terms of incidence and the second most common cancer in terms of mortality worldwide [1]. However, CRC can be prevented considerably by early detection and removal of its precursors, colorectal polyps [2]. Computed tomographic colonography (CTC) can be used both for detecting CRCs and for preventing CRCs by early detection of clinically significant polyps that could develop into cancers [3].

Deep learning based on convolutional neural networks (CNNs) has made it easy to obtain state-of-the-art results in various medical imaging tasks [4]. However, one of the limitations of deep learning is that the development of generalizable CNNs requires a large amount and great variety of training data [5]. In CTC, the available training datasets are relatively small and limited in numbers and variations.

The two principal approaches for addressing the issue of small training datasets in deep learning have included (1) transfer learning of CNNs and (2) data augmentation of training datasets. With the first approach, transfer learning, one starts with a CNN that has been pre-trained with a large number of images. The CNN is then adapted to the desired application by continuation of the training with the available domain-specific data. In CTC, transfer learning has been used successfully in training of CNNs for interpreting virtual endoluminal views [6], for detecting contrast-coated serrated polyps [7] and for performing electronic cleansing [8]. However, transfer learning has the limitations that most of the publicly available pre-trained CNNs have not been pre-trained with medical images and that most of these CNNs are 2D CNNs.

With the second approach, data augmentation, the training dataset is enhanced artificially so that the number and variety of training samples are increased. The most common method for implementing data augmentation has been the manipulation of the training samples by basic image processing operations [5]. In CTC, this approach has been used successfully for training of CNNs for detection of polyps [9] and masses [10], and for reduction in the number of false-positive (FP) detections in computer-aided detection (CADe) [11, 12]. However, the approach has the limitation that some of the image transforms that are commonly used with natural images are not appropriate for use with medical images. This can limit the number and variety of the obtainable training samples.

Recently, several studies have explored the possibility of performing data augmentation by use of generative adversarial networks (GANs) [13, 14]. In CTC, 3D GANs have been used for generation of synthetic polyps to improve the training of 3D CNNs in CADe [15]. However, the development of GANs that can generate realistic synthetic images at a high image resolution is known to suffer from various problems, such as non-convergence of the model parameters, mode collapse or training imbalance between the generator and the discriminator [5].

Flow-based generative models have several advantages over GANs [16]. Unlike GANs, they are designed to learn the distribution of the input data explicitly, thereby providing an exact latent-variable inference, log-likelihood evaluation and a meaningful latent space for performance of valid manipulations of the data. They are also highly parallelizable and can be implemented with a small memory footprint [16].

In this study, we investigated the feasibility of applying a flow-based generative model to 3D data augmentation of colorectal polyps in CTC. We hypothesized that (1) a 3D CNN that implements a state-of-the-art flow-based generative model known as Glow [16] (hereafter referred to as 3D Glow) can be trained to generate synthetic polyps that are visually indistinguishable from real colorectal polyps in CTC, and that (2) data augmentation by use of 3D Glow-generated synthetic polyps can be used for improving the performance of deep learning in differentiating between polyps and non-polyps in CADe for CTC. One could use successful development of 3D Glow to overcome the limitations of currently available CTC datasets, thereby enabling the training of generalizable high-performing 3D CNNs for CADe in CTC.

Given an observed (complicated) probability distribution, a flow-based generative model provides a bijective mapping f between the observed distribution and a simple, well-understood target probability distribution, such as a standard Gaussian distribution. The desired computations can then be performed on the simple target distribution rather than on the observed distribution. Such a mapping that is constructed as a sequence of invertible bijective transforms is called the normalizing flow. Formally, let x denote an observation sampled from a training dataset D that has an unknown probability distribution p(x). The normalizing flow can be defined aswhere z is a latent random variable of the desired target distribution. The corresponding flow-based generative model is constructed by optimizing the parameters of the component functions of Eq. (1), \(f_i\), by use of the maximum-likelihood method. The training loss is the negative log-likelihood over D, orFigure 1 illustrates the construction of a normalizing flow. The path traversed by the random variables, \(z_i\), is a flow, and the full chain that is formed by the corresponding probability distributions is the normalizing flow.

Figure 2a illustrates the overall architecture of our 3D Glow model. The model extends the originally proposed 2D Glow framework [16] into a 3D CNN for processing of volumetric CTC images.

The model has three types of layer blocks. Squeeze transforms an input \(w \times h \times d \times c\) tensor, where \(w \times h \times d\) is the input image size and c is the number of channels, into a \(\frac{w}{2} \times \frac{h}{2} \times \frac{d}{2} \times 8c\) tensor; thus, trading spatial size for number of channels [17]. Flow implements one step of the flow. Split splits the input tensor along the channel dimensions [16].

Figure 2b illustrates the design of a Flow block. There are three layers: an activation-normalization (actnorm) layer [16], an invertible \(1 \times 1 \times 1\)-convolutional layer [16] and an affine-coupling layer [18]. The actnorm layer performs an affine transformation of the input tensor, where the trainable per-channel scale and bias parameters are initialized to yield a mean of zero and a standard deviation of one based on the first mini-batch [16]. The invertible \(1 \times 1 \times 1\)-convolutional layer implements a trainable permutation that reverses the order of the input channels, where the weight matrix of the transform is initialized as a random rotation matrix [16]. The affine-coupling layer splits the input into two parts: the first d dimensions remain the same, whereas the latter dimensions, \(d+1\) to D, undergo an affine transformation based on the scale and shift parameters aswhere \(\odot \) denotes an element-wise product and s(.) and t(.) are scale and translation functions that map \({\mathbb R}^d {\rightarrow } {\mathbb R}^{D-d}\).

It can be shown that the transformation above satisfies the basic properties required by efficient computation of the flow transformation [16]. The model is trained with the loss function of Eq. (2).

Figure 3 illustrates the principle of the 3D Glow-based data augmentation. The training of the model involves the learning of a bijective mapping from input polyp VOIs (volumes of interest, \(x_a\) and \(x_b\)) to vectors of the latent space (\(z_a\) and \(z_b\)) by use of the latent variable z and the normalizing flow of Eq. (1).

We interpolate the new samples in the latent space by use of linear interpolation between any two training samples. Given two vectors of the latent space, such as \(z_a\) and \(z_b\), we calculate a new sample as \(\tilde{z} \sim z_a + \alpha (z_b - z_a)\), where \(\alpha \) is sampled randomly from the uniform distribution and scaled linearly to \(\alpha {\in }[0.4, 0.6]\) to avoid generating VOIs that would be nearly identical to the VOI of real polyps in the training dataset.

Although deep learning has made it relatively easy to obtain state-of-the-art results of various medical imaging tasks, in CTC, the available training datasets are relatively small, and the numbers of clinically significant polyps are even smaller. In this study, we investigated the feasibility of developing a flow-based generative model for simulating colorectal polyps to improve the training of 3D CNNs in CADe for CTC. The results of our experiments indicate that this is indeed feasible.

In the first experiment, we performed an observer study to compare the visually observed characteristics of 3D Glow-generated synthetic polyps with those of real polyps. To the best of our knowledge, this is the first study to show that it is possible to generate VOIs of synthetic polyps that are visually equivalent to VOIs of real polyps in CTC.

In the second experiment, we investigated the effect of the training of deep learning (3D DenseNet) with synthetic polyps generated by 3D Glow on the differentiation of polyps from non-polyps in CTC. To the best of our knowledge, this is the first study to show that data augmentation by synthetic polyps can be used to yield a statistically significant performance improvement in deep learning for CTC. We demonstrated that the use of synthetic polyps generated by 3D Glow outperformed the use of traditional nonlinear and GAN-based augmentations, and that the classification performance was highest when 3D Glow was used in combination with a nonlinear augmentation.

The observed improvement in the polyp classification performance by use of the 3D Glow was high for small polyps (6–9 mm in size). This result is important, because it is the detection of small polyps that remains challenging in CTC, whereas many clinical studies have demonstrated that existing CADe systems can detect large polyps in CTC at a high accuracy [31, 32]. Therefore, 3D Glow could have a meaningful impact in improving the detection of small polyps.

We performed the polyp detection by use of volumetric analysis of 3D CT colonography data, because this enables deep learning to review the complete region of the colon regardless of obstructions such as collapsed or poorly distended colon segments. Although other methods such as virtual endoscopic views have been used previously for polyp detection by deep learning [6], such approaches are limited to well-distended regions of the colon, whereas obstructed regions that can hide clinically significant polyps or masses would need to be identified and reviewed by other means. Also, the visual appearance of a virtual endoscopic view depends on the rendering algorithm and its visualization parameters, whereas 3D CT colonography data are readily standardized by Hounsfield units. Thus, the volumetric analysis provides the most uniform approach for the application of deep learning to polyp detection in CT colonography.

This study had several limitations. First, the number of datasets was necessarily relatively small. Second, there are various other data augmentation methods besides the nonlinear image transforms and the GAN-based method that we used as reference augmentation methods in this study. However, a thorough systematic assessment of their comparative performance against 3D Glow will require a separate study. Third, the study could be expanded by consideration of the unique demands of CTC, such as various bowel preparations or polyps that tend to be underrepresented in CTC datasets. Addressing such issues provides topics for future studies.

We developed a 3D CNN based on a flow-based generative model for generating synthetic colorectal polyps in CTC. Our results indicate that the generated synthetic polyps are visually indistinguishable from real polyps and that data augmentation by such polyps can significantly improve the performance of deep learning in CADe for CTC.