Unsupervised anomaly localization in high-resolution breast scans using deep pluralistic image completion

The tasks of anomaly detection (AD), i.e., the classification of entire images as being either normal or anomalous, and anomaly localization/segmentation (AL), i.e., the spatial segmentation of anomalies within images, have received considerable attention within the fields of machine learning and deep learning in particular. Many AD works exist Choi et al. 2018; Ren et al. 2019; Grathwohl et al. 2019; Nalisnick et al. 2018; Nalisnick et al. 2019; Serrà et al. 2019; Du and Mordatch 2019; Schlegl et al. 2019; Choi and Chung 2019; Deecke et al. 2018; Pidhorskyi et al. 2018, but we consider AL, the more challenging task that is also more applicable to clinical practice.

The vast majority of AL methods benchmark on the industrial anomaly detection dataset MVTec-AD Bergmann et al. 2019. Recent works include CutPaste, a self-supervised learning model which trains an encoder neural network to extract features that are useful for differentiating between normal and anomalous data on the proxy task of detecting the “cutting and pasting” of regions of images to another random part of the image Li et al. 2021; PatchSVDD, a patch-based self-supervised model that utilizes support vector data descriptions (SVDDs) to detect and localize anomalies Yi and Yoon 2020; PatchCore, which localizes anomalies within patches by comparing their features to a memory bank of features of patches from normal images based on pretrained neural networks Roth et al. 2021; and PaDiM, a similar patch-based method that estimates the probability distribution of normal class instances Defard et al. 2020. These methods detect anomalies by comparing image features to features from normal training data; our method instead compares directly to a normal “realization” of the image given its surroundings, which is more robust to the high complexity of medical data, in particular breast tissue, where the possible feature similarity between anomalous and normal data can be much higher as compared to other types of data such as MVTec-AD. This is one possible reason for why these other methods, which perform very well on MVTec-AD, have a performance drop on medical data such as DBT, while our method performs significantly better.

Although MVTec-AD serves to model the application of anomaly localization to the industrial setting, the task of AL for medical images is also an important task for multiple reasons. First, we find that AL methods that perform extremely well on MVTec-AD do not necessarily translate well to medical image AL scenarios. This is due to the MVTec-AD data being significantly more controlled, less complex, and much less varied than the data seen in medical images, as well as having visual similarities to images from ImageNet, which have been absorbed by commonly-used pretrained image encoders, i.e., ResNet He et al. 2016. Healthy tissue in medical images often has high semantic variability, uncountably many distinct possible anomalies, and can generally be quite unpredictable, especially in highly variable anatomies like the breast. Simply put, many existing AL methods do not have the ability to fully generalize to the many possible challenging scenarios of medical anomaly detection. We believe that supporting a greater focus of general AD and AL research on the important, safety-critical application of medical imaging is essential for the development of methods that have broader impact.

Another direction of AL is to reconstruct a test image and consider it as anomalous if the reconstruction is distinct from the input. Reconstruction-based methods, e.g., Schlegl et al. 2019; Choi and Chung 2019; Deecke et al. 2018; Pidhorskyi et al. 2018, commonly solve this problem through an encoder-decoder mechanism. These methods can not always be robust at discriminating anomalous data from normal data because given some input image, anomalous data within it may be partially reconstructed even by a normal-trained reconstructor, making anomalies not stand out within the reconstruction error. Image completion-based methods, e.g., Haselmann et al. 2018; Munawar and Creusot 2015; Pirnay and Chai 2021b; Zavrtanik et al. 2021; Pirnay and Chai 2021a; Swiecicki et al. 2021, alleviate this problem by excluding the reconstructed region as input, and creating a normal completion that is more noticeably different than the anomalous ground truth.

Recent works Wan et al. 2021; Liu et al. 2021a; Zhao et al. 2020; Zheng et al. 2019; Dupont and Suresha 2019 approached the goal of producing multiple plausible and diverse completions for a single input. As such, these methods are unnecessary and impractical for our purposes, which we show experimentally in Section 5.2.

Instead, we achieve completion variability by a simple and novel application of spatial dropout layers to a pretrained completion network. Our method requires no additional training, and could theoretically be used on any sort of convolutional deterministic completion network. This keeps our overall anomaly detection method straightforward and intuitive, and importantly, fast.

