A Novel Approach for Integrating Mamba and Diffusion in 3D Medical Imaging

巨大な3次元のジグソーパズルを完成させようとしている場面を想像してみてください。しかし、誰かが絵の大きな塊を切り取ってしまったため、手元にはバラバラになった断片しか残っていません。これが、MRIやCTのような3D医療スキャンを使用する医師たちが日々直面している現実です。患者が動いたり、装置に不具合が生じたり、あるいはスキャン自体が全身を捉えきれなかったりすることで、データに「穴」が開いてしまうことがあります。これらの欠落した部分を埋めることは非常に困難です。なぜなら人体は複雑であり、もし推測を誤れば、架空の腫瘍を作り出したり、逆に実在する腫瘍を見逃したりしてしまう可能性があるからです。長年、科学者たちはこの問題を解決するために、さまざまな種類の「スマート」なコンピュータプログラムを利用しようと試みてきました。あるプログラムは、隣接するピクセルだけを見る「局所的な探偵」（CNN）のようなものです。また別のプログラムは、全体像を一度に理解しようとする「グローバルなプランナー」のようなものですが、膨大なデータ量に圧倒されてしまいます（Transformer）。さらに、欠落した部分がどのようになるかをランダムに推測する「芸術家」のように振る舞うプログラム（拡散モデル）もありますが、これらは解剖学的な構造を誤ったり、自分の推測に対してどの程度確信を持っているかを提示できなかったりすることがあります。大きな問いは、「高速で、全身を理解し、リアルな細部を生成でき、かつ自分がいつ推測を行っているのかを正確に把握できる単一のシステムを構築できるか？」ということです。

本論文では、これら4つの異なる技術を一つの「スーパーチーム」へと統合することで、その問いに答えようとする新システム「MedMamba3D」を紹介します。これは、メンバーそれぞれが特定の重要な役割を担う建設作業員のようなものです。まず、Mamba Encoderがあります。これは、あらゆる角度から3Dボリュームを読み取る「スマートスキャナー」です。従来のメソッドは巨大な3Dブロックを見ると動作が遅くなり混乱してしまいますが、このスキャナーは極めて効率的です。データを直線的に処理（線形計算量）しながらも、体の異なる部位間の長距離のつながりを記憶しています。それは、膨大な図書室のすべての通路を歩き回ることなく、どの本の場所にあるかを瞬時に思い出す司書のようなものです。

次に、作業チームは2つの専門的なブランチに分かれます。第一のブランチは、Deterministic U-Net、すなわち「建築家」です。その役割は、体の基本的な構造が理にかなっていることを保証することです。このブランチは、部分的なスキャンを観察し、欠落した部分の解剖学的に正しい堅牢なスケルトン（骨組み）を構築します。これにより、心臓が突然肺の中に現れるといった事態を防ぎます。第二のブランチは、Conditional Diffusion Model、すなわち「芸術家」です。建築家が構造を構築する一方で、このブッチは細部——テクスチャ、エッジ、そしてスキャンをリアルに見せる高周波ノイズ——を付け加えます。このブランチは、情報の文脈に沿った細部が確実に適合するように、スマートスキャナーからの情報をガイドとして使用します。

しかし、最も重要な部分は、**Fusion Module（融合モジュール）**です。多くの従来の手法では、これらの異なるアプローチを単に平均化するだけであり、それはまるで絵具を混ぜ合わせて、うまくいくことを祈るようなものでした。MedMamba3Dはよりスマートです。これは「信頼度の判定役」として機能します。建築家と芸術家の両方に、「この特定の地点について、あなたたちはどの程度確信を持っていますか？」と問いかけます。もし建築家は自信を持っているが芸術家が確信を持てていない場合、最終的な画像はその地点において建築家の推測に強く傾斜します。その逆も同様です。この「逆分散重み付け（inverse variance weighting）」と呼ばれるプロセスにより、最終的な結果は構造的な整合性とリアルな細部の完璧なブレンドとなり、同時に、コンピュータがどこで確信を持てていないかを示す「懸念マップ（worry map）」を医師に提示します。

著者らは、脳腫瘍スキャン（BraTS 2021）、膝のMRIスキャン（fastMRI）、胸部X線（MIMIC-CXR）を含む実際の医療データセットを用いてこの新システムをテストしました。その結果、MedMamba3Dは大きな前進を遂げたことが示されました。脳腫瘍のテストにおいて、システムは28.34 ± 0.26のPSNR（画像の鮮明度の指標）を達成し、次点の優れた手法を明確な差で上回りました。また、腫瘍のセグメンテーション（腫瘍の正確な形状を見つけること）の精度も、従来の手法と比較して8.9%向上しました。さらに驚くべきことに、データの90%が欠落している状態でも、システムは堅牢性を維持し、24.6 ± 0.4のPSNRを伴う利用可能な再構成画像を生成しました。また、このシステムは他の高品質な手法よりもはるかに高速でした。一部の拡散モデルが画像の生成に3秒以上かかる一方で、MedMBamba3Dはわずか0.52秒で完了しました。

本論文は、単一の手法による解決策に依存することに対して明確に異議を唱えています。標準的なCNNである「建築家」のみを使用すると、細部が限られたぼやけた画像になり、「芸術家」（純粋な拡散モデル）のみを使用すると、解剖学的に不可能な構造を作成したり、不確実性を推定できなかったりするリスクがあることを示唆しています。著者らは、どちらのアプローチ単独でも不十分であることを発見しました。魔法が起きたのは、それらが不確実性を考慮した融合プロセスと共に組み合わされた時だけでした。また、Transformerが3Dボリュームに対して最良の解決策であるという考えについても、高解像度の3Dスキャンに対して計算コストが急激に増大（二次計算量）するため実用的ではないと指摘し、彼らのMambaベースのアプローチの方が効率的（線形計算量）であることを述べています。

これらの実験において、新しいモデルは単に見栄えが良いだけでなく、より信頼性の高いものでした。システムは0.08という較正スコア（ECE）を達成しました。これは、その信頼レベルが現実と非常に近いことを意味しており、標準的な3D U-Netの0.21というスコアと比較して大幅な改善となっています。著者らは、この手法にはペアとなるデータが必要であり、メモリ消費量も多いものの、状態空間モデリング、拡散プロセス、および不確実性推定を統合することが、単に正確であるだけでなく、医師に対して「いつ画像を信頼し、いつ慎重になるべきか」を伝えることができる、臨床使用においてより安全なツールを生み出すことを実証したと結論付けています。