SHOT-CCR: Biologically guided adversarial training for test-time adaptation in cellular morphology

이 논문은 다양한 실험 배치에서 발생하는 기술적 편향을 제거하고 세포 형태학적 신호를 분리하기 위해 세포 불변 그래디언트 반전을 활용한 생물학적 기반 테스트 시간 적응 프레임워크인 SHOT-CCR 을 제안하여 RxRx1 과 JUMP-CP 데이터셋에서 기존 성능을 크게 상회하는 분류 정확도를 달성했다고 요약할 수 있습니다.

핵심

🧪 핵심 문제: "같은 실험인데 결과가 왜 달라?"

약물 개발 연구자들은 수백만 개의 세포를 현미경으로 찍어 (Cell Painting), 약이나 유전자 조작이 세포에 어떤 변화를 주는지 관찰합니다. 마치 세포의 '얼굴'을 사진으로 찍는 것과 같습니다.

하지만 여기서 큰 문제가 생깁니다.

• 비유: 같은 요리 레시피 (유전자 조작) 를 썼는데, **요리사 (실험자), 주방 (실험실), 사용된 냄비 (실험 배치)**가 다르면 요리의 맛 (세포 이미지) 이 미세하게 달라집니다.
• AI 는 이 미세한 차이 (배치 효과, Batch Effect) 를 '요리사의 특징'으로 착각하고, 실제 '요리 레시피 (약의 효과)'를 구분하지 못하게 됩니다. 마치 다른 식당에서 찍은 사진을 보고 같은 요리를 구분하지 못하는 것과 같습니다.

기존의 AI 모델은 같은 실험실 (배치) 안에서는 잘 작동하지만, 새로운 실험실 데이터가 들어오면 엉망이 되는 경우가 많았습니다.

---

💡 해결책: SHOT-CCR (세포 개수 역전 훈련)

이 논문은 SHOT-CCR이라는 새로운 방법을 제안합니다. 이 방법은 두 가지 핵심 아이디어를 섞었습니다.

1. "세포 개수"라는 함정을 피하라 (생물학적 가이드)

기존 연구들은 "배치 효과"라는 거대한 덩어리 전체를 지우려 했지만, 오히려 중요한 정보까지 잃어버리는 경우가 많았습니다.

• 비유: 사진의 배경 (배치) 을 지우려다 보니, 주인공 (세포의 변화) 까지 흐릿해져 버린 셈입니다.
• 저자들의 통찰: 그들은 "배치에 따라 세포의 개수가 달라지는 경향"에 집중했습니다. 실험실마다 세포가 너무 많거나 적게 찍히는 경우가 많기 때문입니다.
• 방법: AI 가 "세포가 몇 개나 있나?"를 맞추는 능력을 의도적으로 방해합니다. (이를 '경사 역전'이라고 합니다.)
  • 마치 시각 장애인을 위해 소음 (세포 개수 정보) 을 차단하고, 오직 '음색 (세포의 형태 변화)'에만 집중하게 만드는 것과 같습니다.
  • AI 는 "세포가 100 개냐 200 개냐"는 건 신경 쓰지 않고, "세포 모양이 어떻게 변했나?"에만 집중하게 됩니다.

2. 시험 시간에도 학습하라 (테스트 시간 적응, TTA)

기존 AI 는 학습을 끝내고 나면 더 이상 배우지 않습니다. 하지만 새로운 실험실 (새로운 배치) 의 데이터가 들어오면, 실제 시험 (테스트) 하는 동안에도 AI 가 스스로 적응하게 합니다.

• 비유: 새로운 식당에 갔을 때, 요리사가 "아, 여기는 소금기가 좀 적네"라고 바로 적응해서 요리를 다시 맛있게 만드는 것과 같습니다.
• AI 는 새로운 데이터의 특성을 빠르게 파악하여, 이전 학습 지식을 새로운 상황에 맞춰 조정합니다.

---

🏆 성과: 얼마나 잘해냈나?

이 방법은 두 가지 거대한 데이터셋 (RxRx1 과 JUMP-CP) 에서 놀라운 결과를 냈습니다.

1. 정확도 대폭 상승: 기존 최고의 기록 (87.1%) 을 깨고 **91.6%**의 정확도를 달성했습니다.
2. 어려운 세포도 해결: 특히 데이터가 적고 예측하기 어렵다고 알려진 U2OS 세포에서 성능이 가장 크게 향상되었습니다. (기존 68.2% → 76.2% 로 상승)
3. 다양한 세포에 적용: 네 가지 다른 종류의 세포에서도 일관되게 좋은 결과를 보여주었습니다.

---

🌟 요약 및 시사점