Our overall anomaly localization method is summarized in Figures 1 (outer, image/slice-level loop) and 2 (inner, patch-level loop). Beginning with some 2D slice of a DBT scan volume, our model creates a “sliding” patch window that rasters through the slice with a fixed stride. For each sliding-window image patch $I$, we apply a mask over the center region $I_{c}\subseteq I$ to obtain an image $I_{m}=I-I_{c}$ with the region missing; we save the missing region $I_{c}$ as the ground truth completion. Next we compare the distribution of predicted normal completions of $I_{m}$ to the ground truth completion $I_{c}$, by examining the $L_{2}$ distance, in a feature space, of the prediction closest to the ground truth. If that distance is above a certain threshold, then $I_{c}$ is anomalous. In Figure 2, $p_{a}(h_{c}|I_{m})$ and $p_{n}(h_{c}|I_{m})$ are the feature space distributions of (1) anomalous and (2) normal completions of $I_{m}$, respectively. The anomaly score for each $I_{c}$ is used for the associated locations of the anomaly heatmap of the entire DBT slice (Fig. 1).

In this section we present the formal definition of the minimum completion distance (MCD) metric, which we use for anomaly detection at the patch level. Anomaly localization is then performed for an entire DBT slice by using the MCD metric on overlapping patches of that slice.

Consider a single test image/patch $I$ with a ground truth completion region $I_{c}$ that has surroundings $I_{m}$, i.e., $I=I_{c}\cup I_{m}$ and $I_{c}\cap I_{m}=\varnothing$. Now, consider sampling $M$ i.i.d. (independent and identically distributed) possible normal completions of $I_{m}$: $\{I_{c}^{1},\ldots,I_{c}^{M}\}\sim p_{n}\left(I_{c}\middle|I_{m}\right)$. Here $p_{n}\left(I_{c}\middle|I_{m}\right)$ is the probability density function (p.d.f.) of the distribution of possible normal completions of $I_{m}$; we denote $p_{a}\left(I_{c}\middle|I_{m}\right)$ as the same but for anomalous completions.

Now, use a pretrained normal data encoder $\phi$ to map the completions to a feature space via $h_{c}^{i}=\phi(I_{c}^{i})$. Assuming $\phi$ to be an injective function, we can construct p.d.f.s of completions within this feature space; i.e., the feature space p.d.f. paired with $p_{n}(I_{c}|I_{m})$ is $p_{n}(h_{c}|I_{m})$. As such, $\phi$ transforms the completion image samples $\{I_{c}^{i}\}_{i=1}^{M}$ to $\{h_{c}^{i}\}_{i=1}^{M}\sim p_{n}\left(h_{c}\middle|I_{m}\right)$.

We then define the MCD anomaly score of $I$ to be

where $h_{c}^{0}=\phi\left(I_{c}\right)$ is the ground truth completion in feature space.

Next, we show that this metric is an arbitrarily powerful anomaly classifier as the sample size $M$ approaches $\infty$, and more practically that the performance improves with high sample efficiency, given reasonable assumptions about how the distributions of anomalous and normal data are distanced from each other. These assumptions are adapted from Proposition 2 of Ye et al. 2021, and are summarized as follows:

We have shown theoretical support for our MCD anomaly metric, and explained why given appropriate assumptions of the distributions of normal and anomalous data, the metric performs well for a reasonably small completion sample size $M$. The central component of this metric is the diverse sampling from $p_{n}(I_{c}|I_{m})$: the distribution of possible normal completions $I_{c}$ of some surroundings $I_{m}$; in other words, obtaining pluralistic completions. We opt for a simple and intuitive approach for creating pluralistic completions that still manages to achieve sufficient feature variability of completions for the MCD metric. Our goal is to make the output of a completion network $G$ trained on normal data variable for some fixed input masked image $I_{m}$. At each $i^{th}$ evaluation of $G(I_{m})$ we wish to obtain a different output completion $I^{i}_{c}$, while still maintaining the ability of $G$ to create fairly realistic normal inpaintings that will be distinct from any anomalous data.

