外観検査向け異常検知手法に関する論文紹介: EfficientAD

LeapMindの花元です。

異常検知、してますか？

学習済み深層学習モデルを使って外観検査向け異常検知をしたいけど、良い感じの速度で推論できなくて困ること、長い人生で、たまにはありますよね。今日はそんなときに使える外観検査向け異常検知の手法であるEfficientADのアルゴリズムを紹介したいと思います。

背景

以前の記事で紹介したとおり、近年、外観検査向け異常検知の手法として、画像の特徴を得るために学習済み深層学習モデルを用いる手法が多く提案されています。これらの手法でより高い異常検知性能を得るためには、WideResNet-101のような非常に大きなサイズの学習済み深層学習モデルを用いる必要があることから、GPUが使えるような環境であったとしても数ミリ秒といった高速で推論することは困難です。また、これらの手法は対象物にある傷やシミといった局所的な異常を検知することは比較的得意ですが、部品は正常だが誤った順番に並んでいる、といった論理的な異常を検知することは比較的苦手です。

EfficientADでは、推論の高速化のために、モデルの蒸留によるモデルサイズの削減を行います。また、論理的な異常を検知できるようにするために、オートエンコーダーを導入します。

アルゴリズム

EfficientADは

• トレーニング時
  • モデルの蒸留を用いた、局所的な特徴抽出モデル（Patch Description Network / PDN）構築
  • 特殊な損失関数を用いた、局所的な異常検知のためのモデル構築
  • オートエンコーダーを用いた、論理的な異常検知のためのモデル構築
• 推論時
  • 異常マップの正規化

の4点から構成されます。これらを順に説明していきます。

モデルの蒸留を用いた、局所的な特徴抽出モデル（Patch Description Network / PDN）構築

EfficientADでは、まず、局所的な特徴を抽出するためのモデル（PDN）を構築します。PDNは、4つ（もしくは6つ）の畳み込み層からなる小さなネットワークです。ネットワークの構造の例を図1に示します。

図1: PDNの構造の例（論文より引用）

畳み込み層の数が少ないことから、このネットワークの出力として得られるベクトルの受容野（ベクトルの値に影響を与える入力の範囲）は33x33となり、ResNet等の通常の深層学習に用いられるモデルの受容野よりかなり小さくなります。受容野が小さいと、入力画像全体を捉えた特徴を抽出することができなくなるため、ImageNetの画像分類等のタスクでは不利になりますが、物体のキズ等の局所的な異常を検知する際には小さな受容野があれば十分で、大きな受容野はかえって悪影響を与える（キズ等の位置と関係が薄い位置の特徴ベクトルの値に大きな影響を与える）ことがあるため、それを避けることをEfficientADでは意図しています。また、ネットワークを小さくすることで、高速に動作させることが可能となっています。異常検知においても、部品は正常だが誤った順番に並んでいるといった、画像全体の特徴を捉えないと検出できない異常が存在しますが、EfficientADでは後述のオートエンコーダーを用いたネットワークにこのような異常の検出をしてもらうことを意図しています。

PDNのパラメータは、Student-Teacherフレームワークを用いた蒸留によって学習します。具体的にはImageNetで学習したWideResNet-101を用いて、ImageNetの画像を入力したときの、PDNとWideResNet-101の出力の平均二乗誤差が小さくなるようにパラメータを学習します。

特殊な損失関数を用いた、局所的な異常検知のためのモデル構築

PDNを構築したら、その構築したStudentモデルを今度はTeacherモデルとして、蒸留によって局所的な異常検知のためのモデルを構築していきます。通常、蒸留では小さなStudentモデルでもTeacherモデルの出力をできるだけ模倣できるように学習していきますが、この手法においては、A.正常画像を入力としたときは模倣できるようにしたい一方で、B.異常画像を入力としたときは模倣できないようにしたいため、特殊な損失関数を用いて学習していきます。

