기술 요약: 그래디언트 정보 기반 로짓 교정을 통한 이산 확산 모델의 플러그 앤 플레이 가이던스

문제 정의

이산 확산 모델(Discrete diffusion models)은 언어 생성, 생물학적 서열 합성, 분자 설계와 같은 영역에서 상당한 잠재력을 보여주었습니다. 그러나 단백질 안정성이나 유전자 발현과 같은 특정 목적 함수에 최적화된 샘플을 생성하는 '제어 가능한 생성(controllable generation)'을 달ert하는 것은 여전히 어려운 과제입니다. 기존 방식은 대개 분류기/분류기 없는 가이던스(classifier/classifier-free guidance) 또는 **광범위한 모델 미세 조정(fine-tuning)**에 의존합니다. 이러한 방법들은 다음과 같은 근본적인 한계를 가집니다:

• 계산 오버헤드: 미세 조정은 추가적인 학습 단계가 필요하며, 새로운 목적 함수가 생길 때마다 생성 모델을 다시 학습시켜야 합니다.
• 일반화 문제: 별도의 분류기나 시간 의존적 네트워크를 학습시키는 것은 각 타겟 조건에 대한 새로운 데이터 수집을 필요로 합니다.
• 이산 공간의 불안정성: 연속 공간 가이던스 방식(예: 그래디언트 기반 샘플링)을 이산 공간에 직접 적용하는 것은 범주형 데이터의 비미분성과 고차원 야코비안(Jacobian)의 수치적 불안정성으로 인해 방해를 받습니다.
• 비효율성: 순차적 몬테카를로(SMC) 또는 중요도 샘플링(importance sampling)에 기반한 대안적인 학습 불필요(training-free) 방식들은 종종 지나치게 높은 계산 비용을 초래하며 조향(steering) 효과도 제한적입니다.

본 논문은 모델 파라미터를 수정하지 않고도 기성 보상 함수(off-the-shelf reward functions)를 사용하여 이산 확산 모델을 효율적으로 가이드할 수 있는 플러그 앤 플레이, 학습 불필요 프레임워크의 필요성을 다룹니다.

방법론: 그래디언트 정보 기반 로짓 교정 (GILC)

저자들은 GILC를 제안하며, 이는 가이던스 작업을 기대 미래 보상을 나타내는 가치 함수(value function)의 그래디언트를 추정하는 문제로 재정의합니다. 이 방법은 세 가지 핵심 구성 요소를 통해 작동합니다.

1. 가치 추정을 위한 변분 프록시 (Variational Proxy for Value Estimation)

별도의 가치 네트워크를 학습시키는 대신(이는 재학습을 요구함), GILC는 사전 학습된 디노이징 네트워크를 변분 프록시로 재활용합니다.

• 이 프레임워크는 프록시 분포와 실제 역과정(reverse process) 분포 사이의 KL 발산을 최소화합니다.
• 평균장 가정(mean-field assumption) 하에서, 이 목적 함수는 원래의 확산 모델을 학습시키는 데 사용되는 음의 로그 가능도 손실(negative log-likelihood loss)과 정확히 일치합니다.
• 결과적으로, 사전 학습된 네트워크의 깨끗한 예측 로짓($x_\theta$)은 프록시 분포의 파라미터 역할을 수행하며, 이를 통해 이 프록시로부터의 몬테카를로 샘플링을 통해 기대 보상을 학습 없이(training-free) 추정할 수 있습니다.

2. 그래디언트 정보 기반 로짓 교정 (Gradient-Informed Logit Correction)

가이던스 신호를 계산하기 위해, 방법론은 가치 함수의 그래디언트를 추정해야 합니다. 저자들은 노이즈가 섞인 상태 $z_t$를 통해 전체 그래디언트를 계산하는 과정에 모델 야코비안($\partial \eta / \partial z_t$)이 포함되며, 이것이 이산 공간에서는 종종 조건이 나쁘고(ill-conditioned) 수치적으로 불안정하다는 점을 확인했습니다.

• 야코비안 프리 메커니즘 (Jacobian-Free Mechanism): SDS(Score Distillation Sampling)에서 영감을 얻은 GILC는 불안정한 모델 야코비안을 완전히 우회합니다.
• 로짓 공간 교정: 노이즈가 섞인 이산 상태를 통해 미분하는 대신, 이 방법은 깨끗한 예측 로짓 공간($\eta$)에서 직접 그래디언트를 계산합니다. 가이던스 신호는 로짓에 대한 교정값으로 적용됩니다: $\eta^r = \eta + g_\eta / \beta$.
• 이는 보상의 정보가 담긴 방향을 보존하면서도 안정적인 그래디언트 흐름을 보장합니다.

3. 미분 가능 및 비미분 목적 함수 처리

GILC는 서로 다른 유형의 보상 함수를 수용하기 위해 두 가지 변형을 제공합니다:

