こんにちは、インフォマティクスでAI・機械学習業務を担当している大橋です。
今回は画像領域分割（セマンティックセグメンテーション）技術の紹介をしたいと思います。ここで紹介する論文 [1]は、近年発展著しいTransformerをセグメンテーションに取り入れたものとなっています。
※本記事はTransformer [2]、ViT [6]等の深層学習の知識がある程度あることを前提として書いています。
セマンティックセグメンテーションとは
セマンティックセグメンテーションは画像処理技術に分類され、画像の各ピクセルにクラスを割り当てるアルゴリズムです。
具体的には画像が与えられたときに、どこに人が写っていてどこに車が写っているのか等をピクセル単位で推論します（図1右, [3]より抜粋）。
画像処理の技術として画像分類や物体検知と対比されますが、次のような違いがあります。
- 画像分類：1枚の画像に対してクラスを割り当てる問題
- 物体検知：画像の中から複数個の物体領域をボックスで予測する問題（図1左, [4]より抜粋 ）
つまり、セマンティックセグメンテーションは、より詳細に物体領域を特定する問題になります。
 |
 |
図1：物体検知（左）・セマンティックセグメンテーション（右）
従来のセグメンテーションとその問題点
セグメンテーションは画像分類技術と関連の強い技術です。
画像分類ネットワークの一部をEncoderとして利用しEncoder-Decoder構造を持った領域分割ネットワークを組み立てるといったことがよく行われます。
そのため、これまでは畳込層をベースとしたモデルが主流でした。例えばDeepLab [5]等がそれに当たります。
畳込層の問題点
畳込層は局所的な情報の集約を繰り返すことで徐々にその受容野を広げ画像の特徴量を捉えるものですが、この特性から大域的な情報を捉えるのが苦手とされています。
離れたピクセルの情報を得ようとした場合何層も畳込層を重ねないといけません。1つの畳込層で集約できる領域を広くすることで、ある程度この問題は緩和できます。
例えばDeepLab等で用いられているAtrous Convolutionがこの方面の改善にあたります。しかし、根本的に解決する訳ではありません。
画像分類は段階的に画像のサイズを小さくしていき最終的に画像1枚に対してクラスを予測します。そのため最終的には局所的な情報と大域的な情報は融合されます。
一方でセグメンテーションはピクセル単位での予測になるため、モデルの出力は入力画像と同程度の解像度になります。つまり、畳込層ベースの画像分類ネットワークを利用する場合、どこかで拡大操作をしないといけなくなります。
その場合、局所的な操作を行う畳込層のみでは大域的な情報を捉えづらくなり、良い精度を出すためには複雑なDecoderが必要になります。
セグメンテーションは離れたピクセルどうしが同じ物体なのかを画像の大域的情報から判断しないといけないことを鑑みるに、畳込層を基本とした構造は最適ではないように思えます。
Transformerを用いた先行研究の問題点
このような状況はViT [6]の登場で大きく変わりました。
ViTは自然言語処理で用いられているTransformerと呼ばれる技術を画像分類に応用したものになります。
畳込層ベースのモデルと大きく異なる点は、（もちろん設計によりますが）Transformer層は1層で画像の全領域から情報を直接取り込めることです。
ViTは画像分類のモデルですが、これをセグメンテーションのEncoderに転用したモデルが多数研究・発表されています。代表的なものとしてはSETR [7]が挙げられます。
Transformerをベースとしたモデルは先に述べた受容野の問題を根本から解決するものですが、問題点がいくつかあります。
一つ目はTransformer特有の位置埋め込みに関するものです。通常の位置埋め込みは固定長のため、学習時と予測時で画像サイズが異なる場合うまく対応できません。
二つ目はAttention機構のメモリー消費や計算コストが大きいことです。
SETRの場合はViTをそのままEncoderとして用いているため、各Transformer層の出力が同じ空間解像度を持っています。これは階層構造を持っているモデルに比べて計算量を増大させる原因になっています。
特徴量の空間解像度は元画像の\(\frac{1}{16}\)と低解像度で、ピクセル単位の詳細な予測をするという観点からも良い構造とは言えなそうです。また、SETRはいくつもの畳込層を利用し計算コストの大きいデコードを行っています。
以上のことを踏まえ、Transformerのモデル構造を援用し、既存のTransformerベースのモデルに比べて計算コストを下げつつ精度向上を目指したSegFormerというモデルを提案したのが、紹介する論文の内容になります。
モデル概要
SegFormerの概要は以下になります。計算を効率化したTransformerをベースにEncoder-Decoder構造をしたものになっています。
図2：SegFormer（[1]より抜粋）
SegFormerは精度向上と計算コスト軽減のために、以下のようないくつかの工夫が取り入れられています。
- パッチ埋め込みの工夫
- Attention機構の効率化
- Positional Encodingを動的に計算
- 特徴量の階層構造化
- Decoderの簡略化
それぞれ以下に説明します。
工夫詳細
1. パッチ埋め込みの工夫
通常のViTは入力画像を16×16を１パッチとして特徴量を構成をしていますが、SegFormerでは4×4を１パッチとしています。
これは、ViTは画像分類のモデルで画像を大雑把に捉えれば良いのに対して、SegFormerはセマンティックセグメンテーションのモデルでより詳細に画像を捉えなければいけないためです。画像の入力時に、空間方向の情報を不必要に潰さないことを目的としています。
また、ViT等では埋め込む際にパッチの重なりが無いようにしていますがSegFormerではある程度の重なりを持たせて埋め込んでいます。これは、計算量が増えてしまいますが近隣のパッチの特徴量の連続性を担保し、精度を向上させることを目的としています。
参考論文では2次元の畳込層でパッチ埋め込みを実現しており、\(K,S,P\)をそれぞれカーネルサイズ、ストライド、パディングとすると、入力画像のパッチ埋め込みに対しては\(K=7,S=4,P=3\)、各ブロック間でのパッチ埋め込みでは\(K=3,S=2,P=1\)としています。
2. Attention機構の効率化
Attention機構は大域的な情報を直接捉えますが、その代償として計算コストが大きいです。そのため、Attention機構の効率化は重要な問題です。
参考論文では、Attention機構の構造はそのままに適切に次元を下げることによって計算コストの削減を行っています。Attention機構は、入力長\(N\)、チャンネル数\(C\)とした時、\((N, C)\)の次元を持ったテンソルKey(\(K\))、Value(\(V\))とQuery(\(Q\))に対して
\[Attention(Q,K,V)=Softmax(\frac{QK^{T}}{\sqrt{C}})V\]
と表されます。この計算の複雑度は\(\mathcal{O}(N^2)\)です。この\(K,V\)に対し、テンソル変形(\(Reshape\))と線形変換(\(Linear\))を用いて入力長を圧縮しています。
つまり、圧縮率を\(R\)としたときに以下のような変換を施します。
\[\hat{K}=Reshape(\frac{N}{R},C*R)(K)\]
\[K=Linear(C*R,C)(\hat{K})\]
このようにすることで\(K,V\)は\((\frac{N}{R},C)\)の次元を持ち、Attentionの計算の複雑度が\(\mathcal{O}(\frac{N^2}{R})\)に削減されます。
3. Positional Encodingを動的に計算（Mix-FFN）
Attention機構では位置情報が取り込まれないため、位置埋め込み(Positional Encoding)と呼ばれるテンソルの加算が行われます。
位置埋め込みは通常固定サイズのため、学習時と予測時で画像サイズが違う場合に対応できません。
予測時の画像サイズに合わせて適切にリサイズを行うことによって対応することもできますが、それが精度劣化を招くこともあるようです。
このような問題に対応するため参考論文では、固定サイズの位置埋め込みを捨てて動的に位置埋め込みを計算できるようにしています。
具体的には以下のように多重線形層(\(MLP\))、畳込層(\(Conv\))とGELU活性化関数(\(GELU\))の組み合わせで計算しています。
\[ x_{out} = MLP(GELU(Conv_{3\times 3}(MLP(x_{in}))))+x_{in}\]
ここで、\(MLP(GELU(Conv_{3\times 3}(MLP(x_{in}))))\)部分が動的なPositional Encodingに対応しており、\(Conv_{3\times 3}\)で相対的な位置情報の集約作業を行っています。
4. 特徴量の階層構造化
画像分類モデルであるViTやそれをセグメンテーションに直接応用したSETRでは、Transformer各層の出力する特徴量の空間解像度は変わりません。
参考論文では畳込層ベースのモデル同様、各ブロックごとに解像度を下げていく構造を採用しています。
これは、解像度を下げることで計算コストを下げるとともに、高解像度の特徴量はより詳細な局所的情報に焦点を当て、低解像度の特徴量はより大域的情報に焦点を当てることを目的にしています。具体的にはパッチの工夫で述べた2次元畳込層で行います。
5. Decoderの簡略化
畳込層ベースのモデルの場合Encoderが大域的な情報を捉えきれないため、Decoderも畳込層を用いて周辺情報を利用する方が有利なため必然的に複雑化します。
一方でSegFormerはエンコードした情報自体が既に大域的情報を十分に捉えているはずという主張のもと、簡素なDecoderを採用しています。
具体的には、4つの各Transformerブロック\(i\in \{ 1,2,3,4\} \)から出力される特徴量\(F_{i}\)を、適切に形を整えて結合・線形変換するだけのものです。
各\(F_{i}\)は空間解像度\(\frac{H}{2^{i+1}}\times \frac{W}{2^{i+1}}\)、チャンネル数\(C_{i}\)を持ちます。ただし\(C_{i+1}\geq C_{i}\)。これらを全て以下のように線形変換で共通のチャンネル数\(C\)に統一し、また共通の空間解像度\(\frac{H}{4}\times \frac{W}{4}\)に拡大します。
そしてチャンネル方向に結合して、さらに線形変換をします。この様な簡素な構造は計算量削減を意図しています。\(N_{cls}\)は対象データセットのクラス数です。
\[\hat{F_{i}}=Linear(C_{i},C)(F_{i}) ,\forall i\]
\[\hat{F_{i}}=Upsample(\frac{H}{4}\times \frac{W}{4})(\hat{F_{i}}) ,\forall i\]
\[F=Linear(4C,C)(Concat(\hat{F_{i}})) ,\forall i\]
\[M=Linear(C,N_{cls})(F)\]
精度
学習概要
SegFormerの学習・精度評価はCityscapes [3]、ADE20K [8]、COCO stuff [9]という3つのデータセットで行われています。EncoderはImagenet-1Kで事前学習を行い、Decoderはランダムに初期化をしたうえで学習を行っています。
データ拡張としては倍率\(0.5\sim 2\)のランダムリサイズ、水平方向のランダム反転、ランダムの切り抜きを行っています。
SegFormerはパラメータの小さい順にMiT-B0からMiT-B5まで試しています。モデルの詳細は論文 [1]を参照してください。
精度
ADE20Kに対して代表的なモデルとのパラメータ数・精度を表したのが以下になります。
図3：計算効率比較（[1]より抜粋）
設計の目論見通り、他の代表的なモデルと比べて少パラメータで高精度が出ていることが確認できます。
ADE20KとCityscapesに対する詳細な結果比較は以下です。
図4：精度比較（[1]より抜粋）
速度面では若干見劣りする部分はありますが、いずれのケースにおいても少パラメータで高精度が出ていることが分かります。
差分検証
各工夫が精度にどれだけ効いているかも検証しています。
先ずは動的な位置埋め込み(Mix-FFN)が固定長の位置埋め込み(PE)に対して優位性を持っているかをCityscapesで確認しています。Cityscapesでは\(1024\times 1024\)の画像サイズでモデルの学習を行っています。
図5は学習時のサイズとは異なる2つのサイズ（Inf Res）（\(768\times 768\)と\(1024\times 2048\)）の画像を入力とした際の予測精度を掲載しています。
PEは入力サイズに合わせて伸縮させて用いています。Mix-FFNを用いた方が精度が高く、PEよりも予測の汎化性能が高いことが分かります。
図5：PEとMixFNNの精度差（[1]より抜粋）
次に、Decoderの簡略化の部分を見てみます。Decoderの簡略化はTransformerベースのモデルはEncoderの出力が既に大域的な情報を取り入れているため複雑なDecoderを必要としないという予想のもと行われました。
その予想を確かめるためにDecoderをそのままに畳込層ベースのEncoderを用いた場合の精度を比較しており、その結果が図6になります。畳込層ベースのモデルに比べて精度が高く、予想を支持する結果になっています。
また図6では\(S1-S4\)は4つすべてのブロックの特徴量を用いたことを、\(S4\)は最終ブロックの特徴量のみを用いたことを表します。
つまり、Transformerを利用した場合4つすべてのブロックを使った畳込ベースのモデルより精度が出ていることになり、いかに大域的な情報を効率的に取り入れることが重要かを示した結果になっています。
図6：バックボーンの差（[1]より抜粋）
まとめ
今回はTransformer構造をセグメンテーションネットワークに効率的に適用したモデルの紹介でした。
Attention機構を導入して精度を向上しつつも、特徴量の階層構造化・Attention機構の効率化・Decoderの簡素化によって計算量をうまく抑えたモデルです。
Attention機構効率化のための入力長方向の圧縮はK,Vに対しては大きなパッチで処理するのと変わらないので、空間解像度の細かさはQのみによって担保されているという構造になっています。
計算の複雑度は入力長の2乗であることには変わらず、本質的な改善ではないという印象があります。
自然言語のモデルで様々な効率的なAttention機構が提案されているので、それらが画像系のタスクの精度にどのような影響があるのか気になるところです。
論文 [1]ではより詳細な分析がなされていますので、気になる方はぜひそちらを見てみてください。
参考文献
[1]Xie, E., Wang, W., Yu, Z., Anandkumar, A., Álvarez, J.M., & Luo, P. (2021). SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers. ArXiv, abs/2105.15203.
[2]Vaswani, A., Shazeer, N.M., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., & Polosukhin, I. (2017). Attention is All you Need. ArXiv, abs/1706.03762.
[3]Cordts, M., Omran, M., Ramos, S., Rehfeld, T., Enzweiler, M., Benenson, R., Franke, U., Roth, S., & Schiele, B. (2016). The Cityscapes Dataset for Semantic Urban Scene Understanding. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 3213-3223.
[4]Everingham, M., Eslami, S.M., Gool, L.V., Williams, C.K., Winn, J.M., & Zisserman, A. (2014). The Pascal Visual Object Classes Challenge: A Retrospective. International Journal of Computer Vision, 111, 98-136.
[5]Chen, L., Zhu, Y., Papandreou, G., Schroff, F., & Adam, H. (2018). Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation. ArXiv, abs/1802.02611.
[6]Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ArXiv, abs/2010.11929.
[7]Zheng, S., Lu, J., Zhao, H., Zhu, X., Luo, Z., Wang, Y., Fu, Y., Feng, J., Xiang, T., Torr, P.H., & Zhang, L. (2021). Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers. 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 6877-6886.
[8]Zhou, B., Zhao, H., Puig, X., Fidler, S., Barriuso, A., & Torralba, A. (2018). Semantic Understanding of Scenes Through the ADE20K Dataset. International Journal of Computer Vision, 127, 302-321.
[9]Caesar, H., Uijlings, J.R., & Ferrari, V. (2018). COCO-Stuff: Thing and Stuff Classes in Context. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, 1209-1218.
GISやAI機械学習を使ったシステムのご相談を承っています。お気軽にお問い合わせください。