まず、Studentモデルが出力を模倣したときに、Teacherモデルの出力と大きく差異がある部分（＝正常画像のうち、特に注目してほしい部分）に注目してもらうための損失関数を導入します。これは、正常画像のうち、背景等の特に注目する必要のない部分まで学習しすぎると、推論時に撮像環境の違い等によって偽陽性が発生する可能性が高まるからです。この損失関数は、入力画像を$I$として、Teacherモデルの出力$T(I)$とStudentモデルの出力$S(I)$の差分を$D_{c,w,h}=\left( T\left( I\right) _{c,w,h} - S\left( I\right) _{c,w,h}\right) ^{2}$と表したときに、差分の大きな上位0.1%までの値の平均として計算されます。

図2に、この損失関数を用いてトレーニングする際に、モデルが注目している場所を表した画像を示します。上段が入力されたトレーニング画像、下段の白い部分がモデルが注目している場所です。画像が示している通り、この損失関数を用いることで、背景のような比較的早期にTeacherモデルに近い出力が得られる部分はそれ以上学習しないようにし、物体の部分にフォーカスして学習を進めることができるようになります。これにより、背景等の異常検知にあまり関係のない部分の過学習を抑制でき、結果として撮像環境の違い等による偽陽性を抑止できることが期待されます。

図2: Hard feature lossを用いた際のモデルが注目する場所の例（論文より引用）

次に、正常画像にないような画像が入力された際に、StudentモデルがTeacherモデルの模倣ができなくなるような損失関数を導入します。TeacherモデルはImageNetで事前学習しているため、このような損失関数を導入しないと、異常画像を入力としてもTeacherモデルの模倣ができてしまう可能性があります。そこで、各トレーニングステップにおいて、ImageNetの画像$P$を入力として、以下の損失を加えて$L_{ST}$を算出します。つまり、正常画像の分布外にあるようなデータが入力された場合には特徴ベクトルのノルムが0になるように学習することで、Studentモデルが正常画像の分布外にあるようなデータの特徴を模倣できないようにします。

オートエンコーダーを用いた、論理的な異常検知のためのモデル構築

ここまでで、①PDN、②局所的な異常を検知するためのStudentモデルの2つのモデルが構築されました。推論時には、①と②の出力の差分で局所的な異常が検知できることになります。一方で、異常には、部品は正常だが誤った順番に並んでいる、のような論理的な異常があり、これらの異常も検知したいです。このような異常はオートエンコーダーを用いると比較的検知できることが経験的に知られており、EfficientADもオートエンコーダーを用いてこれらの異常を検知することを図っています。オートエンコーダーが論理的な異常を検知しやすい理由としては、オートエンコーダーは入力をダウンサンプリングした後にアップサンプリングするため受容野が大きく、入力画像全体を捉えた特徴を抽出できるからだと考えられます。

オートエンコーダー$A$は入力画像$I$に対して$A\left( I\right) \in \mathbb{R} ^{C\times W\times H}$の出力をするように構成され、以下の損失関数を用いてトレーニングされます。

入力画像とオートエンコーダーの出力の差分から異常スコアを算出することも可能ですが、オートエンコーダーでは入力画像をダウンサンプリングするため、局所的な特徴が潰される傾向にあります。そうすると、正常な網目などといったテクスチャ等をうまく復元できず、偽陽性の原因となってしまいます。そこで、EfficientADでは、オートエンコーダーの出力との差分を取るためのStudentモデル$S'$を構築します。これは、前述のStudentモデルのチャネル数を2倍にすることで併せて学習することができます。こちらのStudentモデルではオートエンコーダーを模倣したいため、以下の損失関数を導入します。

最終的に、トレーニング時には$L_{\mathrm{AE}}$、$L_{\mathrm{ST}}$、$L_{\mathrm{STAE}}$を合計した損失が使われることになります。

これによりStudentモデルはオートエンコーダーの模倣が一定程度できるようになりますが、受容野が小さいため、画像全体の特徴を捉える必要のある画像に対しては原理的にオートエンコーダーの模倣ができないことに留意します。この特性により、全体の特徴を捉える必要があるような論理的な異常がある画像に対して、オートエンコーダーとStudentモデルで出力に差異が生じ、異常が検知できることになります。

異常マップの正規化

