AnoFPDM: Anomaly Segmentation with Forward Process of Diffusion Models for Brain MRI
Yiming Che Fazle Rafsani Jay Shah Md Mahfuzur Rahman Siddiquee Teresa Wu




Anomaly Segmentation with Forward Process of Diffusion Models




Y. Che et al.



Arizona State University 

 ASU-Mayo Center for Innovative Imaging



AnoFPDM: Anomaly Segmentation with Forward Process of Diffusion Models for Brain MRI
    Yiming Che 1,2 Fazle Rafsani 1,2* Jay Shah 1,2* Md Mahfuzur Rahman Siddiquee 1,2 Teresa Wu 1,2
    April 29, 2024
==================================================================================================




Weakly-supervised diffusion models (DM) in anomaly segmentation, leveraging image-level labels, have attracted significant attention for their superior performance compared to unsupervised methods. It eliminates the need for pixel-level labels in training, offering a more cost-effective alternative to supervised methods. However, existing methods are not fully weakly-supervised because they heavily rely on costly pixel-level labels for hyperparameter tuning in inference. To tackle this challenge, we introduce Anomaly Segmentation with Forward Process of Diffusion Models (AnoFPDM), a fully weakly-supervised framework that operates without the need for pixel-level labels. Leveraging the unguided forward process as a reference, we identify suitable hyperparameters, i.e., noise scale and threshold, for each input image. We aggregate anomaly maps from each step in the forward process, enhancing the signal strength of anomalous regions. Remarkably, our proposed method outperforms recent state-of-the-art weakly-supervised approaches, even without utilizing pixel-level labels.

Code is available at <https://github.com/SoloChe/AnoFPDM>* Equal contribution








§ INTRODUCTION


Anomaly segmentation often requires a large number of pixel-level annotations. However, acquiring pixel-level annotations is not only costly but also prone to human annotator bias. Hence, weakly-supervised generative methods <cit.>, leveraging image-level labels in training, are gaining attention. Among these methods, diffusion models <cit.> are commonly chosen as the backbone due to their superior performance compared to other generative methods such as generative adversarial networks (GANs) <cit.>  and variational autoencoder (VAE) <cit.>. However, current weakly-supervised methods are not truly fully weakly-supervised. They still heavily depend on pixel-level labels for hyperparameter tuning during the inference stage. For diffusion model-based methods, this tuning includes determining the appropriate amount of noise (noise scale) added to the input image as well as setting the threshold for the anomaly map. Consequently, the need for pixel-level labels in hyperparameter tuning reintroduces the cost and bias.

In this paper, we propose a fully weakly-supervised framework named Ano-FPDM, built upon diffusion models with classifier-free guidance <cit.>, to eliminate the need of pixel-level labels in hyperparameter tuning. Instead of relying directly on pixel-level labels, we determine the optimal hyperparameters, such as noise scale and threshold, by leveraging the unguided forward process as a reference for the guided forward process. Our model follows the standard training process of classifier-free guidance as used in <cit.> with the image-level labels, i.e., healthy and unhealthy. The inference stage is illustrated in Fig. <ref>. Our framework utilizes the forward process of DM instead of the sampling (backward) process. During the forward process, i.e., from _0 to _T, we gradually introduce noise to the input _0. The denoised inputs, representing the prediction of the original input from the noised input, is obtained with respect to the healthy label guidance and no guidance. The denoised inputs at step t without guidance _0,t^∅ will not remove but compress the anomalous regions during forward process, i.e., a simple compression process of input _0. On the contrary, the denoised input with healthy label guidance _0,t^h removes the anomalous regions while compressing the non-anomalous regions. The two MSEs MSE_t^∅ = _0,t^∅ - _0^2 and MSE_t^h = _0,t^h - _0^2 are collected and jointly used for hyperparameter tunning. The anomaly map is derived by aggregating all MSE_t^h for t∈{t_i}_i=0^t_e. The maximal difference between the two MSEs is used to determine the end step t_e for the guided forward process. Here, the end step t_e directly controls noise scale. The threshold of anomaly map is determined by the quantile of the anomaly map, which is individually selected for each input. The quantile is selected by the maximal difference between MSE_t^h and MSE_t^∅ for each input because the maximal difference is roughly linearly related to the size of the anomalous regions. A smaller quantile is selected for a larger anomalous region to include more possible pixels.

