Layer by layer, module by module: Choose both for optimal OOD probing of ViT

이 논문은 사전 학습된 비전 트랜스포머 (ViT) 의 중간 레이어가 분포 변화에 따라 최적의 성능을 내는 위치가 달라지며, 심한 분포 변화 시에는 피드포워드 네트워크 내부 활성화가, 약한 변화 시에는 멀티헤드 셀프 어텐션의 정규화된 출력이 각각 최적의 OOD 탐지 지점임을 규명합니다.

핵심

이 논문은 **"AI 가 새로운 세상을 만나면, 어디를 봐야 가장 똑똑하게 반응할까?"**라는 질문에서 시작합니다.

요약하자면, 이 연구는 거대한 AI 모델 (비전 트랜스포머) 이 훈련된 데이터와 다른 새로운 데이터 (예: 눈이 내리거나 흐릿한 사진) 를 볼 때, 모델의 '마지막 단계'보다 '중간 단계'가 더 잘 작동한다는 사실을 발견했습니다. 그리고 단순히 중간 단계만 보는 게 아니라, 어떤 부품을 확인하느냐에 따라 결과가 달라진다는 놀라운 사실을 밝혀냈습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 배경: 거대한 도서관과 새로운 손님

상상해 보세요. 거대한 **AI 도서관 (기초 모델)**이 있습니다. 이 도서관은 수백만 권의 책 (이미지 데이터) 을 읽으며 지식을 쌓았습니다. 보통은 이 도서관의 **마지막 책장 (최종 레이어)**에 있는 요약본을 보면 모든 것을 가장 잘 이해할 수 있다고 믿었습니다.

하지만 문제는, 이 도서관에 **새로운 손님 (새로운 데이터)**이 들어왔을 때입니다.

• 정상적인 손님 (ID 데이터): 도서관이 익숙한 책과 비슷한 옷을 입은 손님. (예: 맑은 날 찍은 고양이 사진)
• 낯선 손님 (OOD 데이터): 도서관이 본 적 없는 옷을 입거나, 눈보라를 맞고 온 손님. (예: 흐릿하거나 눈이 덮인 고양이 사진)

연구진은 "이 낯선 손님을 볼 때, 마지막 책장의 요약본을 보는 게 정말 최선일까?"라고 의문을 품었습니다.

2. 핵심 발견 1: "마지막 책장은 너무 전문화되어 있다"

실험 결과, **새로운 손님 (데이터 편향)**이 들어오면 **마지막 책장 (최종 레이어)**은 오히려 당황해서 엉뚱한 대답을 하기 시작했습니다.

• 비유: 마지막 책장은 "우리가 배운 대로만 답해야 한다"는 강박에 사로잡혀 있습니다. 눈이 내리는 고양이를 볼 때, "이건 고양이가 아니야!"라고 외치거나 혼란스러워하는 것입니다.
• 반면, 중간 책장 (중간 레이어): 아직 전문화되지 않아서 유연합니다. "아, 고양이는 고양이인데 눈이 좀 묻어 있구나"라고 유연하게 받아들이고 더 정확한 정보를 제공합니다.

결론: 데이터가 훈련 데이터와 많이 다를수록, 중간 단계를 보는 것이 마지막 단계보다 훨씬 똑똑합니다.

3. 핵심 발견 2: "중간 단계의 어떤 부품을 봐야 할까?"

그렇다면 중간 단계의 어떤 부분을 봐야 할까요? 연구진은 트랜스포머 모델이라는 복잡한 기계의 내부 부품을 하나하나 뜯어봤습니다.

이 기계는 크게 두 가지 작업을 반복합니다:

1. 주목하기 (Attention): "이 부분이 중요해!"라고 집중하는 작업.
2. 생각하기 (Feedforward): 집중된 정보를 바탕으로 깊이 생각해보는 작업.

연구진은 이 두 작업의 중간 과정을 살펴봤습니다.

• 상황 A: 눈이 많이 내린 날 (데이터 편향이 심함)

  • 최고의 비법: "생각하기" 작업이 끝난 직후, **활성화 (Activation)**된 상태를 보는 것입니다.
  • 비유: 이는 마치 "생각을 정리해서 핵심 아이디어를 뽑아낸 직후"를 보는 것과 같습니다. 소음 (눈) 이 많을 때, 핵심 아이디어를 필터링한 상태가 가장 선명하게 들립니다.
  • 주의: "생각하기"를 끝내고 다시 압축하는 과정 (FC2) 은 피하세요. 정보가 너무 줄어서 뭉개져버립니다.

