[21′ ICCV] DRAEM : A discriminatively trained reconstruction embedding for surface anomaly detection

내용 요약

DRAEM은 Reconstruction 기반의 Anomaly Detection 방법론입니다. 기존 방식과는 달리 비정상 이미지를 정상이미지로 Reconstruction 하는 방법을 사용합니다. 이를 위해 비정상 이미지를 합성해 내는 방법을 제안합니다. 또한 Discriminative Network를 추가하여 Defect Segmentation을 학습합니다. DRAEM은 mvTec, DAGM 데이터셋 실험을 통해 기존 방법들보다 성능이 좋음을 보였습니다. 다양한 Ablation 실험을 통해 제안 모듈의 타당성을 입증합니다.

1. 들어가며

이번 글에서는 2021년 ICCV에 발표된 DRÆM – A discriminatively trained reconstruction embedding for surface anomaly detection 논문을 리뷰합니다. 이 논문은 DRAEM이라는 별칭으로 불리며, 이번 글에서도 DRAEM으로 지칭하겠습니다. Anomaly Detection에 관한 기본적인 내용은 Anomaly Detection에 관한 모든 것 글을 참고해주세요.

본격적으로 들어가기 전, 제목의 의미를 통해 저자들이 강조하고자 한 내용을 유추해 보겠습니다. 제목에서는 크게 세 가지 포인트를 찾을 수 있습니다.
첫 번째는 Surface Anomaly Detection 방법론이라는 것입니다.
두 번째는 Discriminately Trained Reconstruction 방법을 사용하는 방법론이라는 것입니다.
이를 통해 Reconstruction 방식의 Anomaly Detection에 관한 내용임을 유추할 수 있습니다. 다만 기존의 Reconstruction 방법과는 달리 Discriminate 한 방법이 추가되었음을 알 수 있습니다. 본문에서 하나씩 자세한 내용을 살펴보겠습니다.

이번 글의 순서는 다음과 같습니다.
첫 번째로 기존 Reconstruction Anomaly Detection 방법들의 문제점을 살펴봅니다.
두 번째로 DRAEM 방법을 살펴봅니다. DRAEM의 큰 그림을 살펴본 뒤 Anomaly를 만드는 방법과 Architecture, Loss Function, Anomal Score 측정 방법을 살펴봅니다.
세 번째로 성능 실험 결과를 통해 DRAEM의 효과를 살펴봅니다.
네 번째로 Ablation 실험 결과를 통해 DRAEM 방법론의 타당성을 살펴봅니다.

2. 기존 방법의 문제점

초창기 Anomaly Detection에서는 Reconstruction 방법이 주를 이뤘었습니다. Reconstruction 기반의 Anomaly Detection 방법은 보통 이렇게 구성되어 있습니다.

학습용 정상 데이터를 Auto Encoder 등을 사용하여 동일한 정상 이미지로 Reconstruction 하도록 학습하는 것이죠. 이때 Anomaly Detection 특성상 모델은 정상 이미지만 학습하게 됩니다. 비정상 이미지는 학습 데이터로 사용할 수 없기 때문이죠.

그렇다면 정상 이미지로만 Reconstruction 하도록 학습한 모델이 비정상 이미지를 입력받으면 어떻게 될까요? 제대로 Reconstruction 하지 못하겠죠? 모델 입장에서는 미분만 공부했는데 갑자기 적분 문제를 풀라고 한 것과 같은 상황이니까요.

따라서 입력 이미지와 Reconstruction 이미지의 차이를 근거로 Anomal Score를 측정하자는 게 기본 철학입니다. 그림으로 보면 이렇습니다.

왼쪽은 Reconstruction Anomaly Detection 방법이 의도한 상황입니다.
입력 이미지가 정상일 때는 입력 이미지와 Reconstruction 이미지가 거의 비슷할 겁니다. 따라서 둘 사이의 거리가 가깝고요. 반면 입력 이미지가 비정상일 때는 입력 이미지와 Reconstruction 이미지의 차이가 클 겁니다. 따라서 둘 사이의 거리가 멀겠죠.

