Attention U-Net 徹底解説:セグメンテーション精度を飛躍させるAttention Gateの実力

こんにちは、ゆずかきです。
ここからは、「Attention U-Net(AttentionUnet)」 と呼ばれる手法について、がっつりとまとめてみたいと思います。

今回取り上げる「AttentionUnet」は、論文

Attention U-Net: Learning Where to Look for the Pancreas
arXiv:1804.03999v3 cs.CV
Ozan Oktay ら、Imperial College London を中心とした研究グループによる提案手法です。

U-Net を使った医用画像セグメンテーションは既に定番となっていますが、この「Attention U-Net」は、“Attention Gate (AG)” という仕組みを組み込み、複数の臓器・組織の中でも「どこに着目すべきか」を学習的に強調するというコンセプトを導入しています。特に、膵臓(Pancreas)のように形状や大きさのばらつきが大きい臓器をセグメンテーションする際に強みを発揮すると報告されています。

それでは、論文の内容と実験結果を一通り振り返りながら、「Attention Gates(AGs)」 のアーキテクチャや効果、導入メリットなどをできるだけ網羅的に解説していきましょう。


§本記事の構成

  • はじめに: Attention U-Net とは?
  • 従来のU-Netとの比較・AG導入の背景
  • Attention U-Netのアーキテクチャ
  • Attention Gate(AG)の仕組み:Additive Attentionによるゲーティング
  • 実装上のポイント:3D U-Net・学習設定など
  • 主要な実験結果と考察(Pancreas Segmentationにおける評価)
  • マルチクラス・マルチラベルへの拡張と利点
  • 課題と今後の展望
  • まとめ

医用画像において“注目すべき領域を自動的にハイライトする”という流れは、昨今の自然言語処理や画像認識の世界でもホットな話題です。論文では主に膵臓CTを事例に検証されていますが、Attentionによる領域選択は幅広い応用が見込まれます。順を追って見ていきましょう。


§はじめに: Attention U-Net とは?

まず、Attention U-Net の大まかな位置づけを整理します。U-Net は、医療画像セグメンテーション(特に2次元 or 3次元のCT/MRIなど)の分野で広く使用されるFCN(Fully Convolutional Network)の一種で、エンコードとデコードを対称的な構造にし、スキップ接続によって細かい局所情報を下位層から高解像度のまま受け渡す仕組みが特徴的です。

しかし、臓器や病変の形状・サイズが大きく異なる場合、あるいは背景組織とのコントラストが低い場合などには、「真に注目すべき領域以外の特徴マップが混ざってしまう」 という課題があります。この対策として、過去には「マルチステージのカスケード」などがよく用いられました。

  • 例:まずは粗いローカライゼーションを行い、その領域を切り出して第2段で詳細セグメンテーションを実施する…など

このようなカスケード型手法は精度が高い一方で、計算資源やモデルパラメータを多く消費します。そこで提案されたのが Attention Gate (AG) を活用し、U-Net 内で「注目領域」のモジュールを加えることで、単一モデルでありながら不必要な背景情報をフィルタリングしよう というアプローチです。

Attention Gate は自然言語処理の分野(たとえば機械翻訳におけるBahdanau Attention)や、画像認識分野のsoft-attentionなどを拡張し、各ピクセル/ボクセルごとに注目係数を算出→不要な位置の活性を抑制するという仕組みを実装します。

結論を先に述べると、この Attention U-Net は、膵臓を含むマルチクラスの腹部CTセグメンテーションで高い性能向上を示し、しかもハードウェア的な追加コストがそれほど大きくないことが報告されています。


§従来のU-Netとの比較・AG導入の背景

U-Net の代表的な流れをおさらいすると、

  1. エンコーダ(下り経路)で画像を畳み込み + プーリングしながら特徴マップを抽出
  2. デコーダ(上り経路)でアップサンプリングを行いつつ、スキップ接続された浅い層の特徴と結合して細かい空間情報を補完
  3. 最終的に、元の解像度と同じサイズのセグメンテーションマップを出力

多くの医用画像タスクで成功を収めてきましたが、標的臓器のサイズが小さい場合や、背景とのコントラストが乏しい場合、あるいは複数の臓器が類似の強度で写る場合などでは、「誤検出」や「感度(Recall)の低下」が起こりやすくなります。

