Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design

Motivation

• Known: 확산 모델은 텍스트-이미지, 단백질 서열 생성 등 다양한 분야에서 성공적이며, 분류기 유도(classifier guidance)나 도함수-프리(derivative-free) 방법들이 보상 유도 생성(reward-guided generation)에 활용되고 있다.
• Gap: 기존 방법들은 완전 노이징 상태에서 디노이징 상태로의 단일 패스에만 의존하며, 두 가지 핵심 문제가 존재한다: (1) 보상 유도 디노이징 중 발생한 오류를 수정할 메커니즘 부재, (2) 특히 마스크 확산 모델에서 한번 변경된 토큰은 끝까지 고정되는 문제, (3) 하드 제약조건(hard constraints) 처리의 어려움.
• Why: 생물학적 서열 설계(단백질, DNA)에서는 구조적 안정성, 결합 친화력, 세포타입 특이성 등 복잡한 보상함수 최적화가 필요하며, 추론 시간 계산량을 증가시켜 더 나은 설계를 얻을 수 있을 것으로 예상된다.
• Approach: 테스트 타임에 임의의 양의 계산을 활용하여 설계를 지속적으로 개선하는 반복적 개선 알고리즘(noising + reward-guided denoising) 제안.

Achievement

그림 2: 기존 보상 유도 알고리즘은 소프트 최적 정책 {p⋆_t}로부터 순차적 샘플링으로 볼 수 있으며, 알고리즘의 차이는 p⋆_t 근사 방식에 있다

1. 이론적 기여: 제안된 알고리즘이 exp(r(x))p_pre(·) 분포로부터 샘플링함을 수학적으로 증명하여, 생성된 설계의 자연스러움(naturalness)과 보상 최적화 간 균형을 이론적으로 보장.
2. 방법론 혁신: 단순한 반복적 개선을 통해 마스크 확산 모델의 근본적 한계(한번 변경된 토큰 고정)를 극복하고, 하드 제약조건을 포함하는 복잡한 보상함수 최적화 가능.
3. 실험적 우수성: 단백질 구조 설계(target RMSD 최소화)와 세포타입 특이성 DNA 설계에서 기존 방법들을 능가하는 성능 달성.

How

그림 3: RERD 알고리즘 요약 - 반복적으로 부분 노이징과 보상 유도 디노이징 수행

핵심 알고리즘 구조:

• 반복적 단계: 각 반복 k에서 (1) 현재 샘플 x^(k)에 부분 노이즈 주입 → 중간상태 생성, (2) 보상 유도 디노이징을 통해 x^(k+1) 획득
• 소프트 최적 정책 근사: 식 (2)의 소프트 가치함수 v_t(x_t)를 근사하기 위해 재구성된 x_0 예측값 x̂_0(x_t)의 보상 r(x̂_0(x_t)) 활용
• 대규모 행동공간 처리: 분류기 유도(연속) 또는 중요도 샘플링(이산)으로 p⋆_t 근사
• 하드 제약조건 처리: 초기 시드 시퀀스를 가능 영역 C 내에서 선택하여 제약조건 자동 만족
• 진화 알고리즘과의 연결: 보상 기반 선택과 부분 변이(노이징)의 조합이 유전 알고리즘(genetic algorithms)과 유사한 구조

Originality

• 첫 확산 모델 반복 개선: 언어모델의 반복 개선(BERT-style refinement) 개념을 확산 모델에 처음 적용하여 테스트 타임 보상 최적화의 새로운 패러다임 제시.
• 오류 수정 메커니즘: 기존 단일 패스 방식의 결정 불가역성 문제를 부분 노이징을 통한 재샘플링으로 해결 - 특히 마스크 확산 모델에서 매우 중요.
• 제약조건 통합 설계: 하드 제약조건을 reward function으로 단순 설정하는 대신, 초기 조건 선택을 통해 실질적으로 달성하는 실용적 방안 제시.
• 통합 이론 프레임워크: KL 정규화된 보상 최적화 목표(식 1)와 확산 모델의 수학적 연결을 명확히 하고, 제안 알고리즘의 최적성 보장.

Limitation & Further Study

• 계산비용 증가: 반복적 개선으로 인한 추론 시간 증가는 실제 응용에서 병목이 될 수 있으며, 효율성 개선 방안 필요.
• 부분 노이징 스케줄 미최적화: 각 반복에서 주입할 노이즈 수준(noise level)을 선택하는 전략이 휴리스틱으로 보이며, 이론적 최적화 부재.
• 실제 검증 제한: 단백질 설계의 경우 구조 예측(ESMFold)만 사용하고 실제 생화학적 검증 부재, DNA 설계도 세포 실험 검증 미흡.
• 보상함수 품질 의존성: 알고리즘의 성능이 보상함수의 정확성에 크게 의존하지만, 부정확한 보상함수에 대한 견고성(robustness) 분석 부족.
• 후속 연구: (1) 적응형 노이징 스케줄 학습, (2) 불완전한 보상함수 환경에서의 성능 특성화, (3) 다른 생성 모델(flow-based, autoregressive)로의 확장, (4) 실제 생물학적 검증 수행.

Evaluation

총평: 확산 모델의 테스트 타임 최적화에 혁신적인 반복 개선 접근을 제시하고, 특히 마스크 확산의 토큰 고정 문제 해결과 하드 제약조건 처리는 실질적 기여다. 단백질/DNA 설계에서 일관된 성능 향상을 보이나, 계산 효율성 분석 부재와 실제 생물학적 검증 부족이 한계. 학술적 우수성은 높으나 실무 적용을 위해서는 효율화와 검증이 필요하다.