하지만 실제로 모델을 학습해 보면 의도한 것과는 다르게 반응하는데요. 오른쪽 그림은 실제 상황을 표현한 것입니다. 정상 입력 이미지에 대해서는 의도한 대로 반응합니다. 입력 이미지와 Reconstruction 된 이미지의 차이가 크지 않죠.
문제는 비정상 이미지가 입력으로 들어갔을 때입니다. 의도와는 달리 비정상 이미지도 잘 Reconstruction 하는 모습을 볼 수가 있습니다. 위의 비유로 다시 생각해 보면 미분만 공부했지만 적분까지 문제를 척척 잘 푸는 상황인 거죠.
이는 CNN의 뛰어난 학습 능력 때문인데요. CNN의 표현 능력이 워낙 크다 보니 한 번도 못 본 특성의 이미지도 잘 Reconstruction 할 수 있습니다.

3. 제안 방법

이제 위 문제를 개선한 DRAEM의 방법을 살펴보겠습니다.

3-1. 큰 그림

먼저 DRAEM의 큰 그림을 살펴볼게요.

위 그림은 DRAEM의 핵심 내용을 표현한 것입니다.
기본적으로는 이미지를 Reconstruction 하도록 학습한다는 점에서는 기존 방식과 동일합니다. 그런데 차이점이 있다면 정상 이미지를 Reconstruction 하는 게 아닌, 비정상 이미지를 정상 이미지로 Reconstruction 하는 겁니다. 이 과정에서 정상 이미지와 비정상 이미지의 차이점을 위주로 학습하겠죠.
여기에 추가로 비정상 부분에 더 초점을 맞추기 위한 트릭을 추가해 줍니다. 바로 Segmentation 네트워크를 추가해 주는 거죠. 이 Segmentation 네트워크는 Discriminative 네트워크라고 부르며 비정상 부분을 Segmentation 하도록 학습합니다. 이렇게 되면 모델은 온전히 비정상 부분만 찾아내도록 학습해 줄 수 있겠죠.

DRAEM의 전체 구조를 저자들의 그림으로 표현하면 다음과 같습니다.

가장 왼쪽에 비정상 이미지가 입력으로 들어가는 모습을 볼 수 있습니다. 이를 Reconstructive 네트워크를 사용해 정상 이미지로 Reconstruction 한 모습입니다. 그리고 Discriminative 네트워크는 비정상 부분을 Segmentation 하도록 학습하는 모습입니다.

3-2. Anomaly 만들기

큰 틀에서는 이해가 가는데요. 한 가지 중요한 설명이 빠져있습니다. 어떻게 비정상 이미지를 학습할 수 있었을까요?
Anomaly Detection 문제에서는 기본적으로 비정상 이미지를 학습 데이터로 제공하지 않는데요. DRAEM에서는 이 문제를 해결하기 위해 비정상 이미지를 합성해 냅니다. 실제 비정상 이미지를 사용할 수 없으니 최대한 비슷한 비정상 이미지를 만들어서 학습하는 거죠. 이렇게 DRAEM에서 비정상 이미지를 만들어내는 과정은 다음과 같습니다.

먼저 Perlin Noise Generator로 비정상 모양을 만들어 냅니다. (P)
그리고 Threshold를 정해 일정 부분만을 골라냅니다. (Ma)
학습 데이터에서 샘플 이미지를 선택합니다. (I)
Augmentation 할 이미지를 선택합니다 (A)
Ma와 A를 결합하여 원본 이미지 I에 추가한 최종 비정상 이미지를 합성해 냅니다. (Ia)

이러한 방식으로 합성해 낸 mvTec 데이터셋의 비정상 이미지 샘플을 보겠습니다.