이 논문의 핵심 메시지는 **"모든 잡음을 지우려 하지 말고, 생물학적으로 중요한 '세포 개수'라는 특정 잡음만 골라내어 AI 를 훈련시키라"**는 것입니다.

• 기존 방식: "배치 효과"라는 거대한 벽을 무너뜨리려다 중요한 정보까지 잃음.
• 새로운 방식 (SHOT-CCR): "세포 개수"라는 구체적인 함정을 AI 가 피하도록 훈련시키고, 새로운 데이터가 들어오면 실시간으로 적응하게 함.

이 기술은 약물 개발 속도를 높이고, 더 적은 비용으로 더 정확한 실험 결과를 얻을 수 있게 도와줍니다. 마치 AI 가 실험실의 '소음'을 필터링하여, 진짜 '약의 효과'라는 신호를 선명하게 들어내는 귀를 갖게 해준 것과 같습니다.

기술 요약

1. 문제 정의 (Problem Statement)

• 배경: 고함량 스크리닝 (HCS) 및 'Cell Painting' 기술을 통해 생성된 대규모 세포 이미지 데이터는 약물 발견을 가속화하는 데 핵심적입니다.
• 핵심 문제: 실험 배치 (Batch) 간 기술적 차이 (배지, 장비, 시약 등) 로 인해 발생하는 배치 효과 (Batch Effects) 가 모델의 일반화 성능을 심각하게 저해합니다. 특히, 같은 생물학적 신호 (유전적 교란) 를 가진 데이터라도 실험 배치나 세포 유형이 다르면 모델이 이를 구분하지 못하거나 오분류하는 경향이 있습니다.
• 기존 방법의 한계: 기존 연구 (Sypetkowski et al., 2023) 는 적응형 배치 정규화 (AdaBN) 를 사용하여 일부 개선 효과를 보였으나, 데이터가 적은 세포 유형 (예: U2OS) 이나 복잡한 배치 간 차이에서는 성능이 크게 떨어졌습니다. 또한, 모든 배치 정보를 제거하려는 시도는 오히려 중요한 생물학적 신호까지 손실시킬 수 있습니다.

2. 제안 방법론: SHOT-CCR

저자들은 SHOT-CCR (Cell Count Reversal) 이라는 새로운 테스트 시간 적응 (Test-Time Adaptation, TTA) 프레임워크를 제안합니다. 이는 생물학적 지식을 기반으로 한 적대적 훈련을 통해 실험적 혼란 요인 (Confounders) 을 분리하는 것을 목표로 합니다.

A. 핵심 구성 요소

1. 생물학적 지향적 적대적 훈련 (Biologically Guided Adversarial Training):

   • 기존 연구가 '실험 배치 ID' 자체를 예측하여 제거하려 했던 것과 달리, 이 방법은 세포 수 (Cell Count) 라는 구체적인 생물학적 특징에 초점을 맞춥니다.
   • 세포 수는 실험 배치 간에 큰 편차를 보이지만, 교란 (Perturbation) 의 본질적 신호와는 무관할 수 있는 기술적 아티팩트입니다.
   • 그래디언트 반전 (Gradient Reversal): 특징 추출기 (Feature Extractor) 에서 추출된 특징이 세포 수를 예측하는 데 사용되지 않도록, 세포 수 회귀 헤드의 그래디언트를 반전시켜 역전파합니다. 이를 통해 모델이 세포 수에 과도하게 의존하지 않도록 유도하면서도, 교란 분류에 필요한 잔여 배치 구조는 유지합니다.

2. 테스트 시간 적응 (Test-Time Adaptation, TTA) 프레임워크:

   • SHOT (Source Hypothesis Transfer): 사전 훈련된 모델의 분류기 (Classifier) 는 고정하고, 특징 추출기 (Feature Extractor) 만을 테스트 데이터에 맞춰 업데이트합니다.
   • 손실 함수: 엔트로피 최소화 (Entropy Minimization), 다양성 정규화 (Diversity Regularization), 그리고 신뢰도 임계값을 기반으로 한 의사 레이블 (Pseudo-labeling) 분류 손실을 결합하여 사용합니다.
   • 적용 방식: 테스트 단계에서 레이블 없이 모델이 새로운 배치 데이터에 적응하도록 반복적으로 학습 (Multi-epoch adaptation) 시킵니다.

3. 모델 아키텍처:

   • 기본 백본으로 DenseNet-161을 사용하며, Cell Painting 이미지의 채널 수 (5 또는 6 채널) 에 맞게 초기화됩니다.
   • 훈련 시 CutMix 증강, 수평/수직 뒤집기, 90 도 회전을 적용합니다.

