이산 확산에서의 분할 순차 몬테카를로를 위한 대조적 분포 정합

초록

이산 확산 모델은 구조화된 범주형 데이터를 생성하기 위한 강력한 프레임워크로 부상했다. 그러나 보상이 기울어진 분포(reward-tilted distributions)로부터 효율적으로 샘플링하는 것은 여전히 근본적인 도전 과제로 남아 있다. Twisted 순차적 몬테카를로(SMC)는 이 작업에 대해 점근적 정확성을 제공하지만, 이산 상태 공간에서 최적의 트위스트 함수를 추정하려면 비용이 많이 드는 몬테카를로 근사가 필요하며, 이는 추론 시 심각한 계산 병목 현상을 초래한다. 이러한 한계를 극복하기 위해 우리는 양성 및 음성 샘플을 통해 매개변수화된 트위스트 함수를 학습함으로써 SMC 추론 비용을 분할 상환하는 새로운 프레임워크인 대비 분포 매칭(CDM)을 소개한다. 효율적인 훈련을 위해, 우리는 이산 확산 모델의 폐쇄형 순방향 커널을 활용하도록 그래디언트 추정기를 재구성한다. 실제로 학습된 트위스트 함수를 평가하는 데는 기본 모델의 단일 순방향 전달에 비해 5% 미만의 추가 계산 오버헤드만 발생한다. 광범위한 실증 평가를 통해, 우리는 CDM이 동일한 벽시계 시간 조건에서 기존의 기준 모델들을 일관되게 능가함을 보여준다. 우리는 유해 텍스트 생성, 조절 DNA 서열 설계, 단백질 설계 가능성, 확산 대규모 언어 모델 정렬 등 다양한 응용 분야에서 우리 접근법의 효과성과 다용도성을 검증한다.

English

Discrete diffusion models have emerged as powerful frameworks for generating structured categorical data. However, efficiently sampling from reward-tilted distributions remains a fundamental challenge. While Twisted Sequential Monte Carlo (SMC) offers asymptotic exactness for this task, estimating the optimal twist function in discrete state spaces necessitates costly Monte Carlo approximations, resulting a severe computational bottleneck at inference. To overcome this limitation, we introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples. For efficient training, we reformulate the gradient estimator to leverage the closed-form forward kernels of discrete diffusion models. In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model. Through extensive empirical evaluations, we demonstrate that CDM consistently outperforms existing baselines under matched wall-clock time. We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.