SegFormer徹底解説 — 位置エンコーディングなしのTransformerはなぜセグメンテーションで機能するのか

セマンティックセグメンテーションにTransformerを導入する試みは、SegFormer登場以前から存在したが、多くは位置エンコーディングと複雑なデコーダに依存しており、テスト時の解像度が学習時と異なると精度が落ちる、計算コストが高いという課題を抱えていた。SegFormer(Xie, Wang, Yu, Anandkumar, Alvarez, Luo、NeurIPS 2021)は、この2つの問題を「位置エンコーディングを使わない階層型エンコーダ」と「畳み込みを一切使わないAll-MLPデコーダ」という、驚くほどシンプルな設計で同時に解決したモデルである。この記事では、原論文(arXiv:2105.15203)の記述に基づいて、その設計がなぜ機能するのかを一つずつ検証する。

この記事は原論文「SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers」(Xie et al., NeurIPS 2021, arXiv:2105.15203)の記述に基づいて構成している。

0. この記事で分かること

• SegFormerが解決しようとした、従来のTransformerベースセグメンテーションの課題
• 位置エンコーディングなしで動く階層型エンコーダMiT(Mix Transformer)の仕組み
• なぜ「畳み込みを一切使わないMLPだけ」のデコーダで十分な精度が出るのか
• Cityscapes/ADE20KでのB0〜B5の実測mIoU・パラメータ数・計算量
• ロバスト性(画像の乱れへの強さ)を含めた実用上の選び方

1. まず結論：SegFormerとは何か

SegFormerは、位置エンコーディングを持たない階層型Transformerエンコーダ(MiT)と、畳み込みを一切使わない軽量なAll-MLPデコーダを組み合わせたセマンティックセグメンテーションモデルである。エンコーダが複数解像度の特徴を出力する設計にしたことで、デコーダ側の複雑な機構を大幅に削減できる、という「エンコーダとデコーダの役割分担を設計し直す」ことが核心のアイデアである。

2. なぜこの設計が必要だったのか

Semantic Segmentationにおいて、ViT(Vision Transformer)をエンコーダとしてそのまま使うと、2つの問題が生じる。1つは、ViTが出力する特徴マップが単一の解像度しか持たない点である。物体のサイズは画像内でさまざまであり、従来のCNNベースのセグメンテーションモデル(FCN、U-Netなど)は、複数解像度の特徴を組み合わせることでこの問題に対処してきた。単一解像度しか出さないViTをそのまま使うと、この多スケール表現力を失ってしまう。

もう1つは、位置エンコーディングへの依存である。ViTは画像をパッチに分割してTransformerに入力するが、Transformer自体には「どのパッチが画像のどこにあるか」という位置情報がないため、固定長または学習可能な位置エンコーディングを別途足す必要がある。この位置エンコーディングは学習時の入力解像度に対して作られているため、推論時に学習時と異なる解像度の画像を入力すると、位置エンコーディングを補間しなければならず、精度が落ちるという実用上の弱点を持つ。セグメンテーションは検出以上に高解像度な入力を要求されることが多く、この解像度依存はより深刻な問題になる。

SegFormerは、この2つの問題を「階層型のエンコーダ設計」と「位置エンコーディングそのものを廃止する」ことで、根本から解消しようとした設計である。

3. 入力は何か

入力は他のセグメンテーションモデルと同様、RGB画像 x \in \mathbb{R}^{H \times W \times 3} である。SegFormerはCityscapesでは高解像度(1024×2048)、ADE20Kでは512×512程度の画像で評価されている。ViTが16×16という粗いパッチでダウンサンプルするのに対し、SegFormerのエンコーダは最初のパッチサイズを4×4まで細かくしており、より細部を保った特徴抽出から出発する。

4. 何を求めるのか／出力は何か

出力は、入力画像の1/4解像度で、各画素がどのクラスに属するかを表すセグメンテーションマップ \hat{y} \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times N_{cls}} である(N_{cls}はクラス数)。これを最近傍補間などで元の解像度に戻して最終出力とする。核心のアイデアは、「多スケールな特徴をエンコーダ側で十分に作り込んでおけば、デコーダは単純な線形層の組み合わせだけで、それらを統合し高精度なマスクへ変換できる」という設計思想である。

5. 基本構成

SegFormerは大きく2つのブロックからなる。画像から4段階の解像度の特徴を抽出するMiT(Mix Transformer)エンコーダと、それらを統合してセグメンテーションマスクを出力するAll-MLPデコーダである。

図1 — MiTエンコーダが4段階の解像度(1/4, 1/8, 1/16, 1/32)で特徴を出力し、All-MLPデコーダが線形層だけでそれらを統合してセグメンテーションマスクを作る。畳み込みを持つのはエンコーダ内のMix-FFNだけであり、デコーダ全体は文字通り線形層のみで構成される。

