Now You Have My Healthy Attention: A U-DiT for Brain-MRI Inpainting

あなたは、人間の脳の巨大な3次元ジグソーパズルを解こうとしていると想像してください。しかし、誰かがその真ん中から大きな塊を抜き取り、空白のグレーの空洞に置き換えてしまいました。これは、医師が患者のMRIスキャンを見る際に直面する課題そのものです。腫瘍（およびそれを除去するための手術）によって画像に穴が開いてしまい、その下にあるはずの「健康な」脳がどのような姿であったかを見ることが困難になります。健康な状態の明確な画像がなければ、腫瘍がどれほどの大きさだったかを測定したり、次の治療ステップを計画したりすることが難しくなります。科学者たちは、これらの欠落した脳のパーツを魔法のように「埋める」ことができるコンピュータプログラムを構築しようと試みてきました。欠損のない、患者自身の解剖学的構造の完璧なコピーを作成することが目標です。目的は、単に見た目を綺麗にすることではなく、医師が疾患をより深く理解するための数学的に正確な地図を作成することです。

あなたがこれから読む論文「Now You Have My Healthy Attention」は、この特定のパズルを解くために設計された新しいコンピュータプログラムについて記述しています。著者であるダニロ・ダネゼ、アンジェラ・ロンバルディ、トマソ・ディ・ノイアは、熟練の修復師のように振る舞うスマートなAIを構築しました。このAIは、単にランダムに推測するのではなく、欠落した部分に何が入るべきかを判断するために2つの巧妙なトリックを使用します。第一に、それは厳格なルールを持っています。それは、欠落している空間に何が入るかを判断するために、すでに目に見えている脳の健康な部分のみを見ることが許可されています。他の空の場所を見ることは拒否されます。これにより、混乱したり偽のディテールを作り出したりすることを防ぎます。第二に、人間の脳はおよそ蝶のように対称的であるという事実を利用します。もし左側が欠けていれば、AIは右側（健康な鏡像）を見て、欠けているピースがどのような形であるべきかを確認します。これら2つのアイデアを組み合わせることで、コンピュータは欠落した脳組織を驚くほど正確に埋めることができます。

問題点：穴の開いた脳

脳のMRIを、患者の頭部の3D写真と考えてみてください。患者に腫瘍がある場合、その腫瘍は暗い点や明るい点として現れますが、多くの場合、医師は腫瘍が支配する前の脳がどのような状態であったかを知る必要があります。しかし、彼らには「以前の状態」の写真がありません。最初のスキャンは通常、患者がすでに病状を呈している時に撮影されます。つまり、「健康な」ベースラインは失われており、損傷した背景に対して疾患を測定しようとすることは、シャツがきれいだった時にどのような状態だったかを知らずに、シャツについた汚れの大きさを測ろうとするようなものです。

これを解決するために、研究者たちは「BraTS Local Synthesis」タスクと呼ばれるチャレンジを作成しました。目標はシンプルです。コンピュータにマスク（隠された領域）のある脳スキャンを与え、そこに現実的で健康な組織のバージョンを描き込ませることです。注意点は、コンピュータの「描き込み」が、SSIM（構造的類似性）、PSNR（ピーク信号対雑音比）、MSE（平均二乗誤差）といった指標を用いて、いかに数学的に実物に近いかで採点されることです。これらは単に「見た目が良いか？」を問うスコアではありません。コンピュータの推測が、実際の健康な組織とどれほど異なっているかを測定する厳格な数学テストなのです。

解決策：厳格なルールブックを持つスマートな画家

著者たちは、「U-DiT」と呼ばれるものに基づいた新しいタイプのAIネットワークを構築しました。U-DiTは、2つの方法を同時に使い分ける画家のようです。まず、彼らは細い筆（畳み込み/コンボリューション）を使用して、脳の隆起のような微細で詳細なテクスチャを描きます。次に、キャンバス全体を俯瞰して（グローバル・アテンション）、脳全体の形状のような大きな全体像を理解します。ほとんどのAIモデルはこれらを効率的に両立させることに苦労しますが、このU-DiTモデルは、より小さなダウンサンプリングされたグリッド上で全体像を見ることで、計算能力を大幅に節約しながら、脳の異なる部分間の長距離のつながりを捉えることができます。

しかし、この論文の真の魔法は、そのアーキテクチャだけではありません。著者が画家をより賢くするために追加した2つの特定のルールにあります。

