Calibrating Generative Models to Feature Distributions with MMD Finetuning

기술 요약: MMD 파인튜닝(kCGM)을 통한 생성 모델의 특징 분포 교정

문제 정의
생성 모델은 개별 샘플은 그럴듯하게 생성하지만, 주요 관심 특징(feature)에 대한 분포적 불일치(distributional misalignment)를 보이는 경우가 많습니다. 예를 들어, 광범위한 화학 공간에서 사전 학습된 모델은 특정 치료 계열(예: 항생제)의 특징 분포와 크게 다른 분자를 생성할 수 있습니다. 이를 직접적인 파인튜닝으로 교정하는 것은 문제가 많습니다. 이는 작은 데이터셋에 대한 과적합 위험이 있고, 전체 타겟 샘플(부족할 수 있음)이 필요하며, 어떤 특정 특징을 일치시켜야 하는지 명시하지 못합니다. 또한, 기존 방법들은 미분 불가능한 이산 값 또는 블랙박스 특징(예: Morgan fingerprint)을 다루는 데 어려움을 겪습니다. 본 논문은 미분 가능한 특징 맵이나 샘플링 절차를 요구하지 않으면서, 오직 특징 수준의 감독(supervision)만을 사용하여 생성 모델을 타겟 특징 분포에 맞게 교정할 수 있는 알고리즘의 공백을 식별했습니다.

방법론: kCGM
저자들은 생성된 특징 분포와 타겟 특징 분포 사이의 최대 평균 편차(Maximum Mean Discrepancy, MMD)를 최소화하면서, 사전 학습된 모델에 대한 쿨백-라이블러 발산(Kullback-Leibler divergence)으로 정규화하는 파인튜닝 방법인 **kernel Calibrating Generative Models (kCGM)**를 소개합니다.

• 목적 함수: 이 방법은 다음 두 항으로 구성된 목적 함수 $L_\theta$를 최적화합니다:

  1. MMD 손실: 모델 유도 특징 분포 $P$ (여기서 $y = h(x)$ 이고 $x \sim p_\theta$)와 경험적 타겟 분포 $Q$ (타겟 특징 $y^{targ}_{1:N}$으로부터 도출됨) 사이의 제곱 MMD에 대한 비편향 유한 샘플 추정치입니다.
  2. KL 정규화: 파인튜닝된 모델이 사전 학습된 분포와 가깝게 유지되도록 하여, 파괴적 망각(catastrophic forgetting)이나 모드 붕괴(mode collapse)를 방지하기 위한 항 $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$입니다.

• 커널 선택: 평균만을 일치시키는 내적 커널(dot-product kernels)에 의존할 수 있는 이전 방식들과 달리, kCGM은 확률 분포에 대한 메트릭을 정의하는 커널(동일한 분포일 때만 0이 되는 커널)을 사용합니다.

  • 연속적인 유클리드 특징의 경우, **에너지 거리 커널(energy-distance kernel)**을 사용합니다.
  • 이산적인 비트 벡터 특징(예: Morgan fingerprints)의 경우, **타니모토/자카드 커널(Tanimoto/Jaccard kernel)**을 채택합니다.
  • 이러한 선택은 추가적인 대역폭(bandwidth) 하이퍼파라미터 도입을 피하게 해줍니다.

• 그래디언트 추정: 특징 맵 $h$와 샘플링 과정이 미분 불가능할 수 있으므로, kCGM은 **스코어 함수 그래디언트(score-function gradients, REINFORCE 스타일)**를 사용합니다.

  • Leave-One-Out (LOO) 베이스라인: 핵심적인 기술적 기여는 MMD 손실을 위한 새로운 LOO 베이스라인의 유도입니다. 표준 LOO 베이스라인은 배치 내 다른 샘플들의 계수를 평균 내지만, MMD 계수는 샘플 간에 결합되어 있습니다. 저자들은 특정 샘플이 미니 배치에 나타나지 않았을 경우의 평균 계수를 근사하여, 편향을 도입하지 않으면서 분산을 줄이는 특수한 "MMD-LOO" 추정치를 유도했습니다.
  • 알고리즘: 과정은 현재 모델로부터 샘플링하고(그래디언트 추적 없이), 특징 벡터를 계산하며, LOO 베이스라인을 사용하여 MMD 및 KL 계수를 계산한 뒤, 스코어 함수 그래디언트를 통해 모델 가중치를 업데이트하는 것을 포함합니다.

주요 기여