従来の工夫:カスケード型のROI局所化

論文内でも言及されているように、これまでの対策としては「2段階(あるいは多段階)のカスケード CNN」を組み合わせる、というやり方がとられていました。

  • Stage1: 低解像度 or よりシンプルなCNNで対象臓器の候補領域をざっくり推定
  • Stage2: 推定された領域だけを切り取り、高解像度のモデルで詳細なセグメンテーション

確かに精度は上がりますが、モデルを2つ以上用意するうえ、それぞれで再び畳み込みを行うため、冗長な計算が増えることがデメリットです。

Attention Gateの狙い

Attention Gate(AG) は、この「不要な背景特徴を抑制する」ことを、単一のネットワーク(U-Net)内で実装してしまおう、という試みです。

  • カスケードの Stage1 のように明示的に領域切り出しを行うのではなく、AG がソフトに背景をフェードアウトさせる。
  • レイヤー途中のスキップ接続にゲートを挿入し、本当に必要な領域に対しては特徴を通し、不要そうな領域の活性は抑える
  • しかも、AG は「タスクに応じて学習される」ので、具体的にどこに注目するかを人手で指示する必要がない。

一般に “Attention” といえば NLP で盛んに研究されていますが、自然画像(ImageNet など)や医用画像での物体領域学習においても、「soft-attention」「self-attention」「channel-wise attention」など、多彩なバリエーションが提案されてきています。今回の論文が特徴的なのは、U-Net のようなエンコーダ・デコーダ構造の skip connection の部分に、局所的(grid 単位)の gating を行うモジュールを埋め込んだ 点と言えます。


§Attention U-Netのアーキテクチャ

論文の図(Figure 1)では、Encoder 側で徐々に空間サイズを縮小しながら特徴マップを取得し、Decoder 側でアップサンプリングしながら段階的に解像度を復元していく構造は、基本的に U-Net と同じです。ただし、各スキップ接続の先に Attention Gate (AG) を用意している、という違いがあります。

Encoder(↓)
    ↓
(AG) -- Skip Connection 
    ↓
Decoder(↑)

このように、AG が挟まることで「本当に必要なピクセル・ボクセルだけを強調する」 という効果を期待します。 なお、論文では 3D CT への適用を想定しているため、3D 畳み込み3D Max-pooling3D Upsampling を組み合わせており、出力は 3D のボリュームセグメンテーションマップとなります。


§Attention Gate(AG)の仕組み:Additive Attentionによるゲーティング

さて、このモデルの核心部分である Attention Gate (AG)。簡単に言うと、入力特徴マップ xl に対し、ゲート信号 g をもとに「注目係数(attention coefficient)α」を計算し、その α を用いて xl を要・不要な部分にスケーリングする仕組みです。

Additive Attention の数式

論文内で提示されている数式は、以下のような「Additive Attention」を用いています(Bahdanau Attention [2] の発想を継承)。

ここで、σ はシグモイド関数、ϕ は ReLU あるいは tanh などの非線形関数、Wx、Wgは学習パラメータを表します。要は、Skip Connection から来た特徴ベクトル Xil と、上位層から来るコンテキスト(Gating Signal) gi を合わせて線形変換し、その出力から [0,1] のスカラー αi を生成 する、という流れ。

αi が大きければ「そこは重要!」、小さければ「そこはあまり見なくていい」という具合に、元の特徴に乗算して抑制/強調を行います。

  • gi は、より粗い空間スケールでの特徴マップ(Encoder の深層部、または Decoder 側の対応レイヤ)から抽出され、“どんな物体を探しているか” の文脈情報を示す。
  • これによって、より高レベルなセマンティクスと低レベルな特徴量を比較し、必要な部位をピンポイントに抜き出す ことが可能になるわけです。

また、grid 単位で計算する という点が重要。自然言語などでは1次元的(単語単位)のattention、画像分類では画像全体を1つのグローバルプーリングベクトルにまとめてチャネル単位でゲーティングするやり方などがありますが、本手法では3次元ボクセル単位での Attention を行うので、極めてきめ細かい局所制御が可能になるわけですね。


AG のForward/Backwardにおける効果

論文でも強調されているとおり、AttentionによるゲーティングはForwardパスだけではなくBackwardパスにも影響します。具体的には、