ルール1：「良いものだけを見る」（Healthy-Only Attention）
通常、AIが穴を埋めようとする際、その穴に何が入るかを推測するために、画像の他の穴を見ることがあります。これは悪いアイデアです。なぜなら、それらの他の穴もまた空であり、未知であるからです。著者たちは、もしAIが他の空の場所を見れば、それは単に堂々巡りの推測をしているだけになると気づきました。そこで、彼らは「Healthy-Only Attention（健康な部分のみのアテンション）」ルールをプログラムしました。彼らはAIにこう命じました。「欠落した箇所を埋めようとする際、他の欠落した箇所を見ることは厳格に禁止されている。君が見ることができるのは、健康で既知の組織のみである」。

さらに、これをよりスマートにするために、特別なバイアスを追加しました。これは、穴の真向かいにある場所に特別な注意を払うようAIに指示するものです。脳は対称的であるため、左側の健康な組織は通常、右側の健康な組織の完璧な鏡像となります。もしAIが左側に穴を埋める必要があるなら、まず右側を見ます。これにより、AIは何かをでっち上げるのではなく、実際に観察された解剖学的構造に基づいて欠落部分を再構築することを強制されます。

ルール2：鏡のトリック（Contralateral-Symmetry Input）
2つ目のトリックは、AIに「カンニングペーパー」を与えることです。著者たちは患者の脳スキャンを左右反転させ（鏡を見ているように）、その反転した画像を元の画像と共にAIに入力します。これにより、AIには欠落している組織がどのようなものであるべきかという直接的なテンプレートが与えられます。なぜなら、脳の健康な側は通常、病んでいる側の鏡像だからです。

しかし、問題があります。鏡に映った部分が脳の外に出てしまったり、別の穴の上に重なってしまったりすることがあります。これを修正するために、彼らは「妥当性チャンネル（validity channel）」を追加しました。これは、AIに対して「この鏡のパーツはリアルで有用だが、あちらの部分は脳の外にあるか、あるいは別の穴であるので無視せよ」と教えるハイライターのようなものです。これにより、AIが安全かつ有用な場合にのみ、鏡像を利用できるようになります。

結果：どの程度うまくいったのか？

著者たちは、219件の脳スキャンのバリデーションセットを用いて新しいシステムをテストしました。結果は目覚ましいものでした。彼らのモデルは、平均構造的類似性（SSIM）スコア 0.864、ピーク信号対雑音比（PSNR）24.7 dB、および平均二乗誤差（MSE）4.6×10⁻³ を達成しました。

これを比較するために、彼らの「U-DiT」モデルを他のアプローチと比較しました。「フル・トランスフォーマー」（より複雑な種類のAI）を使用しようとしたモデルは、SSIMがわずか 0.580 にしか達せず、はるかに劣っていました。著者たちは、畳み込みベースと単一のスマートなグローバル・アテンション・ブロックを組み合わせた彼らの手法が、最適解（スイートスポット）であることを発見しました。

また、特別なルールを取り除いたらどうなるかをテストしました。AIが他の空の場所を見ることを許可した場合（「Healthy-Only」ルールを削除した場合）、スコアは低下しました。鏡の入力を取り除いた場合も、スコアは再び低下しました。これは、「健康な組織のみを見る」という厳格なルールと「鏡のテンプレート」の両方が、高いスコアを得るために不可欠であったことを証明しています。

課題：滑らかさ vs リアリズム

この論文は、興味深い副作用についても指摘しています。AIは数学的な誤差（MSE、PSNR）を最小化しようとするため、非常に滑らかで、わずかにぼやけた画像を作成する傾向があります。これは、多くの可能な健康なテクスチャの「平均」が、滑らかなテクスチャになるためです。これは数学的テストでは非常に高いスコアを獲得しますが、実際の脳組織が持つ微細なノイズを含むディテールのある写真に比べると、少し「リアルではない」ように見えるかもしれません。著者らは、彼らの手法が形状や構造を正しく捉えること（数学的スコアが測定していること）には優れているものの、高周波のテクスチャをいくらか平滑化してしまうことを認めています。彼らは、すでに構築された正確な構造を損なうことなく、それらの微細でリアルなディテールを再び追加するための第2ステップが将来必要になる可能性があると示唆しています。

結論

要約すると、この論文は、規律ある観察者となるようAIを教育することで、欠落した脳組織を補完する巧妙な方法を提示しています。AIに空のスペースを無視させ、代わりに患者自身の健康で鏡写しの側面に頼らせることで、モデルは非常に正確な再構築を実現します。このモデルは単に推測するのではなく、患者自身の既知の健康な部分から欠落したピースを推論するのです。このアプローチは公式リーダーボードで 0.864 の平均SSIMを達成し、時には、パズルを解く最善の方法は、どのピースを見ることが許可されているかを正確に知ることであるということを示しました。