Computation-Efficient Hierarchical Tensor Train Convolutional Architectures for Diffusion-Based Medical Image Segmentation
摘要
Diffusion models have rapidly emerged as a leading approach in generative modeling, recognized for their strong theoretical framework and wide-ranging applications, from image synthesis and 3D scene rendering to medical image segmentation. In particular, diffusion-based approaches have achieved great success in generating accurate masks for medical image segmentation tasks. However, despite advancements, diffusion models remain computationally intensive, requiring large memory, long inference times, and complex operations, which lead to significant challenges for real-time deployment in clinical settings. In this paper, we introduce a Hierarchical Tensor Train (HTT) method that captures the structural characteristics of the denoising process within diffusion models to enable efficient compression and improved computational efficiency. With the HTT, the per-layer model size is reduced by up to 500 \(\times \) while maintaining compression rate and similar accuracy to conventional tensor train decomposition. Compared to traditional compression approaches that incur an additional \(\sim \) 2 GFLOPs per layer, the HTT reduces the per-layer computational load to approximately one-tenth of the original, leading to an overall computational reduction of about one-third across the entire model. This approach offers a promising path toward developing lightweight, scalable diffusion models suitable for diverse real-time applications, including those in clinical and resource-constrained environments.