のように、α が低い画素については、勾配も抑制される ため、不要な背景部分に起因する誤差逆伝搬はあまり重み更新に影響を及ぼさない 形となります。これによって、モデルが主に対象領域の誤差をもとに学習を進めることになり、「より早く、かつ高精度な収束につながる」という理屈です。


§実装上のポイント:3D U-Net・学習設定など

論文では、実験として以下の点が注目されます。

  1. 3Dアーキテクチャ
  • 医用CTは3次元なので、U-Net も 3D Conv, 3D Pool, 3D Upsampling を採用
  • バッチサイズが小さくなることを考慮し、Batch Normalization + Adam オプティマイザなどを用いて安定化
  1. Deep Supervision
  • U-Net のデコーダ各段階で、中間的に「セグメンテーションマップ」の出力を作り、最終出力と合わせて誤差を計算する技法
  • Holistically-Nested系の深層監督 (深層ステージごとに損失を計算) を使うことで、各スケールで取得した特徴がしっかり意味をもつようになる
  1. 学習データと評価データ
  • 主に CT-82 (NIH-TCIA) dataset(82例の膵臓CT)と、CT-150 dataset(150例の腹部CT, 膵臓・肝臓・脾臓などがアノテーションされている)を使用
  • これらのデータセットで従来の U-Net, multi-stage CNN などと比較し、Dice係数やSurface-to-Surface Distance 等で検証している
  1. ソフトウェアおよびリソース
  • PyTorch 実装で公開されており(GitHubリンク 参照)、誰でも再現性をチェック可能
  • 3Dモデルのため GPU メモリ消費量は大きいが、AGのパラメータ数そのものは小さい ため、全体の計算負荷は最小限に抑えられている

以上のようなセッティングで、同じU-Netでも Attention Gate を追加した場合にどれほど精度が改善するかが検証されています。


§主要な実験結果と考察(Pancreas Segmentationにおける評価)

論文の中心的な実験は、腹部CTにおける膵臓セグメンテーション。膵臓は形状バラつきが激しく、背景とのコントラストも低いため、セグメンテーションが難しい部位として知られています。

結果のハイライト

  • Attention U-Net は、標準的な U-Net と比べて Diceスコアが 2~3% 向上(トレーニングデータ数が少ない場合でも改善)
  • カスケード型(複数モデル)で膵臓だけをローカライズしてから詳細にセグメンテーションする、という従来手法に比べても遜色ない or 上回る性能
  • 特に Recall(感度)が大きく改善 → 「膵臓の端っこを見逃しにくくなる」
  • Surface-to-Surface Distance などの表面距離指標でも、Attention 導入による誤差減少が確認された

また、Spleen(脾臓)や Kidney(腎臓)といった他の臓器についても同時にマルチクラスセグメンテーションを実行した際、Attention が入ることで特徴選択がうまく機能するため、モデルのロバストネスが増すという結果が示されています。

AGの学習可視化

論文の図(Figure 4)には、学習の初期~後期にわたるAttentionマップの変化が示されています。最初は画面全体に均一気味だった注意重みが、学習が進むにつれ、膵臓周辺や隣接する腎臓・脾臓といった領域に集中していく様子が視覚的に示されています。これはAGが“必要な領域だけ通す”フィルタとして機能している証拠と言えます。


§マルチクラス・マルチラベルへの拡張と利点

論文中で興味深いのは、AG で複数の臓器を同時にセグメンテーションさせる場合にも対応できる、という点です。

  • AG は単一スカラーを学習するだけでなく、マルチチャネルの attention map を学習可能(論文内で “multi-dimensional gating” と言及)。
  • たとえば膵臓に着目するチャンネル、脾臓に着目するチャンネル… というように分割して学習できる。
  • 結果として、膵臓以外の臓器を同時にアノテーション学習させることで、タスク間の相補情報が得られ、単一臓器のみの学習よりも精度が上がる 可能性が高い。

さらに注目すべきは、一度AGが搭載されたモデル構造であれば、同じネットワークで他の臓器やタスクにも使い回せる点です。マルチステージのカスケードモデルを臓器ごとに構築する手間を省きながら、高精度なセグメンテーションを実現できる点は大きな魅力と言えるでしょう。


