RF-HiT: Rectified Flow Hierarchical Transformer for General Medical Image Segmentation

1. Problem Statement

Medical image segmentation requires pixel-level precision and the ability to model both long-range contextual dependencies and fine-grained boundary details. Existing approaches face significant trade-offs:

• CNNs (e.g., U-Net): Limited by local receptive fields, making it difficult to capture global spatial dependencies.
• Standard Transformers (e.g., ViT): Capable of global attention but suffer from quadratic computational complexity ($O(n^2)$), which is prohibitive for high-resolution medical images.
• Diffusion Models: Offer state-of-the-art generative capabilities and uncertainty estimation but require hundreds of inference steps, leading to prohibitive latency for real-time clinical deployment.

The authors aim to bridge this gap by creating a model that achieves linear computational complexity, fast inference (few steps), and high accuracy without sacrificing global context or boundary precision.

2. Methodology: RF-HiT Architecture

The proposed RF-HiT (Rectified Flow Hierarchical Transformer) is a generative framework that integrates a Rectified Flow model with an efficient Hourglass Transformer backbone.

A. Core Framework: Rectified Flow

Instead of the stochastic forward/reverse process of standard Diffusion Probabilistic Models (DDPMs), RF-HiT utilizes Rectified Flow.

• Mechanism: It learns an Ordinary Differential Equation (ODE) that transports noise ($X_0$) to data ($X_1$) along straight-line trajectories.
• Benefit: This allows for high-quality generation with significantly fewer discretization steps (as few as 3–15) compared to the hundreds required by diffusion models, drastically reducing inference time.

B. Backbone: Efficient Hourglass Transformer (HDiT)

The velocity network ($v_\theta$) is built upon the Hourglass Diffusion Transformer (HDiT) architecture, adapted for medical segmentation:

• Hierarchical Structure: Follows an encoder-decoder (U-Net) design with a bottleneck.
• Attention Mechanism:
  • High-Resolution Levels: Uses Neighborhood Self-Attention to capture local features with linear complexity ($O(n)$).
  • Bottleneck Level: Uses Global Self-Attention to capture comprehensive global context.
• Efficiency: This hybrid approach avoids the quadratic cost of full global attention while maintaining global reasoning capabilities.

C. Dual-Pathway Design & Feature Conditioning

A novel aspect of RF-HiT is its Dual-Pathway structure:

1. Main Flow Model (HFM): Processes the concatenation of the noisy segmentation map and the original image to predict the velocity field.
2. Hierarchical Feature Encoder (HFE): A separate encoder that processes the original image to extract multi-scale anatomical representations.
3. Learnable Linear Interpolation Fusion:
   • Features from the HFE are fused with the corresponding levels of the HFM encoder.
   • Instead of simple addition, the model uses a learnable scalar parameter ($\alpha_l$) to interpolate between the main features and conditioning features:
     $$\tilde{F}^{(l)} = (1 - \alpha_l)F^{(l)} + \alpha_l C^{(l)}$$
   • This mechanism adaptively injects multi-scale anatomical priors into the generative process, guiding the velocity prediction for precise boundary delineation.

3. Key Contributions

1. Novel Architecture: Introduction of RF-HiT, the first framework to combine Rectified Flow with a hierarchical Hourglass Transformer for medical segmentation.
2. Dual-Pathway Conditioning: A specific design featuring a Hierarchical Feature Encoder (HFE) and a learnable linear interpolation fusion mechanism. This effectively injects anatomical context without increasing computational overhead significantly.
3. Extreme Efficiency: The model achieves linear complexity and requires only 13.6 million parameters and 10.14 GFLOPs.
4. Fast Inference: The model produces high-quality segmentation masks in just 3 sampling steps, making it viable for real-time clinical use.

4. Experimental Results

The model was evaluated on two challenging benchmarks: ACDC (Cardiac MRI) and BraTS 2021 (Brain Tumor MRI).

A. Performance on ACDC (Cardiac Segmentation)

• Mean Dice Score: 91.27% (Ranking 2nd overall, just behind the much larger nnFormer at 92.06%).
• Efficiency: Achieves this with 13.6M parameters compared to nnFormer's 149.6M parameters.
• Specifics: Achieved state-of-the-art performance on the Myocardium (89.74% DSC).

B. Performance on BraTS 2021 (Brain Tumor)

• Mean Dice Score: 87.40% across tumor sub-regions (WT, TC, ET).
• Comparison: While slightly lower in absolute Dice than massive 3D models like MedSegDiff-V2 (90.8%), RF-HiT achieves this with a fraction of the computational cost.
  • MedSegDiff-V2: 46M params, 983 GFLOPs.
  • RF-HiT: 13.6M params, 10.14 GFLOPs.
• Inference Speed: Operates in 3 steps, whereas diffusion baselines often require 100+ steps.

C. Ablation Studies

• HFE Impact: Removing the Hierarchical Feature Encoder dropped the mean Dice from 91.27% to 90.69%, proving the value of explicit multi-scale anatomical conditioning.
• Fusion Mechanism: The proposed Learnable Linear Interpolation outperformed simple feature addition (91.27% vs. 91.09%), demonstrating the importance of adaptive weighting.

5. Significance and Conclusion

RF-HiT represents a significant advancement in efficient deep learning for medical imaging.

• Clinical Viability: By reducing inference time to a few steps and lowering computational requirements, it addresses the primary bottleneck preventing generative models from being deployed in real-time clinical workflows.
• Paradigm Shift: It demonstrates that Rectified Flow is a superior alternative to standard Diffusion Models for segmentation tasks, offering comparable quality with massive speed gains.
• Scalability: The architecture proves that high performance does not require massive parameter counts; a compact, linear-complexity design can outperform or match significantly heavier architectures.

In summary, RF-HiT offers a robust, computationally efficient foundation for automated medical image segmentation, balancing the need for global context, boundary precision, and rapid inference.