In Fig. <ref>, we further demonstrate our idea by an unhealthy sample from BraT21 dataset <cit.> and exhibit four components, i.e., _0,t^∅, _0,t^h,MSE_t^∅ and MSE_t^h, in our framework. During the forward process, the anomalous regions in the path _0,t^h,t∈{t_i}_i=0^T are effectively removed due to healthy label guidance. We found that the anomalous regions are more likely to be removed compared to the non-anomalous regions during forward process. In our method, we aggregate all MSE_t^h for t∈{t_i}_i=0^t_e as our final anomaly map to increase the signal strength of anomalous regions. The rationale behind this lies in the idea that all anomaly maps in the path _0,t^h,t∈{t_i}_i=0^T contribute to the same anomalous region while not all of them contribute to the same non-anomalous region. In contrast, anomalous regions persist in the path _0,t^∅,t∈{t_i}_i=0^T. The high frequency components are compressed before low frequency components, e.g., anomalous regions. Initially, the difference between MSE_t^∅ and MSE_t^h increases because the unguided path tends to compress high frequency details (non-anomalous region) while the healthy label guided path is more likely to remove low frequency anomalous regions. At the maximal difference, e.g., t=400 in Fig. <ref>, we interpret it as the low frequency anomalous regions are fully removed in the path _0,t^h,t∈{t_i}_i=0^400, while all high frequency details are compressed in the path _0,t^∅,t∈{t_i}_i=0^400. Subsequently, the difference starts to decrease, signifying the compression of non-anomalous regions in the path _0,t^h,t∈{t_i}_i=400^600 and low frequency anomalous regions in the path _0,t^∅,t∈{t_i}_i=400^600. Consequently, we select the step t that achieves the maximal difference between MSE_∅^t and MSE_h^t as the end step t_e. MSE_∅^t serves as a reference, indicating the completion of anomalous region removal, thereby eliminating the need for pixel-level labels. 

Contributions: (i) Eliminating the need for pixel-level labels in hyperparameter tuning by using unguided denoised inputs as reference; (ii) A novel anomaly segmentation framework using the forward process of DM (iii) A novel anomaly map aggregation strategy to enhance the signal strength of anomalous regions.

Related Work
Prior to the development of DM, GAN and VAE dominated the field of anomaly segmentation <cit.>. However, GAN is often criticized for their unstable training, while VAE is not ideal because of their limited expressiveness in the latent space and a tendency to produce blurry reconstructions. DM is a promising alternative due to its successes in image synthesis <cit.>. In the arena of unsupervised anomaly segmentation, innovations such as AnoDDPM <cit.> have applied the unguided denoising diffusion probabilistic model (DDPM) with simplex noise <cit.> to achieve notable segmentation outcomes. Furthermore, Behrendt et al. <cit.> have enhanced model performance by utilizing patched images as inputs to mitigate generative errors, while Iqbal et al. <cit.> have explored the use of masked images in both physical and frequency domains. The latent diffusion model concept, detailed by Rombach et al. <cit.>, has been adeptly incorporated for fast segmentation in Pinaya et al. <cit.>. In the domain of weakly-supervised approaches, pioneering efforts by Wolleb et al. <cit.> have leveraged DDIM with classifier-guidance <cit.>, a technique paralleled in Sanchez et al. <cit.>'s adoption of DDIM with classifier-free guidance <cit.>, similar to the DDIB approach <cit.>





§ BACKGROUND




 §.§ Diffusion Models


During the training, the model ϵ_θ which is parameterized by θ learns the unknown data distribution q(_0) by adding noise to the data and subsequently denoising it. The model here is usually the U-net like architectures <cit.>. To be more specific, the forward process of DDPM is factorized as q(_1:T|_0) = ∏_t=1^Tq(_t |_t-1), which is fixed to a Markov chain with variance schedule β_1,...,β_T. The transition from _t-1 to _t is modeled as Gaussian q(_t |_t-1) = _t |√(1-β_t)_t-1, β_t and the distribution of transition from any arbitrary step t, i.e., forward process, is in closed form

    _t = √(α̅_t)_0 + √(1-α̅_t)ϵ,

where ϵ∼0,, α̅_t = ∏_i=1^t α_i and α_t = 1-β_t.

The sampling process p_θ_1:T = p(_T)∏_t=1^Tp_θ(_t-1|_t) is the transition from _T to _0 with the learned distribution p_θ_t-1|_t which is the approximation of inference distribution q(_t-1|_t, _0) = q(_t|_t-1,_0)q(_t-1|_0)/q(_t|_0). The forms of both distributions q and p_θ are Gaussian. The training process is the minimization of KL-divergence 

    KLq(_1:T|_0)| p_θ(_1:T|_0) = 𝔼__1:T∼ q(_1:T|_0)logq(_1:T|_0) p_θ(_0)/p_θ(_0:T).