エンコーダが4つの解像度の特徴をあらかじめ作り込んでおくことで、デコーダ側は「それぞれをそろえて足し合わせるだけ」というほぼ最小限の処理で高精度なマスクを得られる。この分業設計が、SegFormer全体のパラメータ数を抑える最大の要因になっている。

6. 構成要素の技術詳細

Overlapped Patch Merging — 4段階の階層構造

MiTエンコーダは、ViTのように画像を一度だけパッチに分割するのではなく、4つのステージそれぞれで解像度を段階的に下げていく。各ステージの出力解像度は入力に対して1/4、1/8、1/16、1/32であり、CNNベースのバックボーン(ResNetなど)と同じ階層構造を再現している。

パッチ分割の方法も、ViTのように隣接しないパッチに切り分けるのではなく、重なりを持たせて(Overlap)パッチをマージする。最初のステージではカーネルサイズ K=7、ストライド S=4、パディング P=3、以降のステージでは K=3, S=2, P=1 という設定が使われる。パッチ同士に重なりを持たせることで、パッチの境界をまたぐ局所的な連続性(隣接パッチ間の情報)を保ったまま特徴マップを縮小できる。

効率的自己注意(Efficient Self-Attention)

通常のマルチヘッド自己注意は、系列長 N に対して O(N^2) の計算量を要する。高解像度画像を扱うセグメンテーションでは N が非常に大きくなるため、これは深刻なボトルネックになる。SegFormerは、Key・Valueの系列を縮約率 R であらかじめ圧縮してから注意を計算する、Sequence Reduction機構を導入している。

この K, V を、系列長 N から N/R まで縮約してから計算に使う。縮約率 R はステージ1から4にかけて [64, 16, 4, 1] と設定されており、解像度が高く系列長が長い浅いステージほど強く縮約し、解像度が下がって系列長が短くなる深いステージでは縮約しない、という設計になっている。これにより計算量は O(N^2) から O(N^2/R) まで下がり、高解像度入力でも実用的な速度で自己注意を計算できる。

Mix-FFN — 位置エンコーディングを廃止する仕組み

SegFormerが位置エンコーディングを一切使わずに済んでいる理由は、順伝播ネットワーク(FFN)の内部に3×3の畳み込みを混ぜ込んだMix-FFNにある。

3×3畳み込みは、画像の端でゼロパディングされる。このゼロパディングされた境界が、間接的に「自分が画像のどこに位置しているか」という手がかりを畳み込みに漏らす(leak)ため、明示的な位置エンコーディングを与えなくても、Mix-FFNを通すだけで位置情報が特徴に織り込まれる。この設計により、学習時と異なる解像度でテストしても、位置エンコーディングを補間するという精度を落とす操作が原理的に不要になる。

All-MLPデコーダ — なぜ畳み込みなしで十分なのか

デコーダは次の4段階からなる、文字通り線形層(MLP)だけの構成である。

1. チャネル統一: 各ステージの特徴 F_i を、それぞれ独立した線形層で同じチャネル数 C に写像する。
2. アップサンプルと結合: すべての特徴を1/4解像度に揃えてアップサンプルし、チャネル方向に結合(concat)する。
3. 融合: 結合された 4C 次元の特徴を、線形層で C 次元に融合する。
4. 予測: 最後の線形層で、クラス数分のチャネルを持つセグメンテーションマスクを出力する。

この設計が機能する根拠として、原論文は有効受容野(Effective Receptive Field, ERF)の分析を挙げている。CNNベースのデコーダ(DeepLabv3+など)が高精度を出すには、広い受容野を得るために複雑な拡張畳み込みやマルチスケールプーリングの機構が必要だった。これに対しSegFormerのMiTエンコーダは、Transformerの自己注意によって、浅いステージでは局所的、深いステージでは非局所的(画像全体に及ぶ)なERFを、追加の機構なしに自然に獲得している。エンコーダ自体がすでに「局所から大域まで」の情報を持つ特徴を作り込んでいるため、デコーダは複雑な受容野拡大の仕組みを持つ必要がなく、単純な線形層の組み合わせだけで十分な精度に到達できる、というのがこの設計の理論的な裏付けになっている。

学習レシピ

原論文の学習設定は、特殊なトリックに頼らないシンプルなものである。ImageNet-1KでMiTエンコーダを事前学習した上で、デコーダはランダム初期化から学習する。最適化にはAdamWを使い、学習率は初期値0.00006からpolyスケジュール(べき乗則で減衰させるスケジュール、係数1.0)で減衰させる。反復回数はADE20K/Cityscapesで160K、COCO-Stuffで80K、学習時のデータ拡張は比率0.5〜2.0のランダムリサイズ、ランダム水平反転、ランダムクロップ(ADE20Kは512×512、Cityscapesは1024×1024、B5のみADE20Kで640×640)にとどめている。原論文自身が「OHEM(Online Hard Example Mining)、補助損失、クラスバランス損失といった広く使われるトリックを採用しなかった」と明記しており、これらの工夫なしでもAll-MLPデコーダとMiTエンコーダの組み合わせだけで高いmIoUに到達できることを示す狙いがある。