§課題と今後の展望

論文終盤で指摘されている課題や今後の方向性についても整理します。

  1. Transfer Learning の活用
  • Attention Gate の学習には十分なデータが必要。
  • とはいえ、既存のU-Net重みを初期化として使い、fine-tuning する形で AG パラメータを最適化すれば、さらに効率的に学習できる可能性がある。
  1. ResNet / DenseNet など他アーキテクチャとの併用
  • 本論文では基本的に「U-Net の skip connection + AG」という形だが、近年はResNet, DenseNet, Swin-Transformer系など、さまざまなモジュールが提案されている。
  • それらと AG を組み合わせれば、さらなる性能向上が期待できるだろう。
  1. 大規模3Dモデルへの拡張
  • 本論文の実験は、GPU メモリの都合上、ボリュームをある程度ダウンサンプリングしている。
  • 近年は GPU メモリが増え、また分散学習も容易になっているため、フル解像度の3D CT複数のコンテキストを同時に扱う など、さらなる拡張が見込まれる。
  1. 医用画像以外への応用
  • 自然画像セグメンテーション、衛星画像解析、あるいは3D点群セマンティックセグメンテーションなどでも、Attention Gate の設計思想は通用する可能性がある。
  • 今後は、医療分野に留まらず、“本質的にどこを注目すべきか” という観点が重要なさまざまな領域で研究が進むと考えられる。

§まとめ

Attention U-Net は、従来の U-Net に “Attention Gate (AG)” を導入することで、不要な背景をソフトに抑制し、対象臓器・構造の局所領域を強調する という仕組みを備えたアーキテクチャです。

  1. カスケードモデル不要:従来の “Localization → Segmentation” の二段階手法を単一モデル内で実現し、計算負荷や冗長パラメータを削減
  2. 局所的 Attention:Additive Attention を3D空間の各ピクセル(ボクセル)単位で計算するため、形状のばらつきが大きい臓器でも高い表現力
  3. マルチクラスへの対応:AG を多次元化することで、複数の臓器や部位を同時にセグメンテーション
  4. 軽量な追加パラメータ:AG 導入によるパラメータ増はごくわずかで、実行速度やメモリへの影響も軽微

論文実験では、膵臓CTセグメンテーション(TCIA/NIH Dataset や CT-150 Dataset)で顕著に性能が向上し、Dice スコアが数%上昇、特にRecallの改善に寄与するという結果。深層学習ベースの医用画像セグメンテーションの中でも、今後も十分に注目され続けるであろう“古典的だけど有用”な拡張と言えるでしょう。


§おわりに

以上、Attention U-Net: Learning Where to Look for the Pancreas 論文を中心に、Attention Gate を組み込んだ医用画像セグメンテーション手法を振り返ってみました。

個人的に面白いと感じるのは、「大域的なコンテキスト(ゲート信号)で、小領域ごとの特徴活性を動的に抑制/強調する」という発想。自然言語処理における自己注意(Self-Attention)とはまた一味違い、スキップ接続上で“必要な場所だけ”を可変的に通すことで、カスケード不要の ROI抽出を実現している点は非常に革新的です。

医療画像解析に限らず、小さな物体や形状の不定な領域をターゲットにするセグメンテーションや検出タスクでも、Attention Gate のアイデアはさらに広がっていくのではないでしょうか。実際の実装は GitHub に公開されていますので、興味のある方は是非覗いてみて下さい。

最後までお読みいただき、ありがとうございました!


🔑この記事のポイントまとめ

  • Attention U-Net:従来のU-NetにAttention Gate(AG) を組み込み、不要な背景をソフトに抑制して対象構造を強調する
  • カスケード不要:従来の多段モデルのような外部ローカライゼーションが不要で、単一CNNで高精度セグメンテーション
  • 局所単位のAdditive Attention:3Dボクセル単位の注意重みを計算し、Forward/Backward両面で重要度制御
  • 膵臓セグメンテーションで有効性を実証:DiceスコアやRecallが顕著に改善、Surface Distanceも縮小
  • 今後の方向性:Transfer Learning やさらに大規模な3Dモデルへの適用、他のネットワークとの組合せが期待

§参考文献(論文リンク)


画像解析

Posted by yuzukaki-dialog