Dinomaly: The Less Is More Philosophy in Multi-Class Unsupervised Anomaly Detection (CVPR 2025)

저번 포스트에서는 anomaly detection task 를 소개하면서 의료 영상에서의 활용 사례를 살펴보았습니다. 그런데, 이전 연구에서는 번거로운 점이 하나 있었습니다. Brain MRI 와 wrist X-ray 에 대해 따로 모델을 훈련한 것입니다. 하지만, MVTec 데이터셋 같이 다양한 종류가 있을 때 (무려 15 classes!), 혹은 그 이상의 클래스가 있을 때 그 갯수만큼 모델을 훈련할 수 있을까요? 시간과 자원이 엄청 낭비될 일입니다.

???: 그럼 그냥 모든 정상 이미지를 넣어놓고 하나의 모델만 쓰면 되는거 아님?

맞습니다. 여러 클래스에 대해 anomaly detection 을 하는 task 를 multi-class anomaly detection (MUAD) 라고 하는데, 학계에선 이 문제를 해결하기 위해 unified model 을 꾸준히 개발해왔습니다. 그럼에도 불구하고, MUAD 와 class-separated unsupervised anomaly detection (UAD) 사이에는 "non-negligible gap" 이 존재했습니다.

그럼 Dinomaly는 이 간극을 어떻게 좁힐까요?

Dinomaly

먼저, 왜 dinomaly 일까요? DINOv2 로 사전훈련된 ViT를 썼기 때문입니다 (...). 여기에 4가지의 engineering component 를 추가했다고 합니다:

1. Scalable foundational Transformers
2. Noise Bottleneck
3. Linear Attention
4. Loose Reconstruction

하나씩 살펴봅시다...

Architecture

Dinomaly 는 attention 과 MLP 만 사용했다고 합니다. 아래 피규어를 봅시다.

Encoder

인코더는 DINOv2-Register 로 사전 훈련된 ViT-Base/14 모델을 활용했습니다. 총 12 layers 중 마지막 두개는 버리고, weight 을 얼린 채로 활용합니다. 총 3 그룹의 벡터가 추출되는데, 3-6 번째 layer 의 합, 7-10 번째 layer 의 합, 그리고 3-10번째 layer 의 합이 다음 단계로 넘어갑니다.

Q: DINOv2-Register 이 뭐에요?

DINOv2 는 LVD-142M (images only) 에 사전훈련된 self-supervised learning scheme입니다.
다음과 같은 student-teacher architecture를 활용합니다:

- Teacher network: updated as student 의 exponential moving average (EMA)
- Student network: learns to match teacher's feature representations

두 네트워크 모두 같은 image 의 augmentation (crop) 이 입력으로 들어갑니다.
이 때, teacher network 는 global crop (큰 crop) 만 받는 반면, student network 는
각종 크기의 crop 을 입력받을 수 있습니다.

두 network 모두 random initialization 에서 시작하며, training objectives 는
다음과 같습니다:

- Image-level: teacher network 와 student network 의 global image embedding
  간의 cross entropy loss
- Patch-level: student network 에 입력한 patch 를 일부 masking 하고, 그 patch의
  feature 과 teacher network 의 patch feature (masking 안됨) 사이의 cross entropy
- 그 외 등등...

그런데, 이런식으로 훈련된 ViT 모델은 attention map 이 굉장히 noisy 합니다 (아래 그림 왼쪽).

즉, ViT 모델이 엉뚱한 곳에 유별나게 높은 norm 을 보이는데, 그 원인은 ViT가 추출한 정보가 너무
많으니 별로 안 중요한 패치에다가 "임시 저장" 을 때려버렸기 때문이라고 추측합니다.

이를 해결하기 위해서 ViT 에 "레지스터" 토큰을 추가하여, 멀쩡한 이미지 패치 외에 따로 정보를 
저장할 토큰을 만들어 주었습니다. 이렇게 훈련된 DINOv2 는 attention map 이 굉장히 깔끔해진 
것을 확인할 수 있습니다 (오른쪽)

