Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV
Transformers have revolutionized medical image restoration, but the quadratic complexity still poses limitations for their application to high-resolution medical images. The recent advent of RWKV in the NLP field has attracted much attention as it can process long sequences efficiently. To leverage...
Gespeichert in:
Hauptverfasser: | , , , , |
---|---|
Format: | Artikel |
Sprache: | eng |
Schlagworte: | |
Online-Zugang: | Volltext bestellen |
Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
Zusammenfassung: | Transformers have revolutionized medical image restoration, but the quadratic
complexity still poses limitations for their application to high-resolution
medical images. The recent advent of RWKV in the NLP field has attracted much
attention as it can process long sequences efficiently. To leverage its
advanced design, we propose Restore-RWKV, the first RWKV-based model for
medical image restoration. Since the original RWKV model is designed for 1D
sequences, we make two necessary modifications for modeling spatial relations
in 2D images. First, we present a recurrent WKV (Re-WKV) attention mechanism
that captures global dependencies with linear computational complexity. Re-WKV
incorporates bidirectional attention as basic for a global receptive field and
recurrent attention to effectively model 2D dependencies from various scan
directions. Second, we develop an omnidirectional token shift (Omni-Shift)
layer that enhances local dependencies by shifting tokens from all directions
and across a wide context range. These adaptations make the proposed
Restore-RWKV an efficient and effective model for medical image restoration.
Extensive experiments demonstrate that Restore-RWKV achieves superior
performance across various medical image restoration tasks, including MRI image
super-resolution, CT image denoising, PET image synthesis, and all-in-one
medical image restoration. Code is available at:
\href{https://github.com/Yaziwel/Restore-RWKV.git}{https://github.com/Yaziwel/Restore-RWKV}. |
---|---|
DOI: | 10.48550/arxiv.2407.11087 |