Our intuition is from the dropout Srivastava et al. 2014 mechanism, which is commonly used during training to combat overfitting, but can also be used during inference to produce variable outputs Kendall and Gal 2017. We apply this general prescription to a completion network to induce variability for conditional generative models. This idea is briefly introduced in Wieluch and Schwenker 2019, but they only present it as a proof-of-concept extension of their work, while we manage to implement it in a real application setting.

In particular, we perform completions using the model of Yu et al. 2018, which includes a Wasserstein generative adversarial network (GAN)-based fully-convolutional completion network $G$ and critic/encoder $\phi_{W}$ Arjovsky et al. 2017. $G$ creates a fixed completion $I_{c}$ of some input masked image $I_{m}$, while $\phi_{W}$ learns to discriminate between real vs. fake normal completion data; we use $\phi_{W}$ as the completion feature encoder $\phi$ described in Section 3.2. Wasserstein generative adversarial networks (GANs) are signicantly more reliable to train than traditional GANs Goodfellow et al. 2014: they converge reliably, remove problems such as mode collapse, and have interpretable loss functions, among other benefits. The critic $\phi_{W}$ is the Wasserstein GAN’s version of the traditional GAN’s discriminator.

In our setting, since the input to $G$ is $I_{m}$, a fixed variable, the network has no inherent stochasticity by default. A simple method of adding stochasticity to the input, i.e., $G(I_{m},z)$ where $z$ is sampled from some noise distribution, would not work as well because the network would simply learn to ignore $z$ Isola et al. 2017; Mathieu et al. 2015. In order to allow $G$ to create semantically diverse yet sufficiently high-quality completions at each evaluation of a single $I_{m}$, we propose using spatial, or channel-wise dropout Tompson et al. 2015 within $G$. This type of dropout randomly makes entire channels of convolutional layer activation maps zero with some probability, rather than individual neurons. We give examples of pluralistic completions using our method in Figure 3. We note that using dropout on $G$ naturally results in reduced visual quality of individual completions, due to the variability that dropout adds to the network. However, we found that the benefit of having access to multiple possible completions outweighs this, still resulting in improved tumor detection performance over the single-completion case (Table 2).

Conceptually, because convolutional layer activation maps carry spatial correlations between adjacent pixels, dropping out individual activations randomly can result in low-quality completions, which we found to be the case in practice. On the other hand, dropping an entire channel of an activation map with spatial dropout, can be thought of as inducing a change in the global feature information of the resulting completion, while avoiding any such negative spatial effects. This makes intuitive sense at a high level: for a given layer of a fully connected neural network, the individual neuron’s activations are the key global features that affect the downstream inference; on the other hand, for a convolutional neural network layer, the key global features are different channels of the given activation map. Lee and Lee 2020 in fact found that spatial dropout used on convolutional neural networks (CNNs) can be functionally similar to using regular dropout on fully connected neural networks.

Intriguingly, we did not find any benefit in the quality or diversity of pluralistic completions between the options of (1) using dropout on $G$ during both training and testing or (2) only testing, over the course of many experiments. As such, for the sake of simplicity, we obtain pluralistic completions by only applying dropout at test time to a normally (non-dropout) trained completion network. We also obtained better inpainting quality (on the training set) when dropout is excluded from the shallowest and deepest layers of $G$ (see Appendix A.1 for details). For all experiments we use a dropout probability of $0.5$, a relatively high value that we found suitable for generating sufficiently diverse completions, which was also assisted by dropout being used after the majority of $G$’s layers. We found that different training iterations of $G$ sometimes resulted in differing quality and variability of completions once dropout was applied, but we saw no obvious trends to this, so chose to halt training simply when the $L_{1}$ distance between completions and ground truth was minimized (see Section 5.1 for more training details).

It is also conceivable to explicitly optimize the placement and probabilities of dropout layers to maximize completion diversity and visual quality. However, there is not an explicit, differentiable dependence of a completion diversity metric (e.g., LPIPS Zhang et al. 2018) or quality metric (e.g., the critic/discriminator score) on the dropout layer parameter(s) and/or placement, so it is unclear how these parameters could be efficiently tuned to optimize for these metrics. Doing so would require a non-differentiable optimization method such as Bayesian optimization, which is computationally prohibitive due to the high number of possible dropout parameters to tune (layer-by-layer), and the computational cost of sampling enough completions at each iteration of such a routine to get a reasonable holistic measure of completion diversity and quality. We attempted this Bayesian optimization procedure in early experiments but found that it did not converge; due to these issues, we simply fixed the dropout probability to a fixed value (0.5) for all layers, which we found sufficient for completion diversity and quality.

