CTS: A Consistency-Based Medical Image Segmentation Model

The field of medical image segmentation has always been a hot research direction within the image segmentation domain. Unlike traditional segmentation methods[15, 22], utilizing generative models for image segmentation[20] can also achieve good results. Since diffusion models[17, 10] are a type of generative model that samples from Gaussian noise, the images they generate possess strong noise resistance and smoothness. Consequently, an increasing number of studies are leveraging diffusion models to tackle the non-generative issues of different images. Researchers use masks as the target for generative model sampling, while also incorporating constraints in the generative models to guide the direction of model generation. However, due to the need for extensive resampling during training and prediction, the issue of low computational efficiency in diffusion models urgently needs to be addressed. Consistency model[18] transform multiple samplings into a single sampling by constructing a unique solution by ODE, significantly reducing the time consumed during the sampling process. Moreover, while reducing the number of samplings, consistency models also ensure the effectiveness of the samples. Compared to DDPMs[10], consistency model represent a superior generative paradigm, yet studies applying this model to the field of medical image segmentation are currently lacking. Therefore, this paper proposes constructing a medical image segmentation model based on the consistency model, and designing a loss function according to the segmentation loss and consistency training loss, enabling end-to-end training of the model. CTS code can be obtained in https://github.com/LanHEU/CTS .

The specific contributions of this text are as follows:

• •

  A medical image segmentation model based on a consistency model has been constructed, featuring a newly designed joint loss function.

• •

  During the decoding phase, multi-scale feature supervision signals are utilized to guide the model’s convergence direction.

In this section, we briefly describe the existing lines of research relevant to our work. Diffusion models have been applied to many fields, such as sequence modeling [12, 5], speech processing[14], computer vision[16, 9] to computed tomography (CT) scanning and magnetic resonance imaging (MRI). In computer vision, to reduce the number of sampling times, many methods have made great efforts. There are also some sampling algorithms tailored for conditional generation, such as without classifier guidance[11] or with classifier guidance[6]. Image segmentation is an important task in computer vision, which studies simplifying the complexity of an image by decomposing it into multiple meaningful image segments[2, 8]. Due to the time, cost, and expertise required[1, 4], the number of images and labels for medical image segmentation is limited. For this reason, diffusion models, by synthesizing labeled data and eliminating the need for pixel-level labeled data, have become a promising method in image segmentation research. BrainSPADE[7] proposed a generative model for synthesizing labeled brain MRI images, which can be used to train segmentation models. However, diffusion models in medical image segmentation face issues such as a high number of sampling times and long prediction times.

This paper aims to fully leverage the advantages of sampling once with a consistency model, while retaining the benefits of the segmentation model. In consistency model[18], the method of directly training a consistency model is referred to as consistency training loss, which is the origin of the ’CT’ in the name of ’CTS’. The specific process is shown in Fig1.

Similar to the consistency model, the basic framework of this paper includes two parts: model $M$ and target model $TM$. The model’s sampling begins with the mask $x^{m}$ of each image, inputting the corresponding data $x^{d}$as a supervisory signal. Initialize the parameters of the two models, and copy the parameters from $M$ to $TM$.

Step 1. The input to the model is the mask $x^{m}$, and noise $z_{n}$, sampled from a Gaussian distribution at step n, is added to the tensor: $x^{m}_{n}=x^{m}+z_{n}$.

Step 2. Simultaneously, based on the time period $t$, obtain: Model average learning moving parameters: $c_{in}$,$c_{out}$,$c_{skip}$

Step 3. The final input to model $M$ is: $x^{m}_{in}=c_{in}*x^{m}_{n}$

Step 4. Generate multi-scale signals using $\bigcup x_{i}^{d},\hat{y}=h^{T}\left(x^{d}\right)$, where $i\in\left(1,2,...\right)$

Step 5. Here, the feature signal $\bigcup x_{i}^{d}$ is incorporated into the UNet model: $y_{out}^{m}=\ g^{T}\left(x_{in}^{m},\bigcup x_{i}^{d},t\right)$

Step 6. The output of the final model M is: $y_{n}^{m}=c_{out}\ast y_{out}^{m}+\ c_{skip}\ast x_{n}^{m}$

Step 7. Utilize the normal sampling method to obtain the noise$z_{n+1}$ of $\left(n+1\right)$th sampling from the Gaussian distribution. $x_{n+1}^{m}=f_{\mathcal{U}}\left(x^{m},x_{n}^{m},n,n+1\right)$

Step 8: Use $x_{n+1}^{m}$ as the input for model $TM$, and repeat the above Step 2-5. And the output is $y_{n+1}^{m}=c_{out}\ g^{TM}\left(x_{in}^{m},x^{d},t\right)\ +\ c_{skip}\ast x_{n+1}^{m}$

The consistency training segmentation loss is as follows: $\mathcal{L}_{CT}=yn+1m-ynm2$ To expedite the convergence speed and training outcomes, training is conducted on structures generated by multi-scale signals: $\mathcal{L}_{S}=y-xm2$ Overall Loss Function: $\mathcal{L}_{CTS}=\mathcal{L}_{CT}+\alpha\mathcal{L}_{S}$, where $\alpha$ is hyperparameter.