• 상황 B: 날씨가 좋은 날 (데이터 편향이 약함)

  • 최고의 비법: "생각하기"를 시작하기 전, **정규화 (LayerNorm)**된 상태를 보는 것입니다.
  • 비유: 날씨가 좋으면 너무 깊게 생각할 필요 없이, 깔끔하게 정리된 초기 상태만 봐도 충분합니다.

4. 요약: 우리가 배운 교훈

이 논문의 결론은 매우 실용적입니다.

1. AI 를 사용할 때, 무조건 마지막 결과를 믿지 마세요. 특히 환경이 변하거나 (날씨가 나빠지거나, 새로운 도메인), 데이터가 훈련 데이터와 다를 때는 중간 단계를 확인하는 것이 더 안전합니다.
2. 어떤 중간 단계를 볼지는 상황에 따라 달라집니다.
   • 혼란스러운 상황 (OOD): "생각하기" 부품이 작동한 직후의 활성화 신호를 확인하세요. (가장 강력함)
   • 평범한 상황 (ID): 그냥 마지막 결과를 보거나, "생각하기" 전의 정리된 상태를 보세요.

한 줄 요약

"AI 가 낯선 상황을 마주할 때, 마지막 결론을 믿기보다 '생각하는 과정'의 중간에 있는 핵심 아이디어를 확인하는 것이 더 똑똑한 방법입니다."

이 연구는 AI 가 실전에서 더 튼튼하게 작동하도록, 우리가 어디를 집중해서 봐야 하는지에 대한 나침반을 제공해 줍니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 최근 연구들은 사전 학습된 기초 모델 (Foundation Models) 의 최종 층 (final layer) 보다는 중간 층 (intermediate layers) 이 더 우수한 표현력을 가질 수 있음을 보여주고 있습니다. 특히autoregressive(자기회귀) 모델에서 이 현상이 두드러지지만, Vision Transformer(ViT) 에서는 최종 층이 여전히 최적이라는 주장도 존재하여 논쟁이 있었습니다.
• 문제: ViT 를 사전 학습 데이터 (예: ImageNet) 와 다른 분포 (Out-of-Distribution, OOD) 의 다운스트림 작업에 적용할 때, 기존 연구들은 왜 중간 층이 더 잘 작동하는지, 그리고 어떤 모듈 (Module) 을 프로빙 (probing) 해야 하는지에 대한 명확한 가이드가 부족했습니다. 또한, 분포 변화 (Distribution Shift) 가 심할 때 최종 층의 성능이 급격히 저하되는 현상의 원인과 해결책을 규명할 필요가 있었습니다.

2. 연구 방법론 (Methodology)

• 실험 설정:
  • 모델: ImageNet-21k 로 사전 학습된 86M 파라미터 ViT (Vision Transformer) 사용.
  • 데이터셋: 11 개의 다양한 분류 벤치마크 (CIFAR-10/100, CIFAR-10-C 의 5 가지 변형, DomainNet 의 2 개 도메인, Flowers102, Pets) 를 활용하여 In-Distribution (ID) 과 Out-of-Distribution (OOD) 시나리오를 모두 테스트했습니다.
  • 평가 방식: Linear Probing (선형 프로빙) 을 사용했습니다. 각 층의 특정 지점에서 CLS 토큰 임베딩을 풀링 (pooling) 하고 로지스틱 회귀 (L-BFGS 솔버) 를 적용하여 분류 정확도를 측정했습니다.
• 세부 분석:
  1. Layer-wise 분석: Transformer 블록의 출력 (RC2, Residual Connection 2) 을 기준으로 각 층 (Layer) 의 깊이에 따른 성능 변화를 ID 와 OOD 환경에서 비교했습니다.
  2. Module-wise 분석: 각 Transformer 블록 내부의 8 가지 연산 (LN1, MHA, RC1, LN2, FC1, Act, FC2, RC2) 중 어느 지점에서 임베딩을 추출하는 것이 가장 효과적인지 미세하게 분석했습니다.

3. 주요 기여 및 발견 (Key Contributions & Findings)

A. 분포 변화 (Distribution Shift) 가 최종 층 성능 저하의 주원인

