CompDiff: Hierarchical Compositional Diffusion for Fair and Zero-Shot Intersectional Medical Image Generation

Each language version is independently generated for its own context, not a direct translation.

🏥 문제: "편향된 요리사"와 "없던 재료"

지금까지의 인공지능 (생성 모델) 은 마치 대규모 식당의 요리사와 같았습니다.

상황: 이 요리사는 매일 수천 명의 '젊은 백인 남성' 환자 사진을 보고 배웠습니다. 하지만 '80 대 이상 아시아계 여성' 환자에 대한 사진은 거의 없거나 아예 없습니다.
문제점: 요리사가 '젊은 백인 남성' 요리는 완벽하게 만들지만, '80 대 아시아계 여성'을 주문하면 요리를 엉망으로 만들거나 아예 못 만듭니다.
기존 해결책의 한계: 기존 연구자들은 "요리사에게 '젊은 남성' 요리를 덜 만들고 '여성' 요리를 더 많이 만들어라"라고 **지시 (손실 함수 재가중)**만 했습니다. 하지만 **아예 존재하지 않는 재료 (데이터)**로 요리를 하라고 지시하는 것은 소용이 없습니다. 요리사는 그 재료를 본 적이 없기 때문입니다.

💡 해결책: CompDiff (컴포지티브 디퓨전)

이 논문은 "지시만 바꾸는 게 아니라, **요리사의 사고방식 (구조)**을 바꾸자"고 제안합니다. 바로 CompDiff입니다.

1. 레고 블록 비유 (조립적 사고)

기존 방식은 "80 대 아시아계 여성"이라는 완전한 인형 하나하나를 통째로 외우게 했습니다. 하지만 CompDiff 는 레고 블록처럼 접근합니다.

블록 1: '나이' (80 대)
블록 2: '인종' (아시아계)
블록 3: '성별' (여성)

이 세 가지 블록을 따로따로 잘 배워두면, **처음 보는 조합 (80 대 + 아시아계 + 여성)**이 들어와도 이 블록들을 잘게 조립해서 새로운 인형을 만들 수 있습니다. 마치 언어 모델이 '사과', '빨간', '맛있다'는 단어를 알고 있으면, 본 적 없는 '빨간 사과가 맛있다'는 문장을 만들 수 있는 것과 같습니다.

2. 전담 관리자 (HCN)

CompDiff 는 **HCN(계층적 컨디셔너 네트워크)**이라는 전담 관리자를 채용합니다.

이 관리자는 나이, 성별, 인종 정보를 따로따로 분석하고, 이 정보들이 서로 어떻게 섞여야 하는지 (예: 나이가 많을수록 인종에 따른 특징이 어떻게 변하는지) 계산합니다.
이렇게 계산된 정보를 레고 블록처럼 조립하여, AI 가 그림을 그릴 때 참고할 수 있게 해줍니다.

🎨 결과: 왜 이것이 중요한가요?

이 방식을 적용한 실험 결과는 매우 훌륭했습니다.

질 좋은 그림: 희귀한 그룹 (예: 고령의 아시아계 여성) 에 대해서도 그림의 질이 훨씬 좋아졌습니다. (기존 방식은 엉망이었지만, CompDiff 는 잘 그렸습니다.)
공정한 결과: 어떤 그룹이든 그림의 퀄리티 차이가 거의 없었습니다. 즉, 특정 인종이나 성별을 차별하지 않는 공정한 AI 가 되었습니다.
보이지 않는 것도 그릴 수 있음: 훈련 데이터에 아예 없던 조합 (Zero-shot) 을 주문해도, 레고 블록을 조립하듯 자연스럽게 그릴 수 있었습니다.
실제 의료 활용: 이 AI 가 만든 그림으로 훈련된 진단 프로그램은 실제 환자 사진을 볼 때 더 정확하게 질병을 찾아냈고, 특정 그룹을 놓치는 실수 (진단 편향) 도 줄었습니다.

📝 한 줄 요약

"기존 AI 는 데이터가 없는 그룹은 못 그렸지만, CompDiff 는 나이, 성별, 인종이라는 '레고 블록'을 따로 배워 조합하는 법을 가르쳐 줌으로써, 어떤 환자 그룹이든 공평하고 정확하게 그림을 그릴 수 있게 만들었습니다."

