Reward-Guided Discrete Diffusion via Clean-Sample Markov Chain for Molecule and Biological Sequence Design

Motivation

• Known: 이산 확산 모델은 분자와 DNA 서열 같은 구조적 데이터 생성에 적합하며, 최근 테스트타임 스케일링을 통한 보상 유도 샘플링이 활발히 연구되고 있다. 기존 방법들(SMC, SVDD 등)은 중간 보상을 활용하거나 Best-of-N 샘플링을 사용한다.
• Gap: 과학 분야의 보상 함수가 비매끄러우므로 (예: 한 글자 변경으로 분자 구조 무효화), 중간 노이즈 샘플에서의 보상 계산은 신뢰할 수 없다. 기존 중간 보상 기반 방법들은 이러한 한계를 제대로 해결하지 못한다.
• Why: 과학 도메인에서 이산 보상 함수의 비매끄러움으로 인해 중간 보상 신호는 노이즈가 많고 오도할 수 있다. 깨끗한 샘플에만 직접 접근하는 방식이 더 효과적인 국소 탐색을 가능하게 한다.
• Approach: Forward process로 깨끗한 샘플을 노이즈화한 후 reverse process를 적용하여 새로운 샘플을 제안하는 forward-backward 조합을 사용한다. Metropolis-Hastings 알고리즘으로 acceptance probability를 계산하고, 이 제안 분포가 acceptance probability를 계산 가능하게 만든다.

Achievement

Figure 1. Left: In scientific applications, the rewards defined on discrete spaces are highly sensitive to small perturb

• 분자 생성 성능 향상: QED, synthetic accessibility, ring count 등 다양한 보상 함수에서 기존 SMC, SVDD 대비 우수한 성능
• 생물 서열 생성: enhancer activity 등 DNA 서열 품질 지표에서 최고 성능
• 일반적 적용 가능성: Uniform masked 양쪽 이산 확산 프레임워크에 적용 가능 (Table 1)
• SMILES 생성: 다른 방법이 실패하는 SMILES 문자열 생성에서도 성공

"

How

Figure 1. Left: In scientific applications, the rewards defined on discrete spaces are highly sensitive to small perturb

• Forward-backward proposal distribution으로 깨끗한 샘플 간 전이 제안
• Metropolis-Hastings acceptance probability로 제안 수락 여부 결정
• 깨끗한 샘플에만 직접 보상 계산 적용
• Stationary distribution이 목표 분포가 되도록 마코프 체인 구성

"

Originality

• 이산 확산에서 깨끗한 샘플만으로 마코프 체인 구성하는 새로운 접근
• Forward-backward 결합을 proposal distribution으로 사용하여 intractable acceptance probability를 tractable하게 변환
• 중간 보상 계산을 완전히 회피하는 방식은 기존 방법들과 차별화

"

Limitation & Further Study

• CSMC의 수렴 속도나 샘플 효율에 대한 이론적 분석 부족
• Forward-backward 반복 횟수에 따른 계산 비용 증가 상세 분석 미흡
• 극도로 높은 차원의 생물 서열에서의 성능 평가 제한적
• Reward function 설계 방식에 따른 성능 민감도 미분석

"

Evaluation

총평: CSMC Sampler는 이산 확산 모델의 보상 유도 샘플링 문제를 창의적으로 해결하는 방법이다. 깨끗한 샘플 기반 마코프 체인 구성으로 과학 도메인의 비매끄러운 보상 함수 문제를 우아하게 극복하며, 분자·생물 서열 생성에서 일관된 성능 향상을 보인다. 다만 이론적 수렴성 분석과 계산 복잡도 상세 분석이 강화되면 더욱 우수할 것으로 판단된다.