Our method also has the added benefit that only negligible additional computational load is needed to create a pluralistic completion compared to an ordinary deterministic completion network, as each completion is created by an independent forward pass through $G$. This also makes pluralistic completion sampling easily parallelizable.

We can now determine whether some $d_{p}\times d_{p}$ patch of a DBT scan includes an anomaly within the center square $d_{m}\times d_{m}$ region by sampling $M$ possible normal completions of that region given the surroundings, and using the minimum completion distance (MCD) metric (Eq. (1)) to compare the completions to the missing region ground truth. The final portion of our model is to use this new metric to localize, or segment, anomalies within a full size DBT slice.

Anomaly localization requires synthesizing an anomaly heatmap for a given DBT slice $X$ that is the same size as that slice, where each pixel of the heatmap corresponds to the model’s prediction confidence of the corresponding slice pixel containing anomalous data. To do so, we begin with the $d_{p}\times d_{p}$ image patch at the top left of $X$—our “window”—and apply the MCD metric to that patch, with the aforementioned masked region chosen a priori, to obtain an anomaly score associated with that patch. We then shift the window by some stride according to a basic overlapping raster scan order, perform the same procedure to obtain an anomaly score for this next window, and repeat until all raster windows have been scored, ending with the patch at the bottom right of $X$. The heatmap is all of these scores arranged with the same spatial orientation of the corresponding raster patch centers that created the scores. Finally, we use bicubic interpolation to upsample the heatmap until it is of the same size as $X$. The overall anomaly localization procedure is summarized in Algorithm 1. In practice, the two for loops are parallelized to take full advantage of GPU memory; i.e., multiple completions are sampled, for multiple inputs, all at once. This feature creates a large decrease in computation time compared to our previous work of Swiecicki et al. 2021, where completions were simply made one-at-a-time.

Now we compare previous leading unsupervised/self-supervised anomaly localization methods to our work: quantitative (pixel AUC) results are summarized in Table 2, while qualitative (anomaly heatmap) results are given in Figure 5. We also provide the pixel-wise average precision (AP) score for each method in Table 2. The AP summarizes the precision-recall curve, and is the weighted mean of the precision achieved at each possible scoring threshold along the precision-recall curve, according to scikit-learn’s sklearn.metrics.average_precision_score in Python. Further details and results of these comparison studies are given as follows, where we first explore other state-of-the-art methods, followed by our model.

As few research considers the topic of pluralistic image completion, we compare our dropout pluralistic completion method (Section 3.3) to the state-of-the-art method of Wan et al. 2021, which presents a two-stage, transformer-based model for pluralistic image completion. We trained this method on the same random normal DBT patch dataset as the dropout inpainter, and example inpaintings created by the trained model are shown in Figure 3. Although the completions are slightly more detailed (but still not anatomically valid) than our dropout inpainter, in practice we find that this method is significantly slower than the dropout method for creating multiple completions, such that it becomes impractical for anomaly localization.

On a single 24 GB RTX 3090 GPU, it takes 2.9 days for HFPIC and 8 minutes for our method to generate a single heatmap with the default setting of an $M=10$ completion sample size. This difference is simply due to the significant margin between the size of the two models: ours has about $3.6$ million trainable parameters, while HFPIC has about $450$ million. We further evaluated this difference by (1) fixing $N=5$, the number of input patches to complete, and testing a range of $M$ for both inpainting methods, and (2) fixing $M=10$ and testing a range of $N$. The computation time results are shown in Figure 6; each datapoint was averaged over six possible input DBT slice patches from the test set. We see that in general, HFPIC is slower than our method by about three orders of magnitude.