• GILC-DB (직접 역전파, Direct Backpropagation): 미분 가능한 보상의 경우, Gumbel-Softmax 트릭과 Straight-Through (ST) 추정기를 결 함께 사용합니다. 이를 통해 소프트 샘플(soft samples)을 역전파에는 사용하되, 순전파(forward pass) 시에는 보상 함수에 하드 원-핫 벡터(hard one-hot vectors)를 전달함으로써 이산 샘플링 과정을 통해 그래디언트가 흐를 수 있게 합니다.
• GILC-PG (정책 그래디언트, Policy Gradient): 미분 불가능하거나 블랙박스 형태의 보상의 경우, 정책 그래디언트 공식을 활용합니다. 몬테카를로 샘플을 사용하여 그래디언트를 추정하고 그룹 상대적 이점(group relative advantages, GRPO와 유사)을 계산함으로써, 프레임워크가 보상 함수의 미분 가능 여부에 구애받지 않도록 만듭니다.

주요 기여

1. 학습 불필요 프레임워크: 사전 학습된 디퓨전 네트워크를 변분 프록시로 사용하여 플러그 앤 플레이 방식의 가치 함수 추정이 가능하다는 것을 입증하였으며, 이를 통해 부가적인 학습의 필요성을 제거했습니다.
2. 안정적인 로짓 교정: 깨끗한 예측 로짓을 직접 교정하는 야코비안 프리 업데이트 메커니즘을 도입하여, 고차원 이산 공간에 내재된 그래디언트 불안정성 문제를 해결했습니다.
3. 보편성: 정책 그래디언트와의 공식적인 연결을 확립하여, 비미분 목적 함수에 대한 프레임워크의 적용 가능성을 확장했습니다.
4. 실증적 검증: 다양한 과학적 도메인에 걸쳐 추가 학습 없이 최첨단 성능을 보여주는 광범위한 증거를 제공합니다.

실험 결과

저자들은 세 가지 과학적 도메인과 이산 이미지 생성에 대해 GILC를 평가했습니다.

• 조절 DNA 서열 설계 (Regulatory DNA Sequence Design): GILC-DB는 가장 높은 예측 활성도(Pred-Activity)와 염색질 접근성(ATAC-Acc)을 달성하였으며, 이는 학습 불필요 베이스라인(SMC, SVDD)을 크게 상회할 뿐만 아니라 DRAKES와 같은 미세 조정된 방법들조차 능가했습니다. 또한 높은 서열 충실도(fidelity)를 유지했습니다.
• 단백질 서열 설계 (Protein Sequence Design): 역폴딩(inverse folding) 작업에서, GILC-DB는 우수한 열역학적 안정성(최고의 Pred-ddG)을 가진 서열을 생성하였으며, 안정성과 구조적 제약을 모두 만족하는 최고 수준의 전체 성공률을 달성하여 DRAKES보다 약 4%포인트 높은 성능을 보였습니다.
• 멀티모달 분자 생성 (Multimodal Molecular Generation): QM9 데이터셋에서 GILC-DB와 GILC-PG는 표적 양자 특성(예: 편극도, 쌍극자 모멘트)을 가진 분자를 생성하는 데 있어 비지도 학습 베이스라인보다 일관되게 우수한 성능을 보였습니다. 특히 GILC-PG는 예측기를 블랙박스로 취급함으로써 높은 유연성을 입증했습니다.
• 이미지 생성 (Image Generation): 본 프레임워크는 CIFAR-10에서의 클래스 조건부 생성과 Meissonic 모델을 사용한 텍스트-투-이미지 합성에 성공적으로 확장되었으며, 아키텍처 변경 없이 고품질의 클래스 일관적이고 미적으로 뛰어난 이미지를 생성했습니다.

효율성: GILC는 단일 샘플링 궤적(trajectory)으로 작동하며, 여러 궤적을 유지해야 하는 Best-of-N 또는 SMC와 같은 방법들에 비해 보상 함수와 디노이징 모델 호출 횟수를 현저히 적게 요구합니다.

의의 및 주장

본 논문은 GILC가 **제어 가능한 이산 생성(controllable discrete generation)**의 중요한 진전임을 주장합니다. 재학습의 필요성을 없애고 기존의 학습 불필요 방식들과 관련된 계산 비용을 완화함으로써, GILG는 데이터가 희소하거나 목적이 빈번하게 바뀌는 과학적 응용 분야에 다각적인 솔루션을 제공합니다.

저자들은 다음과 같은 점을 강조합니다:

• 보상 해킹(Reward Hacking) 방지: 미세 조정 방식과 달리 GILC는 모델 파라미터를 변경하지 않으므로, 데이터 분포의 충실도를 희생하면서 보상에 과적합되는 위험을 줄입니다.
• 연속 및 이산 가이던스의 가교 역할: 새로운 로짓 교정 메커니즘을 통해 연속 공간 디퓨전의 그래디언트 기반 가이던스 직관을 이산 공간으로 성공적으로 전이시켰습니다.
• 광범위한 적용 가능성: 미분 가능 및 비미분 목적 함수를 모두 처리할 수 있는 능력은 유전자 치료 및 약물 설계와 같은 광범위한 실제 과학 문제에 적합합니다.

결론적으로, GILC는 계산 오버헤드를 줄이지만, 향후 연구는 필요한 몬테카를로 샘플 수를 더욱 최소화하고 자연어와 같은 구조화된 도메인을 더 잘 다루기 위해 토큰 독립성 가정을 완화하는 데 집중할 수 있다고 밝히고 있습니다.