Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

기술 요약: 소수 샷 의료 분류를 위한 클래스 대조적 영향력을 통한 확산 모델 제어

문제 정의
소수 샷(few-shot) 의료 영상 분류에서는 레이블이 지정된 데이터가 부족하기 때문에 데이터 증강이 필수적입니다. 오프더쉘(off-the-shelf) 확산 모델은 합성 훈련 샘플을 생성할 수 있지만, 기존 방식은 주로 생성된 이미지의 현실성, 다양성 또는 도메인 적응에 초점을 맞춥니다. 저자들은 이러한 지표들이 다운스트림 분류 작업의 유용성을 보장하지 않는다고 주장합니다. 생성된 이미지의 시각적 품질이 비슷하더라도 분류 성능은 크게 달라질 수 있는 상당한 불일치가 존재합니다. 핵심 문제는 현재의 관행이 생성 품질 자체의 문제가 아니라, 분류기가 결정 경계(decision boundary)를 정교화하는 데 구체적으로 도움이 되는 샘플을 생성하는 데 실패하고 있다는 점입니다.

방법론
제안된 솔루션인 **클래스 대조적 영향력(Class-Contrastive Influence, C2I)**은 최적화 목표를 이미지 품질에서 분류기에 대한 합성 샘플의 그래디언트 기반 영향력으로 전환합니다.

1. 클래스 대조적 영향력(C2I) 기준:

   • 이 방법은 그래디언트 기반 영향력 추정을 활용하여 훈련 샘플이 검증 손실(validation loss)에 어떻게 영향을 미치는지 정량화합니다.
   • 저자들은 가공되지 않은 클래스 불가지론적(class-agnostic) 영향력 점수가 분류 성능과 상관관계가 없음을 관찰했는데, 이는 검증 세트에 서로 다른 클래스로부터 오는 상충하는 신호가 포함되어 있기 때문입니다.
   • 대신, C2C는 **영향력 차이(influence gap)**를 기반으로 정의됩니다. 즉, 샘플이 동일한 클래스의 검증 그래디언트와는 양의 방향으로 일치하고, 다른 클래스의 그래디언트와는 음의 방향으로 대조되는 경우에 유용하다고 간주됩니다.
   • 이진 분류의 경우, C2I는 양성(positive) 및 음성(negative) 검증 세트의 평균 영향력 사이의 제곱 차이를 분산으로 나눈 값으로 공식화됩니다. 다중 클래스 작업의 경우, 모든 클래스에 대한 매칭 클래스의 평균 영향력에 대한 소프트맥스(softmax) 점수로 정의됩니다.

2. 샘플 난이도와의 이론적 연결:

   • 이론적 분석(Theorem 1)은 로지스틱 회귀 설정에서 C2I를 최대화하는 것이 생성된 샘플을 검증 세트의 전역 평균(global mean)으로 유도함을 보여줍니다.
   • 두 클래스 클러스터의 전역 평균은 일반적으로 결정 경계 근처에 위치하므로, C2I를 최대화하는 것은 본질적으로 "어려운(hard)" 예시, 즉 분류하기 어렵고 경계에 인접한 샘플을 선택하게 됩니다.
   • 이러한 어려운, 경계 근접 예시들로 훈련하는 것은 결정 경계를 정교화하여 일반화 성능과 강건성(robustness)을 개선하는 데 도움을 줍니다.

3. 강화 학습(RL) 미세 조정:

   • 이를 실행하기 위해 저자들은 사전 훈련된 텍스트-이미지(T2I) 확산 모델(특히 Stable Diffusion 2.1)을 위한 미세 조정 기법을 제안합니다.
   • 보상 함수(Reward Function): C2I 점수는 강화 학습 프레임워크 내에서 보상 신호 역할을 합니다. 생성기는 기대 C2I 보상을 최대화하도록 미세 조정됩니다.
   • 최적화: 저자들은 C2I 보상을 바탕으로 확산 모델의 파라미터(특히 LoRA 가중치)를 업데이트하기 위해 디노이징 확산 정책 최적화(Denoising Diffusion Policy Optimization, DDPO)를 사용합니다. 이는 생성기의 구조적 변경 없이 생성 과정을 "클래스 정보가 풍부한" 데이터 매니폴드 쪽으로 유도합니다.
   • 워크플로우: 먼저 분류기(ViT-B/16)를 소수 샷 훈련 세트로 미세 조정합니다. 그 후, 홀드아웃(held-out) 검증 세트에 대한 분류기의 그래디언트를 사용하여 생성된 배치에 대한 C2I 보상을 계산하고, 이를 통해 확산 모델의 RL 미세 조정을 가이드합니다.