ここまでで、①PDN、②局所的な異常を検知するためのStudentモデル、③論理的な異常を検知するためのStudentモデル、④オートエンコーダーの計4つのモデルが構築されることになります。推論時には、①と②の出力の差分で局所的な異常を検知し、③と④の出力の差分で論理的な異常を検知し、それらを合計することで最終的な異常レベルを出力したいのですが、これらの出力のスケールが合っているとは限りません。スケールが合っていないと、片方の出力のみが重視されることになるので、正規化をしてスケールを合わせる必要があります。EfficientADでは、トレーニングデータセットの画像の一部をモデルの学習ではなくこの正規化のために使います。

正規化のプロセスでは、トレーニングデータセットの画像の一部を使って、①と②の出力の差分の分布と、③と④の出力の差分の分布を得ます。その後、差分の大きさの特定の2点のクオンタイルにおける値の大きさが同一になるように値を線形変換することでスケールを合わせます。

図3に、モデル全体の動作イメージを示します。左側は、ケーブルにメタルワッシャーが付いているという異常画像を入力した際の、①、②、③、④それぞれの出力イメージと差分を表しています。このような局所的な異常の特徴は、③、④では潰れてしまって捉えられていない一方で、②で再現できないことから、①と②で大きな差異が生じ、結果として異常を検知できていることが分かります。次に、右側は、ケーブルが2本付いているという論理的な異常がある画像を入力した際の①、②、③、④それぞれの出力イメージと差分を表しています。このような異常は局所的な領域だけを見ると異常には見えないため、①、②でほとんど差異が生じず、結果として異常を検知できていない一方で、④が画像の再構成に失敗していることから、③と④で大きな差異が生じ、結果として異常を検知できていることが分かります。

図3: モデル全体の動作イメージ（論文より引用）

性能比較

論文中では、MVTec ADを対象とした実験で、畳み込み層を6つ用いたネットワークをトレーニングした際の平均AUROCが99.1であり、PatchCoreを、Wide-ResNet50-2等をアンサンブルさせて動作させた際のAUROCの99.3に近しい値が出ていると主張されています。また、局所的な異常だけでなく論理的な異常が含まれたデータセットであるMVTec LOCO ADに対する平均AUROCは90.7であり、他の手法より高い値を達成していると主張されています。

PatchCore等の、学習済み深層学習モデルをそのまま特徴抽出モデルとして用いる手法と大きく異なる点は、トレーニング時間と推論時間です。EfficientADでは、トレーニング時にニューラルネットワークのパラメータを学習させる必要があります。一方で、PatchCoreでは特徴ベクトルを抽出し、貪欲法に基づき特徴ベクトルのサブセットを抽出するだけでトレーニングが終了するため、一般にトレーニング時間はPatchCoreの方が速いと言えます。一方で、推論時間に関しては、EfficientADは小さなモデルを使うため、RTX 2080 Tiを搭載したマシンでレイテンシが数msと、PatchCoreと比較して10倍以上高速に動作します。

また、これらの手法の差異として、トレーニングの相対的な難しさがEfficientADにはあるだろうと考えられます。EfficientADでは、計4つのモデルを学習するとともに、異常マップの正規化等の推論のためのハイパーパラメータを決める必要があります。従って、PatchCore等の手法と比べると異常検知をするために決めなければならないハイパーパラメータが増えることから、ハイパーパラメータチューニングに必要な実験の回数が増えることが想定されます。従って、EfficientADは、推論を高速で動作させる必要があったり、論理的な異常を検知する必要があったりする場合に、モデルの構築に時間をかけることがそれほど問題にならない場合に導入を検討するのが良いと考えられます。

まとめ

この記事では、学習済み深層学習モデルを用いた外観検査向け異常検知手法であるEfficientADを紹介しました。この手法では、モデルの蒸留とオートエンコーダーを用いることで、推論時に高速に動作することと、キズ等の局所的な異常だけでなく論理的な異常を検出することを図っています。

ところで、LeapMindでは異常検知モデルだったり、異常検知とは全く関係ないモデルだったりを開発する機械学習エンジニアを絶賛募集中です！興味のある方は以下からぜひご応募ください！

https://leapmind.io/careers/