이 기술은 의료 AI 가 특정 인종이나 성별에 치우치지 않고, 모든 환자에게 공평하게 도움을 줄 수 있는 미래를 여는 중요한 첫걸음입니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem Statement)

의료 영상 생성을 위한 생성 모델 (Diffusion Models) 은 데이터 부족 문제를 해결하고 공정한 AI 를 훈련시키기 위해 널리 사용되고 있습니다. 그러나 기존 연구들은 다음과 같은 근본적인 가정을 간과하고 있습니다: "생성 모델이 모든 인구통계학적 그룹 (인종, 성별, 나이 등) 에 대해 균일하게 고품질의 이미지를 생성하는가?"

불균형 생성기 문제 (Imbalanced Generator Problem): 훈련 데이터가 불균형할 경우, 생성 모델은 평균적인 품질은 높게 유지하더라도 드문 하위 그룹 (Rare Subgroups) 이나 훈련 데이터에 존재하지 않는 인구통계학적 교차점 (Intersectional Groups, 예: '80 대 이상 + 아시아계 + 여성') 에 대해서는 품질이 현저히 저하되거나 생성이 불가능합니다.
기존 방법의 한계:
- 손실 가중치 재조정 (Loss Reweighting, 예: FairDiffusion): 최적화 수준에서 작동하며, 훈련 데이터에 아예 존재하지 않는 조합에는 학습 신호를 제공할 수 없습니다.
- 텍스트 프롬프트 의존성: 기존 모델은 CLIP 임베딩 내에서 인구통계학적 토큰과 임상적 토큰이 제한된 토큰 수 (77 개) 를 공유하며 경쟁하게 되어, 희귀한 조합에 대한 표현력이 부족합니다.

2. 제안 방법론 (Methodology: CompDiff)

저자들은 인구통계학적 정체성이 구성적 (Compositional) 이라는 통찰을 바탕으로, 최적화 수준이 아닌 표현 (Representation) 수준에서 문제를 해결하는 CompDiff를 제안합니다.

핵심 구성 요소: 계층적 조건부 네트워크 (Hierarchical Conditioner Network, HCN)

기존의 단일 텍스트 프롬프트 방식 대신, 인구통계학적 속성을 별도로 처리하여 구조화된 토큰을 생성합니다.

단일 속성 임베딩 (Single-Attribute Embeddings): 나이, 성별, 인종 등 각 속성을 공유 잠재 공간에 매핑합니다.
쌍별 상호작용 (Pairwise Interactions): 속성 간의 비가산적 (Non-additive) 관계를 포착하기 위해 모든 쌍 (나이 - 성별, 나이 - 인종, 성별 - 인종) 에 대한 전용 MLP 를 사용하여 상호작용을 모델링합니다.
완전한 구성 (Full Composition): 쌍별 상호작용을 결합하여 최종 인구통계학적 표현 ( $h_{demo}$ ) 을 생성합니다.
토큰 생성 및 연결:
- $h_{demo}$ 를 가우시안 잠재 변수로 매핑한 후, Cross-Attention 컨텍스트 차원으로 투영하여 인구통계학적 토큰 ( $c$ ) 을 생성합니다.
- 이 토큰 $c$ 를 CLIP 의 임상 텍스트 임베딩 ( $E_{text}$ ) 에 연결하여 UNet 에 입력합니다.

학습 목표 (Training Objective)

모델은 다음 총 손실 함수로 학습됩니다:
$L = L_{diff} + \lambda_{comp}L_{comp} + \lambda_{aux}L_{aux} + \lambda_{KL}L_{KL}$

$L_{diff}$ : 표준 확산 모델 손실.
$L_{comp}$ (구성 일관성): 복잡한 상호작용이 단순한 가산적 합과 너무 멀어지지 않도록 하는 소프트 앵커 (Soft Anchor) 역할.
$L_{aux}$ (보조 분류 손실): UNet 이 실제로 받는 토큰 ( $c$ ) 에서 인구통계학적 정보 (나이, 성별, 인종) 가 유지되도록 분류기를 통해 감시합니다. (임의의 잠재 변수 $\mu$ 가 아닌 투영된 토큰에 적용하는 것이 중요함).
$L_{KL}$ : 변분 추정을 위한 정규화.