A variation of DDPM is DDIM <cit.>. It is non-Markovian, incorporating the input _0 into the forward process. The inference distribution is factorized as q_σ(_1:T|_0) = q_σ(_T|_0)∏_t=2^Tq_σ(_t-1|_t,_0) for all t > 1.




The forward process can be obtained through Bayes' theorem in closed form:

    q_σ(_t+1|_t,_0) = √(α̅_t+1)_0 + √(1-α̅_t+1-σ_t^2)_t-√(α̅_t)_0/√(1-α̅_t),σ_t^2I,

where σ_t is set to 0. They both share the same objective function in training.











 §.§ Classifier-free Guidance


The classifier-free guidance exhibits superior performance in generative tasks compared to classifier guidance <cit.>. Additionally, both training and sampling processes are simplified as it eliminates the need for an external classifier. In the sampling process, the predicted noise is replaced by the noise with guidance. For any t≥0, the predicted noise ϵ_θ^i,t with guidance c_i is ϵ_θ^i,t = (1+w)ϵ_θ(_t, t, c_i) -wϵ_θ(_t, t, ∅). The parameter w controls the strength of guidance. For null label ∅, the predicted noise ϵ_θ^i,t=ϵ_θ(_t, t, ∅). The model ϵ_θ with guidance is implemented by utilizing the extra attention mechanism <cit.>.



§ METHODOLOGY

We utilize the forward process to extract the denoised inputs _0,t^h and  _0,t^∅ for t>0 in terms of healthy label and null label. Then, we collect the mean square error MSE_t^h and MSE_t^∅ for further analysis. Our first step involves determining the health status of the input _0 by assessing the cosine similarity between MSE_t^h and MSE_t^∅. The details can be found in supplementary material. If the input is deemed unhealthy, we determine the optimal end step t_e∈{t_j}_0^T and threshold for segmentation. 



 §.§ Denoised Inputs with Classifier-free Guidance


 We can add noise to the inputs in forward process with either DDPM style using Eq. <ref> or DDIM style using Eq. <ref> without label information. After adding the noise on step t, we extract the denoised inputs with healthy guidance and without guidance by Eq. <ref> and Eq. <ref> respectively.

    _0,t^h = _t - √(1-α̅_t)(1+w)ϵ_θ(_t, t, h) -wϵ_θ(_t, t, ∅)/√(α̅_t)


    _0,t^∅ = _t - √(1-α̅_t)ϵ_θ(_t, t, ∅)/√(α̅_t)

Then, we collect MSE_t^h = _0,t^h-_0^2 and MSE_t^∅ = _0,t^∅-_0^2.



 §.§ Dynamical Noise Scale and Threshold


We use the unguided reference path to determine the appropriate end step t_e, i.e., noise scale, for each input. It is chosen based on the maximal absolute average difference between two MSEs, i.e., 

    t_e = max_t∑_d=1^D∑_i=1^H∑_j=1^W| MSE_t^h - MSE_t^∅|/D× H× W, t∈{t_i}_0^T.

Note that MSE_t^h, MSE_t^∅∈ℝ_+^D× H× W. To obtain the anomaly map H for segmentation, we average MSE_t^h for  t∈{t_i}_0^t_e for all modalities, i.e., H = ∑_d=1^D∑_t=0^t_eMSE_t^hd/t_e× D. Similarly, the segmentation threshold is individually determined for each input. We observe that the value of M = max∑_d=1^D∑_i=0^H∑_j=0^W| MSE_t^h - MSE_t^∅|/D× H× W for t∈{t_i}_0^T is roughly linearly related to the size of tumor. Intuitively, a larger anomalous region corresponds to a larger difference value M as more regions are removed. A smaller quantile is selected for a larger anomalous region to include more possible pixels. For a comprehensive understanding of our selection process, please refer to Algorithm 1, detailed in the supplementary material



§ EXPERIMENTS

We train and evaluate the proposed method on BraTS21 dataset <cit.>. BraTS21 dataset comprises of three-dimensional Magnetic Resonance brain images depicting subjects afflicted with a cerebral tumor, accompanied by pixel-wise annotations serving as ground truth labels. Each subject undergoes scanning through four distinct MR sequences, specifically T1-weighted, T2-weighted, FLAIR, and T1-weighted with contrast enhancement. Given our emphasis on a two-dimensional methodology, our analysis is confined to axial slices. There are 1,254 patients and we split the dataset into 939 patients for training, 63 patients for validation, 252 patients for testing. We randomly select 1,000 samples in validation set for parameter tuning and 10,000 samples in testing set for evaluation. For training, we stack all four modalities while only FLAIR and T2-weighted modalities are used in inference. For preprocessing, we normalize each slice by dividing 99 percentile foreground voxel intensity and then, pixels are scaled to the range of -1, 1. Finally, all samples are interpolated to 128× 128. Our model is trained as proposed in <cit.> with 2 Nvidia A100 80GB GPUs, and the hyperparameters are demonstrated in supplementary material. The backbone U-net is from the previous work <cit.>. 



 §.§ Results: Weakly-supervised Segmentation

