Calibrating Generative Models to Distributional Constraints

Motivation

• Known: 생성 모델의 miscalibration(예: mode collapse, 사회적 편향 재생산, 단백질/DNA 서열 통계의 부정확성)은 여러 도메인에서 널리 관찰되는 현상이며, supervised learning 분야에서는 Platt scaling, conformal prediction 등 post-hoc 보정 기법이 잘 알려져 있다. 또한 개별 샘플 단위의 reward 기반 preference fine-tuning(RLHF, DPO 등) 방법들도 널리 사용되고 있다.
• Gap: 기존 보정 방법들은 예측 모델 출력을 개별적으로 보정하거나 개별 샘플에 대한 reward를 부여하는 데 그쳐, 분포 자체의 통계량(distributional constraint)을 직접 제약하지 못한다. 분포 수준 제약을 다루는 소수의 선행/동시 연구(Khalifa et al., 2021; Shen et al., 2024; Khalafi et al., 2026)는 단일 모델 클래스에만 적용되고 저차원 제약에서만 검증되어, 고차원·다중 제약·다양한 모델군에 걸친 일반적 해법이 부재하다.
• Why: 생성 모델의 통계적 편향을 원칙적인 KL 최소화 관점에서 교정하는 일반적 프레임워크는 이미지 생성, 언어 모델링, 단백질/DNA 설계 등 다양한 응용에서 안전성·공정성·과학적 타당성을 확보하는 데 중요하며, likelihood를 계산 가능한 어떤 생성 모델 클래스에도 적용 가능하다는 점에서 실용적 파급력이 크다.
• Approach: 저자들은 calibration을 base model과의 KL divergence를 최소화하면서 기댓값 제약 Epθ[h(x)]=h*를 만족시키는 constrained optimization 문제로 정식화하고, 제약을 정확히 부과하는 것이 어려운 문제를 완화하기 위해 relax loss(제약 위반에 대한 penalty)와 reward loss(maximum entropy 해에 대한 reward fine-tuning)라는 두 surrogate objective를 제안한다.

Achievement

Figure 4. A: Samples from the Genie2 before and after calibration with CGM-relax (λ=10−3). B: CGM-relax reduces the dist

1. CGM-relax: 제약 위반 항 Lviol과 KL 항 LKL의 가중합으로 정의되는 relax loss를 제안하고, importance weight와 leave-one-out(LOO) baseline을 활용해 unbiased gradient 추정치를 계산하는 stochastic optimization 알고리즘(Algorithm 1)을 개발했다.
2. CGM-reward: maximum entropy 문제와의 연결을 통해 이론적으로 정당화된 reward loss를 제안했다(Theorem 2.1). base model로부터의 샘플만으로 dual parameter α*를 추정하고, 이를 reward로 사용해 fine-tuning하는 알고리즘(Algorithm 2)을 개발했다.
3. 광범위한 실증 검증: 최대 90억(nine billion) 파라미터 모델, 수백 개의 동시 제약 조건에 걸쳐 diffusion model(1D Gaussian mixture, 이미지), protein structure model(Genie2), 조건부 TarFlow, 언어 모델 등 다양한 도메인에서 miscalibration을 크게 감소시킴을 보였다.

How

Figure 3. A: CGM successfully reweights the probability of a rare mode in a 1D Gaussian mixture. B: CGM-relax calibrates

• Calibration 문제를 pθ=argmin DKL(pθ‖pθbase) s.t. Epθ[h(x)]=h로 정식화(Equation 1).
• Relax loss: Lrelax(θ)=‖Epθ[h(x)]−h*‖²+λDKL(pθ‖pθbase)로 정의하고, M개 샘플로부터 bias-corrected unbiased estimator bLviol, bLKL을 구성하여 stochastic gradient descent로 최적화.
• Reward loss: 비모수적 maximum entropy 문제(Equation 4)의 해 pα(x)∝pθbase(x)exp{αᵀh(x)}가 존재함을 증명(Theorem 2.1)하고, Wainwright & Jordan의 dual 결과를 이용해 base model 샘플만으로 α*를 추정(Equation 6).
• 두 알고리즘 모두 importance weight wm=pθ(xm)/pstop-grad(θ)(xm)와 leave-one-out(LOO) baseline을 사용해 KL loss와 (제약 위반 혹은 reward) loss의 unbiased gradient 추정치를 계산.
• protein design(Genie2), image generation(diffusion model, TarFlow), language modeling 등 세 가지 응용 도메인에서 실험을 수행하여 calibration error 감소를 정량적으로 평가.

Originality

• Calibration을 KL divergence 기반 constrained optimization으로 통일적으로 정식화하여, 기존에 개별 도메인·모델 클래스에 국한되던 분포 수준 보정 방법들을 일반화함.
• Maximum entropy 문제와의 연결을 통해 reward fine-tuning과 calibration을 이론적으로 연결하는 새로운 관점(Theorem 2.1) 제시.
• 비미분 가능한 제약 함수 h(·)도 다룰 수 있는 실용적인 unbiased gradient estimator(LOO baseline 포함)를 설계하여, likelihood를 계산할 수 있는 임의의 생성 모델 클래스에 폭넓게 적용 가능하도록 함.

Limitation & Further Study

• Relax loss는 하이퍼파라미터 λ 선택에 민감하며, λ가 유한한 값으로 고정되어야 하기 때문에 정확한 제약 만족과 base model과의 근접성 사이의 trade-off가 발생하는데, 이에 대한 heuristic만 제시되고 엄밀한 자동 선택 방법은 부재하다.
• 이론적 보장(Theorem 2.1)은 h*가 attainable moment set의 relative interior에 존재하고 제약 간 중복이 없다는 가정에 의존하는데, 실제 응용에서 이러한 가정의 검증이 어려울 수 있다.
• Monte Carlo 기반 unbiased estimator는 유한 샘플 M에서 분산이 클 수 있으며, 특히 rare mode나 고차원 제약에서 학습 안정성 문제가 제기될 수 있어 후속 연구에서 분산 감소 기법이나 다른 divergence 척도로의 확장이 필요해 보인다.
• 다양한 모델 클래스(diffusion, autoregressive, flow 기반)에 적용했으나, likelihood 계산이 어려운 일부 생성 모델(예: 일부 GAN 변형)에는 직접 적용이 제한적일 수 있다.

Evaluation

총평: 생성 모델의 분포 수준 보정 문제를 이론적으로 명확히 정식화하고, 다양한 모델 클래스와 대규모 응용에 걸쳐 실용적이고 확장 가능한 두 가지 fine-tuning 알고리즘을 제시한 견고한 연구이다. 하이퍼파라미터 선택과 이론적 가정의 실제 검증 가능성에 대한 추가적 논의가 있다면 더욱 완성도가 높아질 것이다.