기술적 요약: 분포 제약 조건에 대한 생성 모델 보정

문제 정의

생성 모델은 종종 보정 오류를 겪습니다. 이는 샘플링 분포의 통계량(예: 클래스 빈도, 이차 구조 비율, 또는 성별 표현)이 원하는 값에서 체계적으로 벗어날 때 발생합니다. 이 문제는 다양한 분야에서 나타납니다: 이미지 모델은 모드 붕괴를 보이고, 언어 모델은 사회적 편향을 강화하며, 단백질/DNA 모델은 비정형적인 하부 구조 빈도를 가진 구조를 생성합니다. 이러한 오류는 종종 데이터셋 불균형, 비최적의 훈련 역학, 또는 저온 샘플링과 같은 사후 조정에서 비롯됩니다.

저자들은 보정을 제약 최적화 문제로 규정합니다. 기본 생성 모델 $p_{\theta_{base}}(x)$, 통계량 $h(x)$, 그리고 목표 기대값 $h^*$가 주어졌을 때, 목표는 $E_{p_\theta}[h(x)] = h^*$를 만족하면서 기본 모델에 대한 쿨백-라이블러 (KL) 발산을 최소화하는 새로운 분포 $p_\theta$를 찾는 것입니다:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
이 공식화는 보정 제약 조건을 만족하는 기본 분포에 가장 가까운 모델을 찾습니다. 그러나 이 문제를 정확하게 해결하는 것은 비실용적입니다. 왜냐하면 목적 함수와 제약 조건 모두 고차원 생성 모델의 경우 계산하기 어려운 기대값을 포함하기 때문입니다.

방법론

비실용성을 해결하기 위해, 이 논문은 CGM-relax와 CGM-reward라고 명명된 미세 조정 (fine-tuning) 을 위한 두 가지 대리 목적 함수를 소개합니다. 두 방법 모두 샘플에서 유도된 편향 없는 기울기 추정을 사용하는 확률적 최적화에 의존합니다.

1. CGM-relax (Relax Loss)

이 접근법은 하드 제약 조건을 페널티 항으로 대체합니다. 목적 함수는 다음과 같습니다:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• 제약 위반 ($L_{viol}$): 경험적 통계량과 목표 값 사이의 차이의 제곱 $\ell_2$ 노름입니다. 저자들은 편향 수정된 표본 분산 공식을 사용하여 이 항에 대한 편향 없는 추정량을 유도합니다.
• KL 발산 ($L_{KL}$): 몬테카를로 샘플링을 통해 측정됩니다.
• 최적화: 이 방법은 중요도 샘플링을 사용하여 편향 없는 기울기를 계산합니다. 구체적으로, 현재 모델 $p_\theta$(stop-gradient 적용) 에서 샘플을 추출하고 샘플을 재가중하여 KL 항과 제약 위반의 기울기를 추정합니다. 하이퍼파라미터 $\lambda$는 제약 조건 충족과 기본 모델 근접성 사이의 균형을 조절합니다.

2. CGM-reward (Reward Loss)

이 접근법은 보정 문제를 최대 엔트로피 원리와 연결합니다. 최대 엔트로피 문제 (모멘트 제약 조건 하에 기본 모델에 대한 KL 을 최소화) 의 해는 기본 모델의 지수 기울기 (exponential tilt) 입니다: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• $\alpha^*$ 추정: 진정한 최대 엔트로피 분포는 비실용적이므로, 이 방법은 먼저 기본 모델에서 $N$개의 샘플을 사용하여 이중 최적화 문제를 풀어서 최적의 보상 파라미터 $\alpha^*$를 추정합니다.
• 손실 함수: 미세 조정 목적 함수는 현재 모델 $p_\theta$와 추정된 최대 엔트로피 분포 $p_{\hat{\alpha}_N}$ 사이의 KL 발산입니다:
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  이는 "보상"이 제약 조건에서 유도된 보상 미세 조정 문제와 동일합니다.
• 최적화: CGM-relax 와 유사하게, 분산을 줄이기 위해 leave-one-out 기준선을 사용하는 스코어 함수 기울기 추정기를 사용합니다.

기울기 추정

두 방법 모두 중요도 샘플링과 결합된 스코어 함수 기울기 추정기 (REINFORCE 스타일) 를 사용합니다. 기대값이 $\theta$에 의존하는 문제를 처리하기 위해, 저자들은 현재 모델의 고정된 버전에서 샘플을 추출하고 밀도 비율 (스코어 함수) 을 통해 기울기를 계산하는 "stop-gradient" 트릭을 사용합니다. 분산 감소는 leave-one-out 기준선을 사용하여 달성됩니다.

주요 기여