Register token (ICLR 2024 oral!!) 에 관한 더 자세한 설명은 추후 다시 다루겠습니다. 이 논문 재밌더라고요.

Bottleneck

Bottleneck block 은 MLP 로, 3-10 번째 layer 의 출력 vector 의 합이 입력됩니다.

논문에선 이전의 MUAD 가 실패했던 이유로 "identity mapping" 에 대해 설명합니다. 즉, 이상한 이미지가 들어와도 anomaly 를 그대로 복원하는 능력인데, 이는 인공신경망이 multi-class 데이터를 활용하게 되면서 한 개의 class 의 데이터만 학습했을 때 보다 일반화 능력이 강화되었기 때문입니다.

다른 연구에선 identity mapping 을 해결하기 위해 일부러 normal 데이터에 anomaly 를 추가해주었습니다. 하지만, 이는 augmentation 과 비슷하게 heuristic 한 접근이라, 모든 도메인에서 일반화하기는 힘든 접근입니다.

그 대신에, 저자들은 dropout 을 bottleneck layer 에 추가합니다. 원래 dropout 은 overfitting 을 해소하기 위해 만들어졌으나, 이러한 상황에서는 일부러 normal representation 에 perturbation 을 주는 것으로 "pseudo feature anomaly" 를 만들어 낼 수 있습니다.

Linear Attention (Decoder)

8개의 Transformer layer 로 이루어져 있습니다만, softmax attention 대신에 linear attention을 활용합니다. 그 이유는 linear attention 이 query 에 연관된 영역에 집중하지 못하고 attention 이 전역에 퍼지는 단점 (!!!) 때문이였는데요, 이걸 역이용해서 디코더가 anomaly 에 집중하지 말고 이미지 전체를 normal 로 복원하라는거죠.

(Figure 3) Softmax Attention vs. Linear Attention

실제로 decoder 의 각 layer 을 시각화 해 보았을 때, Softmax Attention 을 활용했을 때에는 anomaly 영역을 query 했을 때 그 영역만 높은 attention 이 있는 반면, Linear Attention 은 layer 이 점차 앞으로 올 수록 anomaly 의 흔적이 줄어드는 것을 볼 수 있습니다 (l=7 에는 그래도 조금 있는 느낌?)

Q: 왜 softmax 는 anomaly 에 잘 집중해요?
A: Softmax 는 exponential weighting 을 하게 되어서 activation 이 작은 녀석
   들은 신호가 suppress 됩니다.
   
Q: Linear Attention 의 computational complexity 가 O(Nd^2) 로 줄어든다는데?
A: 아래 참조

1. $(\phi(\bold{Q})\phi(\bold{K^T}))\bold{V}$

• $(N\times d)$ times $(d\times N) = N\times N$ matrix, $O(N^2d)$
• $(N\times N)$ times $(N\times d) = N\times d$ matrix, $O(N^2d)$

2. $\phi(\bold{Q})(\phi(\bold{K^T})\bold{V}$)

• (오른쪽 먼저) $(d\times N)$ times $(N\times d) = d\times d$ matrix, $O(Nd^2)$
• $(N\times d)$ times $(d\times d) = N\times d$ matrix, $O(Nd^2)$

뭐 근데.. N = 784 (28 x 28), d = 768 (token dimension) 인 마당에 그렇게 큰 차이가 있을까 싶긴 하다. 물론 줄일 수 있으면 좋지.

Loose Reconstruction

이 부분은 훈련중에만 사용됩니다. Encoder 과 decoder 으로 knowledge distillation 을 했다고 생각하시면 되는데, 아래의 reconstruction scheme 중에 논문이 고른 scheme 은 (e) 입니다:

(e) 2 Group:

• Group 1: encoder layers 3–6 vs decoder layers 1–4
• Group 2: encoder layers 7–10 vs decoder layers 5–8
• 각각 그룹 사이에서 cosine distance 줄이기