3. 주요 기여 (Key Contributions)

표현 수준의 해결책: 데이터 불균형 문제를 최적화 가중치 조절이 아닌, 인구통계학적 표현의 구조적 계층화 (Hierarchical Factorization) 를 통해 해결합니다.
Zero-Shot 교차 일반화: 훈련 데이터에 존재하지 않는 드문 인구통계학적 조합 (예: 훈련 데이터에 없는 특정 인종 - 나이 - 성별 조합) 에 대해서도, 학습된 단일 속성 및 쌍별 임베딩을 구성 (Compose) 하여 고품질 이미지를 생성할 수 있습니다.
HCN 아키텍처: 인구통계학적 속성을 분리하여 계층적으로 처리하고, 이를 CLIP 임베딩과 결합하는 전용 네트워크를 설계하여 파라미터 공유와 데이터 효율성을 극대화했습니다.
하류 작업 (Downstream) 공정성 향상: 생성된 데이터로 훈련된 질병 분류기의 성능과 공정성 (Demographic Bias 감소) 을 동시에 개선함을 입증했습니다.

4. 실험 결과 (Results)

데이터셋: 흉부 X-ray (MIMIC-CXR) 및 안저 영상 (FairGenMed).
비교 대상: 표준 파인튜닝 (Baseline), FairDiffusion.

주요 성과 지표:

이미지 품질 (FID): CompDiff 는 두 모달리티 모두에서 가장 낮은 FID 를 기록했습니다.
- 흉부 X-ray: 64.3 (Baseline: 82.8, FairDiffusion: 75.1).
- 안저 영상: 54.6 (Baseline: 72.2, FairDiffusion: 64.3).
공정성 (ES-FID): 인구통계학적 하위 그룹 간의 품질 편차를 측정하는 ES-FID 에서 CompDiff 가 가장 낮은 값을 보여 그룹 간 편차가 가장 적음을 입증했습니다.
Zero-Shot 교차 일반화: 훈련 데이터에서 완전히 제거된 5 개의 드문 교차 그룹에 대해 평가했을 때, CompDiff 는 최대 21% 의 FID 개선을 보였습니다. 반면 FairDiffusion 은 훈련 데이터가 없는 경우 오히려 성능이 저하되기도 했습니다.
하류 분류기 성능: CompDiff 로 생성된 데이터로 훈련된 분류기는 실제 데이터에서 더 높은 AUROC 를 보였으며, 인종/성별/연령에 따른 진단 편차 (Underdiagnosis rate, Equalized Odds Difference) 가 감소했습니다.

Ablation Study (성분 분석):

계층적 구조의 중요성: 단순 MLP(Flat) 를 사용한 경우 인구통계학적 제어력이 급격히 떨어졌으나, HCN 을 사용한 경우 높은 정확도를 유지했습니다.
보조 손실의 위치: 보조 분류 손실을 투영 전 ( $\mu$ ) 이 아닌 투영 후 토큰 ( $c$ ) 에 적용해야 UNet 이 올바른 정보를 학습함을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 의료 영상 생성 분야에서 인구통계학적 편향을 해결하기 위해 단순히 데이터 재가중치를 하는 것을 넘어, 모델 아키텍처 자체의 설계 (Representation Design) 가 얼마나 중요한지를 강조합니다.

구성적 일반화: 언어 모델이 알려진 단어로 새로운 문장을 만드는 것처럼, CompDiff 는 알려진 인구통계학적 속성을 조합하여 훈련되지 않은 드문 그룹에 대한 이미지를 생성할 수 있습니다.
실용적 가치: 생성된 데이터가 실제 임상 진단 AI 의 공정성과 성능을 향상시키는 데 직접적으로 기여함을 입증했습니다.
한계점: fairness 평가가 정량적 지표에 의존하며, 연속적이거나 비구조화된 속성에는 적용이 제한적일 수 있습니다. 또한 표현 수준의 해결책이 데이터 불균형의 영향을 완전히 제거하지는 못합니다.

결론적으로, CompDiff 는 공정하고 제로샷 (Zero-Shot) 교차 일반화가 가능한 의료 영상 생성을 위한 새로운 패러다임을 제시하며, 향후 공정한 의료 AI 개발을 위한 중요한 방향성을 제시합니다.