[논문 리뷰] Class-Guided Image-to-Image Diffusion: Cell Painting from Brightfield Images with Class Labels

Jan Cross-Zamirski, Praveen Anand외 4명|arXiv (Cornell University)|2023. 03. 15.

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

이 논문은 브라이트필드 현미경 입력에서 세포 페인팅 이미지를 생성하기 위해 이산형 교란 레이블을 활용하는 클래스 가이드드 이미지-투-이미지 확산 모델을 제안한다. 확산 기반 이미지-투-이미지 프레임워크에 클래스 레이블을 통합함으로써, 특히 생물학적으로 활성인 서브셋에서 훈련된 경우, 형태적 특징의 정확도와 후속 약물 작용 기전 예측 성능이 향상된다. 이는 특징 상관관계와 구조적 유사도 측면에서 가이드되지 않은 모델보다 뛰어나다.

ABSTRACT

Image-to-image reconstruction problems with free or inexpensive metadata in the form of class labels appear often in biological and medical image domains. Existing text-guided or style-transfer image-to-image approaches do not translate to datasets where additional information is provided as discrete classes. We introduce and implement a model which combines image-to-image and class-guided denoising diffusion probabilistic models. We train our model on a real-world dataset of microscopy images used for drug discovery, with and without incorporating metadata labels. By exploring the properties of image-to-image diffusion with relevant labels, we show that class-guided image-to-image diffusion can improve the meaningful content of the reconstructed images and outperform the unguided model in useful downstream tasks.

연구 동기 및 목표

• 약물 발견 분야에서 저비용의 브라이트필드 현미경 입력으로부터 고해상도 형광 세포 페인팅 이미지를 생성하는 데 도전하는 것.
• 이산형 클래스 레이블(예: 교란 유형)이 이미지-투-이미지 확산 모델을 통해 생물학적으로 의미 있는 재구성을 생성하는 데 도움이 될 수 있는지 조사하는 것.
• 표준 이미지 메트릭스 외에도 약물 작용 기전 예측 및 클러스터링과 같은 후속 생물학적 작업을 통해 모델 성능을 평가하는 것.
• 레이블 가이드드 확산이 생물학적으로 관련성이 있는 활성 화합물 서브셋에서 훈련된 경우, 가이드되지 않은 이미지-투-이미지 기반 모델보다 특징 허용도와 이미지 재구성 정확도에서 뛰어나다는 것을 입증하는 것.

제안 방법

• 이 방법은 클래스 가이드드 노이즈 제거를 통해 팔레트 이미지-투-이미지 확산 프레임워크를 확장하며, 이산형 교란 레이블에 조건을 부여하기 위해 분류기 가이던스 메커니즘을 사용한다.
• U-Net 기반의 노이즈 제거 네트워크를 사용하여 브라이트필드 입력 이미지와 클래스 레이블 임베딩을 동시에 입력으로 받아 해당하는 형광 채널 이미지를 예측한다.
• 모델은 JUMP-CP Target2 데이터셋의 일부에 대해 훈련되며, 입력으로 3채널 브라이트필드 이미지를, 목표로 5채널 세포 페인팅 이미지를 사용한다.
• 클래스 레이블은 학습된 프로젝션 헤드를 통해 임베딩되며, 분류기 가이던스를 통해 노이즈 제거 과정에서 사용되어 원하는 표현형 출력으로 방향을 조정한다.
• 레이블 품질과 데이터 희소성 영향을 평가하기 위해 전체 플레이트 및 활성 서브셋 훈련 루틴을 모두 지원한다.
• 이미지 메트릭스(FID, SSIM), CellProfiler와의 특징 상관관계, 약물 작용 기전 예측 작업에서의 타겟 매칭을 통해 성능을 평가한다.

Figure 2: Given input Brightfield (3 channels) our model is able to generate 5 Cell Painting channels. Incorporating meaningful labels can improve biological feature quality and performance on downstream tasks without significantly reducing image quality or adding background noise. Columns left to r

실험 결과

연구 질문

• RQ1 이산형 클래스 레이블은 현미경 이미지-투-이미지 확산 재구성의 품질과 생물학적 관련성에 향상 효과를 줄 수 있는가?
• RQ2 교란 레이블을 통합함으로써 약물 프로파일링 작업, 예를 들어 약물 작용 기전 예측 및 클러스터링 성능이 향상되는가?
• RQ3 훈련 데이터의 품질—특히 비활성 또는 노이즈가 많은 샘플의 포함 여부—는 클래스 가이드드 확산 모델의 성능에 어떤 영향을 미치는가?
• RQ4 클래스 가이드드 모델이 특징 허용도와 이미지 재구성 메트릭스 측면에서 가이드되지 않은 이미지-투-이미지 기반 모델보다 얼마나 뛰어나게 성능을 내는가?

주요 결과

• 활성 교란 서브셋에서 훈련된 클래스 가이드드 모델은 CellProfiler에서의 실제 형태적 특징과 가장 높은 상관관계를 기록했으며, 가이드되지 않은 팔레트 모델보다 뛰어났다.
• 작은 훈련 세트로도 구조적 유사도(SSI)와 약물 작용 기전 예측의 타겟 매칭 정확도가 향상되어 데이터 부족에 대한 강건성을 입증했다.
• 교란 레이블을 조건으로 사용함으로써 타겟 매칭 정확도가 가이드되지 않은 모델 대비 15% 향상되었으며, 이는 더 나은 표현형 일반화 능력을 의미한다.
• 전체 플레이트에서 교란 레이블을 사용해 훈련한 모델는 가장 낮은 FID 점수(14.2)를 기록했지만, 약간의 배경 노이즈가 발생하여 이미지 정확도와 생물학적 관련성 사이의 상충 관계를 보였다.
• 분류기 가이던스는 타겟 매칭을 향상시켰지만 이미지 및 특징 품질은 약간 저하시켜, 레이블 선택 시 신중함이 필요함을 시사했다.
• 결과적으로 높은 품질의 생물학적으로 정보적인 레이블이 필수적임을 보여주며, 열악하거나 정보가 없는 레이블은 이미지 품질을 떨어뜨리고 모델을 잘못 이끌 수 있음을 시사한다.

Figure 3: Images generated by models trained with the active subset. The labelled images were sampled with both AdaGN and CG. Cropped to $100\times 100$ pixels. Columns left to right: Brightfield (input), DNA, RNA, ER, Mito, AGP.