Adaptive spectral methods improve 3d medical image segmentation efficiency

Learning Spectral Allocation: A Fractional Diffusion Framework for Adaptive Volumetric Segmentation

Computer Vision and Pattern RecognitionArtificial Intelligence

Summary

Segmenting 3D medical images is computationally intensive, and the authors found a way to let the model decide how much spectral processing each part needs. They created a special operator, called FHEAT, that can smoothly range from doing nothing to applying a sharp smoothing filter, controlled by learnable parameters. Integrating this into a lightweight network, their model learns to skip unnecessary computation steps, reducing effort while keeping or improving accuracy on common medical imaging tasks. This approach turns a manual design choice into something the training process figures out automatically.

What this means in practice

  • For medical imaging engineers: Design segmentation systems that automatically optimize spectral processing stages for accuracy and efficiency in volumetric scans.
  • For machine learning practitioners: Incorporate learnable spectral operators into training pipelines to reduce computation without manual stage design in volumetric data tasks.

Authors

Yi-Hui Shen, Tie-Qiang Li

Abstract

We address adaptive computation in 3D medical image segmentation: instead of designing another backbone, we ask how much spectral mixing each network stage needs and let optimization answer. We derive FHEAT, a two-parameter operator family, from the discrete cosine transform (DCT) solution of a fractional heat equation. A fractional order alpha and a diffusion strength D govern the operator, and at D=0 it is exactly the identity. Reparametrized by the semigroup time tau = D*alpha, same-resolution instances compose exactly, so any distribution of diffusion across same-resolution stages amounts to a single Sobolev-type regularizer of learned strength. This identity limit lets the optimizer of each layer, not the designer, decide whether global mixing is needed and how sharp it should be. We instantiate FHEAT in a lightweight U-shaped architecture (Light-UNETR) paired with a Kolmogorov-Arnold mixer (KAN3D) with adaptive rational activations, yielding FHEAT-Seg. At 5% to 20% label rates on three public benchmarks, training produces gradient-driven spectral sparsification: seven of the eight stage-level operators drive D to zero, and the survivor saturates at the sharpest low-pass (alpha ~ 0.9) in the decoder layer feeding the semi-supervised attention map. The retired layers become exact identity shortcuts at inference, cutting FLOPs from 4.29G to 0.90G (a 79% drop) at 0.975M parameters. Under a standard semi-supervised protocol, FHEAT-Seg reaches Dice scores of 90.47% (left atrium), 78.79% (Pancreas-CT), and 81.90% (BraTS 2019), ahead of five semi-supervised methods and the Light-UNETR baseline. The large variant also surpasses Light-UNETR-L under full supervision (Dice 93.09%, 85.11%, and 87.19%) with 2.851M parameters and 55.75G FLOPs. These results suggest that the allocation of spectral computation is a learnable property of optimization dynamics, not a manual design commitment.