1. 공식화: 생성 모델 보정을 제약 KL 최소화 문제로 규정함.
2. 알고리즘: 확률적 최적화를 사용하여 제약 문제의 해를 근사하는 두 가지 실용적인 미세 조정 알고리즘 (CGM-relax 및 CGM-reward) 제안.
3. 확장성: 이러한 방법이 수백 개의 동시 제약 조건과 최대 90 억 개의 파라미터를 가진 모델을 처리할 수 있음을 입증.
4. 이론적 연결: 보정, 최대 엔트로피, 그리고 보상 미세 조정 간의 연결을 확립하여 CGM-reward 가 최대 엔트로피 해를 목표로 하고 CGM-relax 가 정규화된 버전을 목표로 함을 보여줌.

결과

저자들은 세 가지 다양한 응용 분야에서 이 방법들을 평가했습니다:

1. 단백질 설계 (Genie2 및 ESM3):

   • 목표: 이차 구조 함량 (알파-나선 및 베타-스트랜드) 을 자연 단백질 통계 (CATH 데이터셋) 와 일치하도록 보정.
   • 제약 조건: 최대 99 개의 동시 제약 조건 (이변량 CDF).
   • 결과: CGM-relax 는 Genie2 에 대해 자연 단백질 분포와의 거리를 4 배 이상, ESM3 에 대해 2 배 이상 감소시켰습니다. CGM-reward 는 고차원 제약 조건으로 $\alpha^*$를 추정하는 어려움으로 인해 다소 완만한 개선을 보였습니다. 두 방법 모두 생체 물리학적 타당성을 유지했지만, "설계 실패"가 약간 증가했습니다.

2. 이미지 생성 (TarFlow):

   • 목표: AFHQ 데이터셋의 "야생동물" 카테고리 내 클래스 비율 균형 (예: 사자, 호랑이, 여우, 늑대의 균등화).
   • 제약 조건: 5 개의 동시 제약 조건.
   • 결과: CGM-relax 는 클래스 비율의 총변동 거리를 0.306 에서 0.101 로 감소시켰습니다. 샘플의 현실성 (FID) 이 약간 감소했지만, 보정된 모델은 더 다양하고 균형 잡힌 생성물을 생산했습니다. CGM-reward 는 제약 위반을 유의미하게 감소시키지 못했습니다.

3. 언어 모델링 (Gemma-2-9B-IT):

   • 목표: 다양한 직업에 대한 생성된 어린이 이야기의 성별 편향 감소 (예: 남성/여성 변호사의 동등한 확률 보장).
   • 제약 조건: 8 개의 직업 조건부 제약 조건.
   • 결과: 두 방법 모두 성별 불균형을 감소시켰습니다. CGM-relax 는 기준선 대비 3 배 이상 불균형을 감소시킨 반면, CGM-reward 는 통계적으로 유의미하지만 더 작은 개선을 보였습니다. 미세 조정된 모델은 이야기 품질을 유지했습니다 (로그 확률 저하 없음).

합성 실험:

• 희귀 사건: CGM 은 기본 확률이 $10^{-3}$까지 낮을 때도 가우시안 혼합 모델에서 희귀 모드를 성공적으로 상향 가중했습니다.
• 고차원: CGM-relax 는 최대 $10^3$개의 제약 조건을 효과적으로 처리했습니다. 반면, CGM-reward 는 경험적 최대 엔트로피 문제 (식 7) 를 해결하는 어려움으로 인해 30 개 이상의 제약 조건에서는 비실용적이 되었습니다.

중요성 및 주장

이 논문은 CGM 이 모델 아키텍처 변경이나 처음부터의 재훈련 없이 분포 제약 조건을 충족하도록 고차원 생성 모델을 보정하는 실용적인 프레임워크를 제공한다고 주장합니다.

• 강건성: CGM-relax 는 최대 엔트로피 파라미터 추정이 비실용적이 되는 고차원 제약 조건 설정에서 특히 CGM-reward 보다 더 강건한 것으로 나타났습니다.
• 트레이드오프: 이 방법들은 보정 정확도와 기본 모델에 대한 충실도 (하이퍼파라미터 $\lambda$로 제어) 사이의 탐색 가능한 트레이드오프를 제공합니다.
• 한계: 저자들은 잔류 오류가 남아있으며, 특히 매우 희귀한 사건의 경우 그렇다고 인정합니다. 또한, 이 프레임워크는 현재 확률 밀도가 계산 가능한 모델 (확산 모델, 정규화 흐름, 자기회귀 모델) 에 의존하므로 GAN 과 같은 암시적 모델에는 직접 적용이 제한됩니다.

이 연구는 제약 조건 자체를 분포 정렬을 위한 구동 신호로 사용하여 보정을 표준 보상 미세 조정과 구별되는 독특하고 유익한 연구 방향으로 위치시킵니다.