Zig-RiR: Zigzag RWKV-in-RWKV for Efficient Medical Image Segmentation.

Area of Science:

• Artificial Intelligence
• Medical Imaging
• Computer Vision

Background:

• Deep learning models, particularly those combining Convolutional Neural Networks (CNNs) and transformers, excel at medical image segmentation by extracting local and global features.
• However, the quadratic computational complexity of transformers limits their efficiency in processing high-resolution medical images.
• The Receptance Weighted Key Value (RWKV) model offers linear complexity for long-distance modeling, presenting an opportunity for more efficient medical image analysis.

Purpose of the Study:

• To address the limitations of existing models in medical image segmentation, this study explores the potential of RWKV for improved efficiency and accuracy.
• To propose a novel nested architecture, Zigzag RWKV-in-RWKV (Zig-RiR), designed to effectively capture both local and global features without sacrificing spatial continuity.
• To enhance the exploration of local features and maintain spatial coherence during the segmentation process.

Main Methods:

• A novel nested structure, Zigzag RWKV-in-RWKV (Zig-RiR), is proposed, comprising Outer and Inner RWKV blocks for hierarchical feature extraction.
• The Outer Zig-RWKV processes local patches as "visual sentences" to capture global context.
• The Inner Zig-RWKV further analyzes subpatches as "visual words" to extract detailed local information, complemented by a Zigzag-WKV attention mechanism to preserve spatial continuity.

Main Results:

• The Zig-RiR model demonstrates superior accuracy and efficiency in medical image segmentation across four diverse datasets (2D and 3D modalities).
• When tested on high-resolution 1024 × 1024 images, Zig-RiR achieved a 14.4x speed improvement and an 89.5% reduction in GPU memory usage compared to state-of-the-art methods.
• The proposed method effectively integrates global and local feature extraction while maintaining spatial continuity, outperforming existing approaches.

Conclusions:

• The Zigzag RWKV-in-RWKV (Zig-RiR) architecture offers a highly efficient and accurate solution for medical image segmentation.
• This novel approach overcomes the computational limitations of traditional transformer models, enabling effective analysis of high-resolution medical images.
• Zig-RiR represents a significant advancement in medical image segmentation, providing a faster and more memory-efficient alternative for clinical applications.