Step 9: Update the $TM$ model parameters; the $TM$ model update adheres to the learning rate: $\theta^{TM}\leftarrow\ stopgrad\left(\mu\left(k\right)\theta^{TM}+\left(1-\mu\left(k\right)\right)\theta^{M}\right)$.

The pseudocode of the CTS algorithm is shown in Alg1.

Multi-scale Feature Supervision Signal. The process of integrating multi-scale feature supervision signals $\bigcup x_{i}^{d}$ is shown in Fig1(a). The decoder stage of the image data encoding network progressively generates feature maps of each size, and combines them with the corresponding supervision signals. In the decoder stage of the image data encoding network, a corresponding supervision signal $x_{i}^{d}$ is gradually generated for each size feature map. The supervision signal $x_{i}^{d}$, as shown in Fig1(b). This process is integrated into the $M$ model through a channel attention mechanism, achieving the addition of multi-scale supervision signals. During the decoder stage of the image data encoding network, these feature maps contain information of various scales, which can assist the model in better understanding the details and contextual information of the image. To better integrate the supervision signals and feature maps, a channel attention mechanism is employed. This automatically learns the importance weights of each channel, thereby making better use of the information from the supervision signals.

This section demonstrates through experimentation the advantages of CTS in medical image segmentation. We started with a thorough comparison of existing alternatives, followed by additional analysis to dissect the reasons behind CTS’s success.

Datasets. We conducted experiments on medical tasks in two different image modalities: MRI image segmentation of brain tumors, and ultrasound image segmentation of thyroid nodules and liver tumor segmentation on the BraTs-2021 dataset[3] as well as in SEHPI datasets[21]. This paper utilized anisotropic diffusion filtering[13], while also removing Poisson noise from medical images, preserving more edge information and effective feature structures. Consequently, this further improved the model’s performance.

Experiment Details. We utilized a $4\times$ UNet. In the testing phase, we employed a single diffusion step for inference, which is significantly smaller than most previous studies. All experiments were implemented using the PyTorch platform and executed on one GTX 4090. All images were uniformly resized to 256×256 pixels. Training was conducted in an end-to-end manner using the AdamW optimizer. $batchsize=8$. The learning rate was initially set to $1\times 10^{-4}$.

Main result. We compared the SOTA segmentation methods proposed for each task with general medical image segmentation methods. The main results are presented in Table 1. Part of the relevant results originates from the work[19]. In our experiments, we trained on each dataset for 700,000 iterations, with the specific training duration being one month. CTS-nM indicates that the model did not use multiscale feature supervision signals. CTM-M denotes the use of multiscale feature supervision signals. For a fair comparison, the Meg method incorporates a Fourier filter, while CTM-FM includes the FFTP structure mentioned same with MedSegDiff. Detailed results can be seen in Tab1. Visualization results are shown in Fig2. The results reveal that CTS-nM can surpass most methods, demonstrating that consistency model can not only reduce the number of samplings but also enhance effectiveness. The performance of the CTM-M model, with the addition of multi-scale feature signals, is further improved. CTM-FM also proves that Fourier filtering can further enhance performance. All methods under CTS were tested with a single sampling, averaging 1.9s. This significantly reduces the time compared to other models. Therefore, the CTS model can guarantee model effectiveness while accelerating sampling.

In Fig4, the convergence process of the model is shown. The trend of the loss value during training. The changing trends of IoU and Dice metrics for model parameters saved at different times on the test set. It can be observed that, as the training time increases, the loss region during training becomes more leveled, but the test results do not show saturation, with a significant growth margin. This is likely related to the strong learning and representational capabilities of the consistency model. The findings of this paper are in agreement with the conclusions of the consistency studies.

Ablation experiment Fig4 compares the convergence speeds between CTS-nM and CTS-M. Since models typically enter a smoothing phase after exceeding 10,000 rounds, only the experimental results of 20,000 rounds are compared here. It is evident that the inclusion of multi-scale feature signals significantly accelerates the convergence speed of the models.Tab2 compares the ablation results between two datasets. Here, to observe the result, each experiment was trained with 500,000 samples. It is evident that incorporating multi-scale signals yields better results. Additionally, FFTP can also enhance performance. Adding either element individually does not significantly differ in improving the model, possibly because both involve modifications to the way supervisory signals are added, with the difference lying in the method. This further indicates that adding better supervisory signals can greatly impact the model’s outcomes. Eventually, the model that incorporated both methods was used to train the overall model.

This paper first establishes a medical image segmentation method based on a consistency model CTS. It not only yields better results but also significantly reduces prediction time. Moreover, by constructing multi-scale feature supervision signals, the training convergence speed is accelerated. Meanwhile, medical image data processed with anisotropic edge enhancement filters can achieve improved outcomes. However, there are still some shortcomings. In our experiments, due to a lack of GTX 4090, we did not train as many as one million rounds as initially planned in the Consistency[18]. Furthermore, we observed that increasing the number of training rounds did not lead to saturation in results, leading us to infer that more training rounds could potentially enhance the model’s performance.