가운데 그림은 Augmentation에 사용된 이미지입니다.
오른쪽 그림은 최종 합성된 비정상 mvTec 데이터셋입니다. 기존 mvTec 특성을 거의 유지하면서 Noise에 가까운 비정상 부분이 추가된 모습입니다.

3-3. Loss Function

다음은 DRAEM의 Loss Function을 살펴보겠습니다.

먼저 Reconstructive 네트워크에서 사용되는 Loss Function을 생각해 보겠습니다.
Reconstructive 네트워크에서는 비정상 이미지를 정상 이미지로 Reconstruction 해야 하는데요. 따라서 Loss로는 L2와 SSIM Loss를 사용합니다.

다음으로 Discriminative 네트워크 입장에서 생각해 보겠습니다.
Discriminative 네트워크는 비정상 부분을 Segmentation 해야 하죠. 따라서 픽셀별 Cross Entropy Loss를 적용할 겁니다.

그런데 단지 이렇게만 적용해 주면 문제가 있습니다. DRAEM이 만들어낸 비정상 이미지는 대부분이 정상이고 조금씩 비정상 부분이 섞여 있는 형태입니다. Class Imbalance 상황인 것이죠. 따라서 Discriminative 네트워크는 Focal Loss를 사용하여 학습합니다.

최종 Loss는 Reconstructive 네트워크의 Loss와 Discriminative 네트워크의 Loss를 합쳐 구성해 줍니다.

3-4. Anomal Score

다음으로 DRAEM이 Anomal Score를 측정하는 방법을 살펴보겠습니다.

Discriminative 네트워크는 비정상 부분을 Segmentation 하도록 학습했는데요. 그 말은 Discriminative 네트워크의 출력 그 자체가 비정상 부분이라는 것이죠. 따라서 Discriminative 네트워크의 출력을 Smoothing 하여  Anomal Score로 사용합니다. 수식으로는 이렇게 표현합니다.

• M₀ : Defect Segmentation Map
• f : Smoothing Filter

4. 효과

다음은 성능 실험 결과를 통해 DRAEM의 효과를 살펴보겠습니다.

4-1. mvTec 데이터셋

먼저 mvTec 데이터셋 Image Level 성능은 다음과 같습니다.

• [11] : PaDiM
• [20] : Mahalanobis AD

위 결과에서 주목해야 할 부분은 padim, mahalanobis AD의 성능입니다. 이 둘은 DRAEM 전까지 SOTA 방법들인데요. Pretrained Feature를 그대로 사용하는 방법이라는 공통점이 있습니다. DRAEM은 Pretrained Feature보다 비정상을 만들어 학습한 Feature가 더 좋은 구분력을 가지고 있음을 보였습니다.

Pixel Level 성능은 다음과 같습니다.

마찬가지로 거의 가장 좋은 성능을 보입니다.

4-2. DAGM 데이터셋

DAGM 데이터셋에 대한 성능은 다음과 같습니다.

위쪽은 Unsupervised Anomaly Detection 방법들의 성능입니다. Unsupervised 방식 중에서는 가장 좋은 성능을 보입니다.
아래쪽은 Supervised Anomaly Detection 방법들의 성능입니다. DRAEM은 Supervised 방식의 성능보다는 조금 떨어지지만 거의 필적한 성능을 보입니다.

5. Ablations

다음은 Ablation 실험 결과를 살펴보겠습니다.

Ablation 실험에서는 DRAEM 모듈을 하나씩 변경해 가며 성능을 비교합니다. 이를 통해 DRAEM 모듈의 타당성을 검증합니다.

5-1. Reconstruction Network 효과

먼저 Reconstruction Network를 제외한 뒤 성능을 비교해 보겠습니다.

• Disc : Reconstruction Network를 제외한 방법

이 말은 Defect Segmentation만 학습했을 때의 성능을 의미합니다. DRAEM 방법과 비교했을 때 성능이 크게 하락하는 모습을 볼 수 있습니다. 따라서 비정상 이미지를 정상 이미지로 Reconstruction 하는 과정은 모델에게 도움이 됨을 알 수 있습니다.