7. モデルバリアント比較

SegFormerは、エンコーダのMiT-B0からMiT-B5まで、6段階のサイズで提供される。以下は原論文(Table 2)が報告するCityscapes/ADE20Kでの実測値である。

出典: 原論文Table 2(Xie et al., NeurIPS 2021)。Cityscapes mIoUはマルチスケール+左右反転(MS)テスト時の値、パラメータ数はエンコーダ・デコーダ込みの全体値(デコーダはB5でも全体の約4%程度)。参考として、単一スケール(SS)テストでのB5のmIoUは82.4であり、MSテストによって84.0まで伸びる。

原論文はこの結果を、当時のSOTA(最良)モデルと比較して「SegFormer-B4は64.1MパラメータでADE20K 50.3% mIoUを達成し、当時の最良手法より5分の1近く小さいモデルサイズで2.2ポイント上回った」と報告している。B0はわずか3.8Mパラメータながら、Cityscapesで76.2 mIoUという実用的な精度を維持しており、エッジ実装を意識した設計であることが数値からも裏付けられる。

8. 何が苦手か／難しい環境

SegFormerの弱点は、多くがTransformerベースモデル一般に共通するものである。まず、大きなモデル(B3以上)ほど学習・推論の計算コストが高く、CityscapesのFLOPsはB0の125.5GからB5の1460.4Gまで、10倍以上に膨らむ。高解像度画像を扱うセグメンテーションでは、この計算量の増大がリアルタイム運用の障害になりやすい。

大規模事前学習データへの依存も無視できない。MiTエンコーダはImageNet-1Kで事前学習されており、事前学習データと大きく異なるドメイン(医療画像、衛星画像など)へ転移する際には、ファインチューニングデータの質と量が精度を大きく左右する。

極小物体や細い構造(電線、標識のポールなど)の検出は、SegFormerでも根本的に容易ではない。Overlapped Patch Mergingは最初のステージでも4×4のダウンサンプルを行うため、数ピクセル幅しかない構造は、深いステージに到達する前に情報として失われやすい。

一方で、原論文が示すCityscapes-Cでのロバスト性評価は、SegFormerの明確な強みでもある。ガウスノイズを加えた劣化画像に対し、SegFormer-B5はDeepLabv3+(Xception-71バックボーン)と比べてmIoUの相対的な改善が最大588%、雪(snow)を模した劣化に対しては最大295%の相対改善を記録しており、通常のCNNベースモデルより画像の乱れに強いことが定量的に示されている。この頑健性は、Transformerの自己注意が局所的なノイズに引きずられにくいという性質に起因すると考えられる。

9. 実用上の選び方

エッジデバイス・リアルタイム処理では、SegFormer-B0/B1が現実的な選択肢になる。B0は3.8Mパラメータ、Cityscapesで76.2 mIoUという、精度と軽量性のバランスが取れた設定であり、組み込み機器やドローンへの搭載も視野に入る。

自動運転の走行可能領域認識のように、高解像度画像を扱いつつ一定の精度が求められる用途では、B2〜B3が妥協点になりやすい。B4/B5は精度こそ最も高いが、Cityscapes FLOPsが1200Gを超えるため、車載GPUでのリアルタイム推論には最適化(量子化、TensorRT変換など)が前提になる。

医療画像・衛星画像などのオフライン精密解析では、リアルタイム性より精度が優先されるため、B4/B5を選び、必要に応じてドメイン固有データでファインチューニングするのが妥当である。

悪天候や照明変化が多い屋外環境(農業ロボット、屋外監視カメラなど)では、SegFormerのCityscapes-Cでの頑健性の高さが実用上の利点になる。同程度の精度を持つCNNベースモデルより、画像の乱れに対して精度低下が小さいことが期待できる。

より基礎からObject DetectionとSemantic Segmentationを学びたい場合は「Object Detection・Semantic Segmentation入門」を、セグメンテーション分野全体の最新動向は「セマンティックセグメンテーションの技術動向」を参照してほしい。

10. まとめ（3行で復習）

• SegFormerは、位置エンコーディングなしで多スケール特徴を出す階層型エンコーダ(MiT)と、畳み込みを使わないAll-MLPデコーダを組み合わせたモデルである。
• Mix-FFNの3×3畳み込みが位置情報を暗黙的に漏らすことで位置エンコーディングを不要にし、Transformerの広い有効受容野がシンプルなMLPデコーダを可能にしている。
• B0(3.8M、Cityscapes 76.2 mIoU)からB5(84.7M、84.0 mIoU)まで幅広いサイズがあり、画像の乱れに対する頑健性の高さも実測で裏付けられている。

参考リンク

• SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers (NeurIPS 2021, arXiv:2105.15203)
• SegFormer 論文 ar5iv版(全文)
• SegFormer 公式実装(NVlabs/SegFormer, GitHub)
• NeurIPS 2021 Proceedings: SegFormer