1. 정식화: 본 논문은 특징 분포 매칭을 오직 타겟 특징 샘플만을 사용하는 KL 정규화된 MMD 파인튜닝으로 정식화하여, 원래의 샘플 공간에서의 전체 타겟 샘플이 필요하지 않도록 했습니다.
2. 비편향 추정기: 저자들은 미분 불가능한 특징 맵과 재매개변수화가 불가능한 샘플러(예: 자기회귀 모델, 이산 확산 모델)를 사용할 수 있도록 MMD 목적 함수에 대한 비편향 스코어 함수 그래디언트 추정기를 유도했습니다.
3. 일반화: kCGM은 특징의 평균만을 일치시키는 CGM을 일반화합니다. kCGM은 MMD를 통해 전체 분포를 일치시킵니다. 논문은 kCGM에서 내적 커널을 사용하면 CGM-relax 목적 함수가 복구됨을 언급합니다.

실험 결과
저자들은 세 가지 도메인(소분자 생성, 단백질 구조 생성, 조절 DNA 생성)에서 kCGM을 평가했습니다.

• 소분자 항생제:

  • 태스크: G2PT 자기회귀 그래프 모델을 교정하여 174개의 알려진 항생제의 특징 분포와 일치시키기.
  • 특징: Morgan fingerprints (국소 부분 구조), Murcko scaffolds (위상), 그리고 물리화학적 기술자(physicochemical descriptors).
  • 결과: kCGM은 직접적인 파인튜닝 및 CGM과 비교하여 타겟 특징 분포 일치(MMD로 측정)와 사전 학습 모델과의 거리 사이에서 우수한 트레이드오프를 달emons 달성했습니다. 결정적으로, 소규모 항생제 세트에 대한 직접 파인튜닝은 높은 비율의 화학적 부적격 분자를 초래한 반면, kCGం은 특징 일치를 개선하면서 동시에 화학적 유효성(validity)을 증가시켰습니다. 또한 kCGM은 핑거프린트만으로 학습했을 때도 스캐폴드 및 기술자 매칭을 개선하는 교차 특징 전이(cross-feature transfer) 능력을 보여주었습니다.

• 단백질 구조 다양성:

  • 태스크: Genie 2(단백질 확산 모델)를 CATH 데이터셋의 이차 구조 분포(알파 헬릭스 대 베타 스트랜드)에 맞게 파인튜닝하기.
  • 결과: 사전 학습된 모델은 주로 알파 헬릭스를 생성했습니다. kCGM은 CGM에 비해 타겟 분포와의 불일치를 약 2배 감소시켰습니다. kCGM은 KL 가중치 $\lambda$에 대해 매끄러운 트레이드오프 곡선을 보여준 반면, CGM은 $\lambda$에 대한 민감도가 낮아 최적화가 어려웠습니다.

• 조절 DNA 생성:

  • 태스크: 세포 유형 특이적 활성 프로파일(AlphaGenome에 의해 예측됨)을 가진 현실적인 DNA 서열을 생성하도록 이산 확산 모델을 파인튜닝하기.
  • 결과: kCGM은 47개의 서로 다른 cis-regulatory 토픽에 대한 활성 프로파일의 전체 분포를 성공적으로 일치시켰으며, 평균만을 일치시키는 CGM보다 뛰어난 성능을 보였습니다. 파인튜닝된 모델은 실제 인핸서(enhancer)와 유사한 활성 프로파일을 가진 서열을 생성한 반면, 사전 학습된 모델은 활성이 낮은 서열을 생성했습니다.

의의 및 주장
본 논문은 kCGM이 특히 타겟 데이터셋이 작거나 특정 분포의 일부 측면만을 보존해야 할 때, 직접적인 파인튜닝에 대한 강력한 대안을 제공한다고 주장합니다. 주요 의의는 다음과 같습니다:

• 특징 수준의 감독만을 사용하여 생성 모델을 사용자 지정 특징 분포에 정렬할 수 있습니다.
• 과학 분야에서 흔히 사용되지만 표준 그래디언트 기반 파인튜닝과는 호환되지 않는 블랙박스 및 미분 불가능한 특징(예: 핑거프린트, 예측기 출력)을 처리할 수 있습니다.
• 직접적인 파인튜닝의 흔한 실패 모드인 샘플 품질(유효성)을 유지하거나 오히려 개선하면서 분포적 미교정을 수정할 수 있습니다.
• 자기회귀 모델, 연속 공간 확산 모델, 이산 확산 모델을 포함한 다양한 아키텍처에 적용 가능합니다.

저자들은 MMD 가설 검정의 차원의 저주로 인해 매우 고차원적인 특징에 대해서는 효과가 떨어질 수 있다는 점과, 트랙터블한 샘플링 궤적 로그 확률(log-probabilities)이 필요하다는 점(flow matching 모델 제외) 등의 한계를 언급했습니다. 향후 연구에서는 원래의 생성 공간에서의 전체 샘플 없이 오직 분포적 제약만이 존재하는 설정으로 이 프레임워크를 확장할 수 있을 것이라고 제안했습니다.