5-2. Discriminative Network 효과

다음은 Discriminative Network를 제외한 뒤 성능을 비교해 보겠습니다.

• Recon-AE : Discriminative Network 없이 Reconstruction 네트워크만 사용

즉 비정상 이미지를 정상 이미지로 Reconstruction만 학습한 네트워크인데요. DRAEM과 비교했을 때 성능이 크게 낮은 모습을 볼 수 있습니다. 따라서 Defect를 Segmentation 하도록 학습하는 과정은 모델에게 도움이 됨을 알 수 있습니다.

5-3. Anomaly 특성에 따른 효과

다음은 다양한 특성의 Anomaly를 사용한 성능을 비교해 보겠습니다.

• DRAEM ImageNet : Augmentation 입력으로 ImageNet 이미지 사용
• DRAEM color : Augmentation 입력으로 다양한 색상 사용
• DRAEM rect : Augmentation 입력으로 직사각형 모양 사용

위 결과를 보면 ImageNet을 사용했을 때 DRAEM과 성능이 가장 비슷한 모습을 볼 수 있습니다. 반면 단순한 color Augmentation 방법은 성능이 크게 하락하는 모습입니다.

이를 통해 두 가지 교훈을 얻을 수 있습니다.
첫 번째는 Anomaly를 합성하여 DRAEM 방식으로 학습하면 성능이 좋다는 것입니다. DRAEM color, DRAEM rect는 실제 Anomaly와는 매우 다른 특성의 Anomaly를 학습했습니다.

그럼에도 불구하고 꽤 좋은 성능을 보였죠. 이를 통해 실제와 다른 Anomaly라도 구분하도록 학습하면 모델에게 도움이 됨을 알 수 있습니다.
두 번째는 실제와 비슷한 Anomaly를 만들어 학습할수록 성능이 좋다는 것입니다. DRAEM에서 사용한 Perlin Noise는 실제와 가장 유사한 Anomaly를 합성해 냅니다. 그 결과 color, rect Augmentation 방법보다 ImageNet, Perlin 방법의 성능이 훨씬 좋았죠.
이를 통해 실제와 비슷한 Anomaly를 구분하도록 학습할수록 모델에게 도움이 됨을 알 수 있습니다.

5-4. Low Perturbation 효과

다음은 Augmentation 합성 강도에 따른 효과를 비교해 보겠습니다.

위 결과를 보면 큰 성능 차이를 보이지 않는 모습입니다. 따라서 Perturbation에 대한 영향은 크지 않음을 알 수 있습니다.

5. 결론

지금까지 DRAEM 내용을 살펴봤습니다.
DRAEM은 Reconstruction 방식의 Anomaly Detection 방법론입니다. 기존 Reconstruction 방법들과는 달리 비정상 이미지를 정상 이미지로 Reconstruction 하도록 학습했습니다. 이를 위해 Anomaly를 합성하는 방법을 사용했는데요. Perlin Noise와 이미지를 합성하여 원본 이미지에 추가하는 방법을 적용했습니다. 또한 비정상 특징에만 집중할 수 있도록 Discriminative Network를 추가해 주었습니다. Discriminative Network는 합성한 Defect를 Segmentation 하도록 학습했죠.
DRAEM은 mvTec과 DAGM 데이터셋에서의 성능 실험을 통해 기존 방법들보다 우수함을 보였습니다. 또한 다양한 Ablation 실험을 통해 DRAEM 제안 모듈의 타당성을 입증했습니다.
DRAEM은 비정상 이미지를 합성하여 Anomaly Detection 성능을 향상할 방법론을 제안했다는 의의가 있습니다. 반면 mvTec, DAGM 등 Variation이 작은 데이터셋에서만 성능을 검증했다는 한계가 있습니다.

Ffightingseok