We report pixel-level DICE score, intersection over union (IoU) and area under the precision-recall curve (AUPRC) in terms of foreground area in Table <ref>. The performance on 10,000 mixed data, comprising 4,656 unhealthy samples and 5,344 healthy samples, along with the performance on all 4,656 unhealthy data are reported separately. For our method, we present results for three setups: (i) DDPM forward (stochastic encoding): Eq. <ref> is used to add noise to inputs; (ii) DDIM forward (deterministic encoding): the inputs are noised by Eq. <ref>; (iii) tuned with labels: the hyperparameter is tuned with pixel-level labels for all inputs (non-dynamical) and noise is added using DDIM forward. For comparison methods, we report the results from AnoDDPM with Gaussian noise <cit.>, pure DDIM, DDIM with classifier <cit.> and DDIM classifier-free <cit.>. The first two methods are only trained on healthy data and the hyperparameters of all methods are tuned by using 1000 slices in validation set. Note that our methods do not require pixel-level labels for tunning. We adopt median filter <cit.> and connected component filter for postprocessing. The details are provided in supplementary material. We surpass the previous methods in a quantitative evaluation. The qualitative results are shown in Fig. <ref>. It shows that our method can enhance the signal strength of the anomalous regions.





 §.§ Ablation Study: Hyperparameter Selection

To validate the effectiveness of our chosen t_e, we randomly select 100 unhealthy samples from the validation set and calculate the AUPRC at different steps. In Fig. <ref> (a), we illustrate the change in mean and standard deviation of AUPRC over different steps. Notably, the maximal AUPRC is achieved at t = t_e, supporting the feasibility of our selection. Another hyperparameter under consideration is the guidance strength w. However, it does not significantly impact performance. The sensitivity of w with respect to DICE is depicted in Fig. <ref> (b).







§ CONCLUSION
 
In this paper, we propose a novel anomaly segmentation framework that eliminates the need for pixel-level labels in hyperparameter tuning. We utilize denoised inputs without label guidance as a reference for selecting the noise scale and dynamic threshold of the anomaly map. By aggregating anomaly maps from each forward step, we enhance the signal strength of anomaly regions, which improves the quality of anomaly maps. Our method surpasses the previous methods in terms of weakly-supervised segmentation on BraTS21 dataset.















splncs04




§ SUPPLEMENTARY MATERIAL



 §.§ A. Hyperparameter Settings

The parameters for the diffusion model used in our method are detailed in Table <ref>.




 §.§ B. Byproducts of MSE Inference

Two byproducts of MSE inference are classification and threshold selection. To determine the label of inputs _0 without the pixel-level labels, we utilize cosine similarity Cos∑_d=1^D∑_i=1^H∑_j=1^WMSE^h/D× H× W, ∑_d=1^D∑_i=1^H∑_j=1^WMSE^∅/D× H× W. Note that MSE^∅, MSE^h ∈ℝ_+^D× T× H × W. In Fig. <ref>,  we compare the MSE of 50 healthy samples and and 50 unhealthy samples. In the Fig. <ref> (a) and (b), we observe that MSE_t^h and MSE_t^∅ in healthy samples are more similar compared to those in unhealthy samples. Consequently, we determine the threshold Cos_thr that achieves maximal accuracy in the validation set.

In the Fig. <ref> (c), we depict the relationship between the size of the anomalous region (number of pixels) and the maximal difference value M. Notably, they exhibit a relatively linear relationship. This observation serves as a rough indication of the size of the anomalous region and can be utilized for the selection of the quantile Q without relying on pixel-level labels. We select the segmentation quantile Q^*∈a,b of the anomaly map H for each input _0 by Algorithm. 1. The input M_max is obtained from the validation set for scaling. A smaller quantile is selected for a larger anomalous region to include more possible pixels. In our case, we select a=0.90, b=0.98. Then, the predicted pixel-level labels is obtained as H ≥ Q^*.
boxruled






 §.§ C. Postprocessing

After we obtain the anomaly map, we apply a median filter with kernel size 5 to effectively enhance the performance. Then, we apply the connected component filter to remove the small connected components which is regarded as noise. We apply the same postprocessing to all methods for fair comparison.

