Tree-NET: Enhancing Medical Image Segmentation Through Efficient Low-Level Feature Training

TL;DR

Tree-NET is a new framework that improves medical image segmentation by using bottleneck feature supervision across multiple training phases, reducing computational costs while maintaining or enhancing accuracy.

Contribution

It introduces a novel multi-phase training framework with bottleneck supervision, significantly reducing FLOPs and memory usage without sacrificing segmentation performance.

Findings

Reduces FLOPs by 4 to 13 times

Decreases memory usage

Maintains or improves accuracy

Abstract

This paper introduces Tree-NET, a novel framework for medical image segmentation that leverages bottleneck feature supervision to enhance both segmentation accuracy and computational efficiency. While previous studies have employed bottleneck feature supervision, their applications have largely been limited to the training phase, offering no computational benefits during training or evaluation. To the best of our knowledge, this study is the first to propose a framework that incorporates two additional training phases for segmentation models, utilizing bottleneck features at both input and output stages. This approach significantly improves computational performance by reducing input and output dimensions with a negligible addition to parameter count, without compromising accuracy. Tree-NET features a three-layer architecture comprising Encoder-Net and Decoder-Net, which are autoencoders designed to compress input and label data, respectively, and Bridge-Net, a segmentation framework that supervises the bottleneck features. By focusing on dense, compressed representations, Tree-NET enhances operational efficiency and can be seamlessly integrated into existing segmentation models without altering their internal structures or increasing model size. We evaluate Tree-NET on two critical segmentation tasks -- skin lesion and polyp segmentation -- using various backbone models, including U-NET variants and Polyp-PVT. Experimental results demonstrate that Tree-NET reduces FLOPs by a factor of 4 to 13 and decreases memory usage, while achieving comparable or superior accuracy compared to the original architectures. These findings underscore Tree-NET's potential as a robust and efficient solution for medical image segmentation.