주요 기여

• 효과성의 공식화: 본 논문은 C2I를 통해 샘플의 효과성을 공식화하며, 높은 C2I 샘플이 단순히 현실적인 샘플이 아니라 어려운 경계 근접 예시라는 점을 입증했습니다.
• RL 기반 생성 체계: 오프더쉘 확산 모델을 태스크에 유용한 클래스 정보가 풍부한 샘플을 생성하도록 유도하기 위해 C2I를 보상으로 사용하는 플러그 앤 플레이(plug-and-play) 방식의 RL 프레임워크를 도입했습니다.
• 경험적 검증: 제안된 방법은 여러 소수 샷 의료 영상 벤치마크(BreastMNIST, DermaMNIST, PneumoniaMNIST)에서 표준 증강 및 기타 확산 기반 베이스라인보다 일관된 개선을 보이며 검증되었습니다.

실험 결과
방법론은 소수 샷 설정(클래스당 16 또는 32개의 레이블된 샘플)에서 세 가지 MedMNIST 벤치마크에 대해 평가되었습니다.

• 성능: C2I 가이드 생성기는 표준 변환(RandAugment, RandomErasing, Mixup) 및 기타 확산 기반 증강 방법(DataDream, Dataset Expansion, DistDiff)을 포함한 베이스라인들을 일관되게 능가했습니다.
  • 평균적으로 제안된 방법은 다양한 백본(ViT-B/16 및 ResNet-18)에 대해 가장 높은 AUC를 달 기록했습니다.
  • 특히, 이 방법은 ResNet-18의 성능을 개선했는데, 이는 합성된 샘플이 광범위하게 정보를 제공하며 모델 아키텍처 간에 전이 가능하다는 것을 나타냅니다.
• 강건성: C2I 가이드 증강으로 훈련된 모델은 다양한 노이즈 조건(솔트-앤-페퍼, 가우시안 블러, JPEG 압축) 하에서 우수한 강건성을 보였으며, 이는 생성된 샘플이 더 안정적인 결정 경계를 형성함을 시사합니다.
• 난이도 상관관계: RL 미세 조정이 진행됨에 따라 C2I 보상이 증가하고 생성된 이미지의 특징이 검증 특징 중심(feature centroids)에 가까워진다는 실험 결과는, 더 어려운 예시들이 생성되고 있음을 확인시켜 줍니다.
• 어블레이션(Ablation): RL 훈련 에포크 수와 합성 이미지 수를 늘리면 다운스트림 AUC가 일반적으로 향상되었으나, Dataset Expansion과 같은 일부 베이스라인은 합성 데이터가 증가함에 따라 성능 저하를 보였습니다.

의의 및 주장
본 논문은 합성 증강이 이미지 품질뿐만 아니라 **태스크 유용성(task usefulness)**에 의해 가이드될 때 가장 효과적이라고 주장합니다. C2I는 범용적인 확산 모델을 레이블된 데이터가 가장 부족한 곳에서 결정 경계를 강화하는 표적 생성기로 변모시킵니다.

저자들은 자신의 방법이 원래의 훈련 데이터를 사용하는 것보다 추가적인 계산 비용(RL 미세 조정 과정으로 인해 A100에서 약 5 GPU 시간 소요)을 발생시킨다는 점을 인정합니다. 그러나 이러한 비용은 낮은 데이터 환경에서 일관된 정확도 및 강건성 이득을 통해 정당화된다고 주장합니다. 이 연구는 증강을 위한 확산 모델 사용에 있어 패러다임의 변화를 제안합니다: 즉, 생성 과정을 분류기에 대한 다운스트림 영향력에 맞춰 최적화하는 것입니다.