• 발견: 사전 학습 데이터와 다운스트림 데이터 간의 분포 변화가 클수록 (OOD), 네트워크의 깊은 층 (최종 층) 일수록 성능이 급격히 떨어지는 것을 확인했습니다.
• 반대 경우: 분포 변화가 없는 경우 (ID, 예: ImageNet 기반 데이터), 최종 층이 가장 높은 성능을 보였습니다.
• 결론: 중간 층의 우월성은 단순히 사전 학습 목표 (autoregression 등) 의 부산물이 아니라, 분포 변화에 대한 중간 층의 강건성 (Robustness) 때문입니다.

B. 모듈별 최적 프로빙 지점의 차이

기존의 표준 관행인 Transformer 블록의 최종 출력 (RC2) 을 프로빙하는 것은 OOD 환경에서 최적의 선택이 아닙니다. 분포 변화의 정도에 따라 최적의 모듈이 달라집니다.

1. 강한 분포 변화 (Strong Shift) 환경:

   • 최적 모듈: Feedforward Network (FFN) 의 활성화 (Act, GeLU) 직후.
   • 이유: FC1 이 차원을 4d 로 확장하고 Act 가 비선형성을 부여하며, FC2 가 다시 d 로 축소하기 전에 위치합니다. 이 지점의 표현은 고차원 공간에서의 특징 분리 (feature disentanglement) 가 잘 되어 있으며, 노이즈가 필터링된 상태라 선형 분류기에 가장 적합합니다.
   • 성능: FC2 (최종 FFN 출력) 는 오히려 가장 낮은 성능을 보였습니다.

2. 약한 분포 변화 (Weak Shift) 또는 ID 환경:

   • 최적 모듈: Feedforward Network 직전의 LayerNorm (LN2) 또는 블록 출력 (RC2).
   • 이유: 분포 변화가 적을 때는 표준적인 블록 출력이 여전히 유효하며, LN2 도 안정적인 성능을 보입니다.

4. 실험 결과 (Results)

• Layer-wise: Fig. 1 과 Fig. 5 에서 보듯, 분포 변화가 심한 데이터셋 (Speckle Noise, Contrast 등) 에서는 깊이가 깊어질수록 (Layer Depth 100%) 성능이 하락하는 반면, 중간 층 (약 50-70%) 에서 정점을 찍습니다.
• Module-wise: Table 1 과 Fig. 3 에서 보듯, OOD 데이터셋에서 Act (FFN 활성화) 모듈을 프로빙했을 때 평균적으로 가장 높은 정확도를 기록했습니다. 반면, FC2는 12 개 데이터셋 중 10 개에서 가장 낮은 정확도를 보였습니다.
• 정리:
  • ID 환경: 최종 층 > 중간 층.
  • OOD 환경: 중간 층의 FFN 활성화 (Act) > 최종 층 > 중간 층의 블록 출력 (RC2).

5. 의의 및 시사점 (Significance)

• 실무적 가이드: ViT 기반 모델을 OOD 환경에서 사용할 때, 미세 조정 (Fine-tuning) 이 불가능하거나 비용이 큰 경우, 어떤 층과 어떤 모듈을 추출하여 사용하는지에 따라 성능이 결정됩니다.
  • 분포 변화를 감지하기 어렵다면, 표준적인 RC2 대신 LN2를 프로빙하는 것이 안전한 대안입니다.
  • 분포 변화가 크다고 판단된다면, **중간 층의 FFN 활성화 (Act)**를 추출하는 것이 최적의 전략입니다.
• 이론적 통찰: Transformer 내부의 정보 흐름을 모듈 수준에서 분석함으로써, FFN 이 단순한 변환이 아닌 'Key-Value 메모리' 역할을 하며, 활성화 함수가 특징의 분리 (disentanglement) 에 핵심적인 역할을 함을 시사합니다.
• 향후 연구 방향: 분포 변화 감지 및 적응을 위한 효율적인 방법론 개발에 기여하며, 정보 이론적, 기하학적 측정을 통한 Transformer 내부 표현에 대한 추가 연구의 토대를 마련했습니다.

요약: 이 논문은 ViT 의 OOD 일반화 성능을 극대화하기 위해, 단순히 '최종 층'을 사용하는 것이 아니라 분포 변화의 강도에 따라 '중간 층'과 '특정 모듈 (FFN 활성화)'을 선택적으로 사용해야 함을 증명했습니다.