Moreover, we have tested the effectiveness of HFPIC by using it to create a heatmap for a DBT slice in the test set. Although the extreme computation time makes it impractical to test PICARD with HFPIC on the entire test set, we tested it on a single image (which took days to compute) shown in Figure 7. Here, we actually see a decrease in anomaly localization performance; this is likely due to the anatomically unrealistic nature of HFPIC DBT completions, as shown in Figure 3.

In Section 3.2 we showed that theoretically, the MCD metric (Equation 1) achieves perfect AUC performance in the limit of inpainting sample size $M\rightarrow\infty$. We evaluate this behavior empirically in order to validate these claims by calculating the tumor localization performance (pixel AUC) of PICARD on a range of values of $M$, $\{1,2,5,10,25,50,100,250\}$, on a set of ten DBT scans randomly sampled from the test set, shown in Figure 8. We use a subset instead of the full testing set due to computation feasibility (it takes almost 6 days to evaluate the entire set with $M=250$). Performance does indeed increase asymptotically as $M\rightarrow\infty$, but not to a perfect AUC of $1$. This is due to the fact that in our derivations, we assume that the pluralistic inpainter is able to perfectly sample from the true distribution of possible completions; in practice, the inpainting method is necessarily imperfect, as it is difficult to capture the broad anatomical variability and complexity of breast tissue. Finally, we note that this analysis was completed after all other experiments, where $M=10$ was chosen a priori.

The superior results of PICARD for DBT breast lesion detection are quite promising, however, further refinements could be applied to make the method even more powerful.

One of the difficulties of generative modeling of breast tissue is the extremely high complexity and natural variability of the tissue, so that it is difficult to obtain anatomically realistic completions, even with state-of-the-art methods like HFPIC Wan et al. 2021. In addition to complex local details, breast tissue can have complicated correlations between distant image locations, which may not be able to be fully captured by the inherent locality of convolutional neural networks models. Indeed, the coarse-to-fine feature hierarchy of traditional convolutional image completion methods–such as the one used in this work–is well-suited for natural images, where low-resolution features are more global, yet it is unable to fully represent the complexities of breast tissue. Visual transformer-based models, e.g., Dosovitskiy et al. 2020; Liu et al. 2021b can model long-range pixel interactions, but even the transformer-based model of HFPIC was unable to produce anatomically realistic content. As such, it is unclear what type of generative model would be able to learn reasonable representations of breast data that preserve both local fine-grained details while maintaining the complex global structure of the tissue. Such a model may need to include some sort of inductive biases for the unique structures seen in visual anatomical data; alternatively, entirely different generative models may prove useful, such as normalizing flows, energy-based, or score-based methods, which we leave for future works.

Having more realistic completions would better approximate sampling from the true distribution of possible completions, theoretically leading to more robust minimum completion distance performance, and therefore better anomaly localization. This would fix some of the issues of false-positive regions that can be seen in some of PICARD’s heatmaps, that have breast tissue that is labeled as healthy, but still possesses visual features that are uncommon in the training set. We found that even training our inpainter(s) on the full DBT training set of normal slices did not improve performance, so it appears to be a limitation of the model structure rather than the dataset size.

Although PICARD’s tumor localization performance does receive a boost from using multiple completions instead of just one (present in the first two rows of Table 2 where $M=10$, vs. the next two with $M=1$), in theory the difference could be higher, again if the sampled completions were more realistic and better approximated the true distribution of possible normal completions. One possible solution to this would be to choose a dropout probability for the completion network that results in optimal anatomical realism. However, such optimization needs access to some cancer images during the validation phase, greatly reducing its range of applications. It may be possible to quantify the anatomical realism of completions generated on some validation set of only healthy cases, and optimize the dropout probability to maximize this quantity, but we leave this nontrivial task for future works.

Another possible future work is that as the completion region $I_{c}$ is our region of interest, we made no assumptions about if the surrounding region $I_{m}$ contains anomalies. Still, it may be worth considering how to detect anomalies within $I_{m}$ as well, which could begun with considering the joint distribution $p(I_{c},I_{m})$ rather than just $p(I_{c}|I_{m})$ as in this work. However, it is unclear if this would improve heatmapping performance, as we use a stride small enough such that all pixels (beyond a “padding region” on the border) within a DBT slice will be included at least once within some evaluated $I_{c}$.