Layer-to-layer 로 할 수도 있었겠지만, 이렇게 했을 때 모델이 "identity mapping" 을 잘 학습한다고 하네요. Layer-to-layer 로 tight 한 constrain 을 넣는 대신, group-to-group 간의 constrain 을 넣어서 decoder 의 layer 들이 encoder 의 layer 들과 다르게 행동할 수 있도록 풀어준 느낌입니다.

또한, loss 자체에도 약간 완화를 해주는데, 아래 식을 보시죠.

$f_D(h,w)$ 가 feature point at h, w 인데, decoder이 너무 reconstruction 을 잘한 batch 의 상위 k% 는 gradient 를 10%로 떨어뜨리겠다는 것입니다. 즉, "외우지마!" 라는 것입니다 (솔직히 추잡해보이는데 결과를 봐보죠뭐).

Results

먼저 multi-class UAD settings performance 부터 봐봅시다. Image-level 에서 MVTec-AD 의 AUROC 를 98.6 -> 99.6 (1%) 으로, VisA에서 95.5 -> 98.7 (3.2%), Real-IAD 에서 86.4 -> 89.3 (2.9%) 로 성능을 올렸습니다. 이정도면 일단 제안한 방법이 성능을 올린다는 결론은 낼 수 있습니다.

다음은 Ablation study 입니다. 본문 테이블은 MVTec-AD 에서 실험을 진행하였는데, 데이터가 너무 쉬워서 성능이 다 비슷한 관계로 supplementary 에서 VisA 결과물을 가져왔습니다.

먼저 가장 눈에 띄는 것은 Noisy Bottleneck 입니다. 단독으로 AUROC를 95.81 에서 97.38 (+1.57) 끌어올렸습니다. 그 다음으로 Loosen Constraint 와 Loosen Loss 를 추가했을 때 각각 (+0.68) (+0.51) 올랐는데, Noisy Bottleneck 때문에 어려워진 reconstruction 과 시너지가 났다고 저자들은 해석합니다.

추가로 Linear Attention 이 0.16 올렸다고 나와있으나, 이는 통계적으로 유의미하지 않아보이며, 실제로도 이거때문에 성능을 올릴 수 있는지는 의문입니다.

Reconstruction constraint scheme 에 대한 ablation 도 제공했습니다. 주목할 점은, layer-to-layer 보다 group-to-group scheme 이 항상 성능이 높았지만, Image-level 에서는 1 group 이 2 groups 보다 성능이 더 좋았습니다. 저자들은 그룹을 한개만 쓰면 high level 과 low level feature 이 섞이게 되어 정확한 anomaly localization (pixel-level) 이 어려워 진다고 해석합니다.

Table 5 에서는 각 방법론에서 다른 resolution 을 사용해서 ablation 을 하는데, 이게 사실 NeurIPS 에 냈을 때 얘만 resolution 높은거 (392 x 392, 다른건 256 x 256임) 써서 성능 높은거 아니냐 해서 리젝먹은 적이 있어서 (...) 보강한 실험결과입니다. 흥미롭게도, 다른 방법론은 input의 resolution 이 올라갔을 때 오히려 성능이 떨어지는데, 이건 Dinomaly 가 392 x 392 에 맞게 architecture 을 설계한거 아니냐? 라는 의문이 해소되지는 않습니다.

마치며..

CVPR 2025 에 발표된 Dinomaly 를 읽어보았습니다. 이 논문은 어떤 혁신적인 방법론을 제시한다기보다는 단순한 engineering 을 많은 실험으로 검증한 느낌을 받았습니다. 두가지 얻어간 것이 있다면,

• Encoder-decoder based reconstruction anomaly detection 을 할 때 bottleneck 에 어떤 noise 를 주는 것은 도움이 될 수 있다
• Vision Transformers Need Registers 라는 논문을 알게 되었다

읽어주셔서 감사합니다 ^^