3. 주요 기여 (Key Contributions)

1. 생물학적 지향 TTA 프레임워크: 컴퓨터 비전의 TTA 기법을 Cell Painting 데이터에 적용하여, 보이지 않는 배치 및 세포 유형 변화 하에서도 강력한 일반화 성능을 달성했습니다.
2. 세포 수 적대적 훈련 (CCR): 세포 수라는 생물학적 특징을 명시적으로 대상으로 한 새로운 적대적 메커니즘을 도입하여, 기존 일반적인 배치 효과 제거 방법보다 우수한 성능을 입증했습니다.
3. 포괄적인 평가: 두 가지 대규모 Cell Painting 데이터셋 (RxRx1, JUMP-CP) 과 네 가지 서로 다른 세포 유형 (U2OS, HUVEC, RPE, HepG2) 에 걸쳐 일관된 성능 향상을 보였습니다.

4. 실험 결과 (Results)

A. 성능 향상 (Accuracy Improvement)

• RxRx1 데이터셋:
  • 기존 SOTA (AdaBN) 대비 4.5% 향상된 정확도 (91.6%) 를 기록했습니다.
  • 기존 1,139 개의 siRNA 교란 클래스를 분류하는 데 성공했습니다.
  • 특히 데이터가 부족하고 예측이 어려웠던 U2OS 세포 유형에서 정확도가 68.2% 에서 76.2% 로 크게 개선되었습니다 (+8.0%p).
• JUMP-CP 데이터셋 (CRISPR 교란):
  • 484 개의 CRISPR 교란 클래스에서 15.7% 의 정확도 향상을 보였습니다 (AdaBN 대비 SHOT-CCR 적용 시).
  • 이 데이터셋은 세포 유형이 단일 (U2OS) 이고 세포 수 분포가 균일하여 CCR 의 효과가 RxRx1 보다 작았으나, 여전히 TTA 기법 자체가 큰 개선을 가져왔습니다.

B. 통계적 유의성 및 분석

• 통계적 유의성: SHOT-CCR 의 성능 향상은 통계적으로 유의미했습니다 (t-test, p < 0.0001).
• 세포 수 분포의 영향: 세포 수 분포가 배치 간에 크게 다른 경우 (RxRx1) CCR 의 효과가 극대화되었으나, 분포가 균일한 경우 (JUMP-CP) 효과는 미미했습니다. 이는 CCR 이 세포 수 이질성이 큰 환경에서 특히 유용함을 시사합니다.
• 생물학적 신호 회복 (Gene Enrichment): 성능이 개선된 유전자들을 분석한 결과, 세포질, 헬리케이스 활성, 엔도막 시스템 등 생물학적으로 의미 있는 경로에서 유의미한 풍부화 (Enrichment) 가 확인되었습니다. 이는 모델이 단순한 아티팩트가 아닌 실제 생물학적 신호를 학습했음을 증명합니다.

C. 모델 백본 비교

• DenseNet-161, EfficientNetV2, ConvNeXt 를 비교한 결과, 아키텍처 변경보다는 TTA 기법과 CCR 의 적용이 성능 향상의 주된 요인이었습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 배치 효과 해결의 새로운 패러다임: 단순히 모든 배치 정보를 제거하는 것이 아니라, 생물학적으로 중요한 특징 (세포 수) 을 식별하고 이를 적절히 제어함으로써 모델이 교란 신호를 더 잘 학습하도록 돕는 전략을 제시했습니다.
• 데이터셋 구성에 대한 통찰: 실험 배치 간 세포 수 분포가 크게 다른 경우, 이를 훈련/테스트 세트에 어떻게 배분하느냐가 모델 성능에 결정적임을 보여주었습니다.
• 확장성: 제안된 방법은 RxRx1 과 JUMP-CP 두 가지 주요 데이터셋에서 모두 유효성이 입증되었으며, 향후 다른 대규모 Cell Painting 데이터셋 및 약물 발견 파이프라인에 적용 가능한 표준 벤치마크가 될 것으로 기대됩니다.
• 미래 방향: 트랜스포머 기반 모델 (Layer Normalization 사용) 로의 확장, 세포 수 외의 다른 교란 요인 (예: 플레이트의 행/열 위치) 에 대한 그래디언트 반전 파라미터 조정 등의 연구가 필요하다고 제안합니다.

이 논문은 생물학적 지식을 머신러닝의 적대적 훈련에 통합함으로써, 실험적 노이즈를 줄이고 약물 발견을 위한 AI 모델의 신뢰성과 일반화 능력을 획기적으로 높인 중요한 연구로 평가됩니다.