Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV

要約

トランスフォーマーは医用画像の復元に革命をもたらしましたが、二次関数の複雑さにより、高解像度の医用画像への適用には依然として制限が生じています。
自然言語処理分野における最近の Receptance Weighted Key Value (RWKV) モデルの出現は、長いシーケンスを効率的に処理できるため、多くの注目を集めています。
その高度な設計を活用するために、私たちは医用画像復元用の初の RWKV ベース モデルである Restore-RWKV を提案します。
元の RWKV モデルは 1D シーケンス用に設計されているため、2D 医療画像の空間関係をモデル化するために 2 つの必要な変更を加えます。
まず、線形の計算複雑さでグローバルな依存関係を捕捉するリカレント WKV (Re-WKV) アテンション メカニズムを紹介します。
Re-WKV には、包括的な受容野の基本として双方向の注意と、さまざまなスキャン方向からの 2D 依存関係を効果的にモデル化するための反復的注意が組み込まれています。
2 番目に、幅広いコンテキスト範囲にわたって全方向からトークンをシフトすることでローカルの依存関係を強化する全方向トークン シフト (Omni-Shift) レイヤーを開発します。
これらの適応により、提案された Restore-RWKV は医療画像復元のための効率的かつ効果的なモデルになります。
Restore-RWKV の軽量版でも、パラメータが 116 万しかなく、既存の最先端 (SOTA) 手法と比較して同等、またはそれ以上の結果を達成します。
広範な実験により、結果として得られる Restore-RWKV は、PET 画像合成、CT 画像のノイズ除去、MRI 画像の超解像度、オールインワンの医用画像復元など、さまざまな医用画像復元タスクにわたって SOTA パフォーマンスを達成することが実証されています。
コードは https://github.com/Yaziwel/Restore-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 the Receptance Weighted Key Value (RWKV) model in the natural language processing field has attracted much attention due to its ability to 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 medical 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. Even a lightweight variant of Restore-RWKV, with only 1.16 million parameters, achieves comparable or even superior results compared to existing state-of-the-art (SOTA) methods. Extensive experiments demonstrate that the resulting Restore-RWKV achieves SOTA performance across a range of medical image restoration tasks, including PET image synthesis, CT image denoising, MRI image super-resolution, and all-in-one medical image restoration. Code is available at: https://github.com/Yaziwel/Restore-RWKV.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google