[論文レビュー] Attention for Image Registration (AiR): an unsupervised Transformer approach

Zihao Wang, Hervé Delingette|arXiv (Cornell University)|May 5, 2021

Advanced Neural Network Applications参考文献 17被引用数 4

ひとこと要約

本稿では、変形可能な画像登錧をシーケンス・トゥ・シーケンス翻訳タスクとして扱う、自己教師付きの変換器ベースのフレームワーク、AiRを提案する。畳み込み型バックボーンを用いずにマルチスケール自己注意メカニズムを活用することで、ベンチマークデータセット上で最先端の性能を達成し、単一スケールの変換器およびCNNベースの手法に比べ、変形精度と構造的類似性において優れている。特に大規模な変形に対して顕著な優位性を示す。

ABSTRACT

Image registration is a crucial task in signal processing, but it often encounters issues with stability and efficiency. Non-learning registration approaches rely on optimizing similarity metrics between fixed and moving images, which can be expensive in terms of time and space complexity. This problem can be exacerbated when the images are large or there are significant deformations between them. Recently, deep learning, specifically convolutional neural network (CNN)-based methods, have been explored as an effective solution to the weaknesses of non-learning approaches. To further advance learning approaches in image registration, we introduce an attention mechanism in the deformable image registration problem. Our proposed approach is based on a Transformer framework called AiR, which can be efficiently trained on GPGPU devices. We treat the image registration problem as a language translation task and use the Transformer to learn the deformation field. The method learns an unsupervised generated deformation map and is tested on two benchmark datasets. In summary, our approach shows promising effectiveness in addressing stability and efficiency issues in image registration tasks. The source code of AiR is available on Github.

研究の動機と目的

• 類似度メトリック最適化に基づく従来の学習なし画像登錧手法の不安定さと高い計算コストを解消すること。
• CNNベースのディープラーニング手法が長距離依存関係およびグローバル構造的変形を捉えることの制限を克服すること。
• 自然言語処理分野で優位を占めてきた変換器アーキテクチャが、アノテーション済み変形フィールドが存在しないエンド・ツー・エンドの自己教師付き画像登錧に応用可能かどうかを検討すること。
• 複雑な変形のモデリングを向上させるために、異なる受容 field サイズに対応する特徴表現を強化するマルチスケール注意メカニズムを開発すること。
• 自己注意メカニズムが弱教師ありの形で変形フィールドを効果的に学習可能であり、アノテーション付き学習データの必要性を回避できることを実証すること。

提案手法

• フレームワークは、固定画像および移動画像のパッチを入力トークンとして扱い、画像登錧をシーケンス・トゥ・シーケンス翻訳タスクとして定式化する。
• 変換器エンコーダは、固定画像パッチをマルチヘッド自己注意により処理し、入力画像のメモリ表現を学習する。
• 変換器デコーダは、エンコーダの出力メモリ特徴量と移動画像パッチの自己注意を併用して変形フィールドを予測する。
• 変形フィールドを用いて、微分可能な空間サンプリングにより移動画像をワープし、再サンプリング操作を介したバックプロパゲーションを可能にする。
• 大規模な変形に対する性能向上を図るため、マルチスケール並列注意（MAPT）モジュールを導入し、異なる受容 field サイズにおける特徴表現を強化する。
• モデルは、ワープされた移動画像と固定画像の画素単位L2距離に基づく再構成損失を用いて、自己教師付きで学習する。

実験結果

リサーチクエスチョン

• RQ1 畳み込み特徴抽出器を一切用いない純粋な変換器アーキテクチャが、画像登錧のための変形フィールドを効果的に学習できるか？
• RQ2 複雑な変形をモデリングする際、単一スケールの変換器とマルチスケール注意メカニズムの性能はどのように異なるか？
• RQ3 アノテーション済み変形ラベルが存在しない自己教師付き変換器ベース手法が、競争力のある登錧精度を達成できるか？
• RQ4 画像登錧タスクにおいて、変換器の自己注意メカニズムがCNNを上回ってグローバル構造的関係を捉えることができるか？
• RQ5 マルチスケール特徴学習の影響は、生成された変形マップおよび画像再構成の品質にどのように現れるか？

主な発見

• マルチスケール並列注意変換器（MAPT-AiR）は、MNISTデータセットにおいてすべての指標で最良の性能を示し、平均PSNRは17.311 ± 2.522、Smooth DICEスコアは0.827 ± 0.034を達成した。
• 単一スケールAiRは、CNNベースのDIRNets（16.966 ± 2.897 vs. 12.941 ± 2.218）をPSNRで上回ったが、マルチスケールバージョンに比べて性能が劣った。
• MAPT-AiRは、『9』『7』『8』などの大規模な変形を的確に捉え、単一AiRはこのような複雑な形状のアライメントに失敗した。
• MNISTでは平均MSEが0.027 ± 0.015を記録し、画素レベルの誤差が最小限に抑えられた高い再構成忠実度を示した。
• 強力な構造的アライメントにもかかわらず、すべての変換器ベース手法は、ワープ画像において幾何的詳細欠損およびノイズを示し、画素レベルの精度に限界があることが示された。
• 結果から、変換器がグローバルな変形モデリングに有効である一方で、細かなテクスチャおよびエッジ保持においてまだCNNに劣ることが確認された。