[논문 리뷰] Puzzle-AE: Novelty Detection in Images through Solving Puzzles

Mohammadreza Salehi, Ainaz Eftekhar외 3명|arXiv (Cornell University)|2020. 08. 29.

Anomaly Detection Techniques and Applications참고 문헌 51인용 수 14

한 줄 요약

Puzzle-AE는 이미지 퍼즐을 자가지도 학습 전문 과제로 삼아 U-Net을 훈련시켜 복원 품질과 정상 데이터에 대한 일반화 능력을 향상시키고 과적합을 방지하는 비정상 탐지 프레임워크를 제안한다. 이 방법은 히ュ리스틱한 조기 정지가 필요 없이 안정적이고 데이터 효율적인 훈련을 통해 CIFAR-10, MNIST 및 MVTecAD, 의료 영상과 같은 실세계 데이터셋에서 최신 기술 수준(SOTA)의 성능을 달성한다.

ABSTRACT

Autoencoder, as an essential part of many anomaly detection methods, is lacking flexibility on normal data in complex datasets. U-Net is proved to be effective for this purpose but overfits on the training data if trained by just using reconstruction error similar to other AE-based frameworks. Puzzle-solving, as a pretext task of self-supervised learning (SSL) methods, has earlier proved its ability in learning semantically meaningful features. We show that training U-Nets based on this task is an effective remedy that prevents overfitting and facilitates learning beyond pixel-level features. Shortcut solutions, however, are a big challenge in SSL tasks, including jigsaw puzzles. We propose adversarial robust training as an effective automatic shortcut removal. We achieve competitive or superior results compared to the State of the Art (SOTA) anomaly detection methods on various toy and real-world datasets. Unlike many competitors, the proposed framework is stable, fast, data-efficient, and does not require unprincipled early stopping.

연구 동기 및 목표

• CIFAR-10과 실세계 영상과 같은 복잡한 데이터셋에서 오토인코더(AE)의 과적합 및 낮은 복원 품질 문제를 해결하기 위해.
• 자기지도 학습(SSL)을 활용해 퍼즐 풀이를 전문 과제로 삼아 AE 기반 모델의 일반화 능력과 의미적 특징 학습 능력을 향상시키기 위해.
• 일반적인 퍼즐과 같은 전문 과제에서의 단기 학습을 방지하기 위해, 모델이 낮은 수준의 통계적 특징을 악용하는 대신 의미 있는 특징을 학습하도록 하기 위해.
• 히ュ리스틱한 초모수 조정이나 조기 정지를 필요로 하지 않는 안정적이고 데이터 효율적이며 재현 가능한 비정상 탐지 프레임워크를 개발하기 위해.
• 고정성 특이도(FPR = 99.5%)에서 높은 TPR를 확보하는 실용적인 평가 기준을 충족시키며, 토이 데이터셋(MNIST, CIFAR-10 등)과 실세계 데이터셋(MVTecAD, 의료 영상 등)에서 뛰어난 성능을 보여주기 위해.

제안 방법

• 입력은 조각으로 나누어진 뒤 뒤섞인 이미지이고, 출력은 원본 이미지로 복원되는 U-Net을 생성기로 훈련시어 4x4 이미지 퍼즐을 해결하도록 한다.
• 복원 품질 향상과 더불어 더 의미 있는 특징을 학습하도록 유도하기 위해 GAN 유사 훈련 설정에서 판별기(discriminator)를 사용한다.
• 모델이 가장자리나 밝기 기반 힌트를 악용할 수 있는 단기 솔루션을 자동으로 제거하기 위해 강력한 적대적 훈련(PGD 및 FGSM)을 적용한다.
• 색조 복원과 퍼즐 풀이를 결합한 다중 과제 전문 학습 전략을 도입하여 특징 추상화 능력과 모델의 강건성을 더욱 향상시킨다.
• 모델이 인간의 레이블 없이도 올바른 조각 배치를 예측하도록 자가지도 학습 목표를 설정한다.
• 복원 손실을 이상 탐지 점수로 사용한다: 이상 입력에서는 손실이 높게 나타나며, 정상 입력은 정확하게 복원된다.

실험 결과

연구 질문

• RQ1 자기지도 학습 전문 과제로 퍼즐 풀이를 활용할 경우, 오토인코더 기반 비정상 탐지 모델의 복원 품질과 일반화 능력 향상에 기여하는가?
• RQ2 특히 복잡한 데이터셋에서, 퍼즐 풀이 SSL 과제의 단기 솔루션을 제거하기 위해 적대적 훈련이 얼마나 효과적인가?
• RQ3 색조 복원과 같은 추가 SSL 과제를 퍼즐 풀이와 결합할 경우, 실세계 비정상 탐지 벤치마크에서 모델의 강건성과 성능 향상에 기여하는가?
• RQ4 제안된 프레임워크는 히ュ리스틱한 조기 정지나 광범위한 초모수 조정 없이도 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ5 의료 및 산업 결함 탐지 데이터셋에서 고정성 특이도(FPR = 99.5%) 조건 하에 높은 TPR 성능을 보이는가?

주요 결과

• 11개 데이터셋 벤치마크에서 Puzzle-AE는 1/12의 데이터로 훈련했을 때 평균 AUC 72.47을 기록했으며, LSA(88.89)와 DSVDD(92.61)를 크게 앞서며, 평균 AUC 94.90을 달성했다.
• MNIST에서 Puzzle-AE는 평균 AUC 94.90을 기록했고, LSA의 88.89과 DSVDD의 92.61을 초월하여 더 뛰어난 데이터 효율성과 성능을 입증했다.
• FPR = 99.5%에서 Puzzle-AE는 평균 TPR 0.2932를 기록했으며, LSA의 0.4261보다 뛰어나 높은 특이도 조건에서도 더 나은 탐지 능력을 보였다.
• 그레이스케일 입력에 대해 강건성을 보이며, CIFAR-10의 '자동차' 클래스에서 성능 저하가 최소한이었고, 색상에 대한 민감도가 감소하고 일반화 능력이 향상됨을 시사했다.
• PGD/FGSM을 활용한 적대적 훈련을 추가하면, 강력한 단기 솔루션이 존재하는 데이터셋(예: MNIST)에서 약 1% 향상되어 모델의 강건성이 향상됨을 확인했다.
• MVTecAD 및 의료 영상 데이터셋에서 최신 기술 수준의 결과를 달성했으며, 이는 이전의 일종의 SSL 방법인 GT조차도 능가했고, 조기 정지가 필요 없음을 입증했다.