Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction

Motivation

• Known: 기존 discrete diffusion의 controllable generation은 classifier/classifier-free guidance나 fine-tuning 기반 방법이 주로 사용되었으며, 연속공간의 plug-and-play guidance(예: diffusion posterior sampling)는 이미지 생성 등에서 성공적으로 적용되어 왔다.
• Gap: discrete diffusion에서는 categorical data의 non-differentiability로 인해 연속공간의 plug-and-play guidance 기법을 직접 적용하기 어렵고, 기존 SMC나 importance sampling 기반 방법은 계산 비용이 크고 steering 효과가 제한적이다.
• Why: 추가 학습 없이 다양한 reward function(differentiable/non-differentiable 모두)을 활용해 DNA, 단백질, 분자 생성 등 과학적 응용 분야에서 controllable generation을 효율적이고 안정적으로 수행할 수 있다는 점에서 실용적, 이론적 가치가 크다.
• Approach: 사전학습된 denoising network를 value function 추정을 위한 variational proxy로 재해석하고, Gumbel-Softmax와 Straight-Through estimator를 결합한 뒤 clean prediction logits에 직접 gradient 기반 보정을 가하는 Jacobian-free 메커니즘을 제안한다.

Achievement

Figure 5. Performance comparison on DNA sequence design. The deterministic estimation (DE) baseline is compared with Mon

1. Training-free guidance framework 제안: 사전학습된 diffusion network를 재학습 없이 value function proxy로 활용해 plug-and-play guidance를 구현.
2. Logits correction guidance: clean prediction logits에 직접 value gradient를 적용해 discrete space에서 발생하는 gradient instability(Jacobian 관련 수치 불안정성)를 해결.
3. Policy gradient와의 이론적 연결: non-differentiable reward function까지 다룰 수 있도록 프레임워크의 일반성을 확장.
4. State-of-the-art 성능 및 효율성 입증: DNA, 단백질 서열, 분자 생성 등 복수 과학 도메인에서 fine-tuning 기반 방법과 대등하거나 이를 능가하는 성능을 추가 학습 없이 달성.

How

Figure 2. Analysis of value estimation and guidance targets. (a)

• Reward-guided 생성을 KL-regularized 최적화 문제(pr_θ(x) = argmax_q E[r(x)] - β D_KL(q||p_pre))로 정식화.
• 사전학습된 denoising network xθ(zt,t)를 재사용해 중간 상태에서의 expected future reward(value function)를 추정하는 variational proxy 도입.
• Gumbel-Softmax trick과 Straight-Through estimator를 결합해 discrete space에서도 gradient flow를 유지.
• Jacobian 계산을 우회하는 방식으로 clean prediction logits에 직접 reward gradient를 적용해 guided prediction xr_θ를 산출, 이를 posterior q(zs|zt, xr_θ)에 대입해 다음 상태 zs를 샘플링.
• differentiable reward에는 gradient 기반 보정을, non-differentiable reward에는 policy gradient 형태의 추정을 적용하여 통합적 guidance 체계 구성.
• DNA sequence design, protein sequence engineering, multimodal molecular generation 등 다양한 도메인에서 실험적으로 검증.

Originality

• discrete diffusion에서 사전학습된 denoising network 자체를 value function 추정용 variational proxy로 재해석한 새로운 관점 제시.
• 기존 연속공간 plug-and-play guidance 기법(gradient 기반)을 discrete space에 안정적으로 이식하기 위한 Jacobian-free logit correction 메커니즘 고안.
• 미분 가능/불가능 reward를 policy gradient 관점에서 통합적으로 다룰 수 있는 이론적 연결고리 제시.

Limitation & Further Study

• Jacobian-free 근사가 실제 gradient와 얼마나 정확히 일치하는지에 대한 이론적 오차 분석이 제한적일 수 있음.
• Gumbel-Softmax 및 Straight-Through estimator에 내재된 bias가 다양한 도메인 및 reward 함수에 대해 어떻게 영향을 미치는지 추가 검증이 필요.
• β(KL trade-off) 등 하이퍼파라미터 민감도 및 다양한 vocabulary 크기·시퀀스 길이에서의 확장성에 대한 추가 실험이 요구됨.
• 후속 연구로 policy gradient 연결을 활용한 non-differentiable reward에 대한 분산 감소 기법이나 다중 reward 결합 시의 안정성 연구가 필요.

Evaluation

총평: discrete diffusion model의 training-free guidance라는 실용적으로 중요한 문제를 참신한 variational proxy와 Jacobian-free logit correction으로 효과적으로 해결한 견고한 연구로 평가된다.