Dynamic Search for Inference-Time Alignment in Diffusion Models

Motivation

• Known: 확산 모델은 다양한 도메인에서 강력한 생성 능력을 보유하고 있으나, 보상 함수와의 정렬이 어려운 문제가 존재. 기존의 분류기 유도(classifier guidance)는 미분 가능한 보상 함수에 의존하는 한계 있음.
• Gap: 실제 과학 응용 분야(약물 설계, 단백질 구조 예측 등)에서 AutoDock Vina, AlphaFold3, DSSP 등의 보상 함수는 물리 시뮬레이션, 룩업 테이블 기반으로 비미분이거나 블랙박스 형태. 기존 그래디언트 프리 유도 방법들도 최적의 정렬을 보장하지 못함.
• Why: 확산 모델의 디노이징 과정(denoising process)이 트리 구조를 형성한다는 통찰을 기반으로, 검색 알고리즘을 통해 더 효율적인 정렬이 가능할 것으로 예상.
• Approach: 사전학습된 확산 모델의 디노이징 과정을 트리로 형식화하고, 동적 빔 탐색(dynamic beam search)을 적용하여 빔 폭과 트리 너비를 시간 단계에 따라 동적으로 조정.

Achievement

1. 탐색 프레임워크 제안: 확산 모델의 디노이징 프로세스를 트리 구조로 형식화하여, 보상 최적화를 체계적인 탐색 문제로 재구성. 이는 기존의 ad-hoc한 유도 방식과 달리 일관된 이론적 기초 제공.
2. 동적 빔 탐색(DSearch) 알고리즘: 고정 너비 빔 탐색의 비효율성을 해결하기 위해, 시간 단계별로 빔 폭 b(t)와 트리 너비 w(t)를 동적으로 조정. 약한 빔의 계산 자원을 다른 빔으로 재할당하여 효율성 극대화(w(t)·b(t) 고정).
3. 다중 도메인 검증: 생물학적 수열 설계(biological sequence design), 분자 최적화(molecular optimization), 이미지 생성 등 다양한 도메인에서 기존 방법 대비 우수한 보상 최적화 성과 입증.

How

Figure 2: DSearch의 트리 너비 확장과 빔 폭 동적 조정. 약한 빔의 자원을 다른 빔으로 재할당하면서 w(t)b(t) 유지

트리 정의 및 너비 제한:

• 나이브 접근의 O(|X|^T) 복잡도를 해결하기 위해 사전학습된 정책으로부터 샘플링하여 트리 너비 w(t) 제한
• w(t)=1일 경우 best-of-N 샘플링으로 축소되는 일반적 형태 유지

휴리스틱 함수(Heuristic Function):

• 중간 노드의 가치를 평가하기 위해 추정된 가치 함수(estimated value function) 도입
• 기존의 단순 근사 ν̂_t(x_t) := r(x̂_0(x_t))를 개선하는 더 정확한 접근법 제시

룩어헤드 휴리스틱(Lookahead Heuristic):

• Algorithm 1에서 K 스텝 선점 탐색을 통해 소프트 가치 함수의 근사 정확도 향상
• 소수의 추가 시뮬레이션으로 더 신뢰할 만한 중간 노드 평가 가능

노이즈 레벨 기반 동적 스케줄링:

• 노이즈 레벨에 따라 적응적으로 트리 확장 일정 조정
• 초기 단계(높은 노이즈)에서 넓은 탐색, 후기 단계(낮은 노이즈)에서 선택적 탐색

Originality

• 탐색 프레임워크의 혁신: 확산 모델의 추론 시간 정렬을 처음으로 체계적인 트리 탐색 문제로 정의. 이는 기존의 기울기 기반 또는 휴리스틱한 유도 방식과 근본적으로 다른 접근.
• 동적 빔 폭 조정: 고정 빔 탐색의 비효율성을 인식하고, 시간 단계와 노이즈 레벨에 따라 자원을 동적으로 재할당하는 새로운 전략 제시. 이는 단순하지만 효과적인 개선.
• 향상된 휴리스틱 함수: 기존의 단순한 x̂_0 기반 근사를 넘어, 룩어헤드 탐색을 통한 더 정확한 중간 노드 가치 추정 방법 개발.
• 비미분 보상 함수 지원: 그래디언트가 필요 없는 완전한 그래디언트 프리 프레임워크로, 실제 과학 응용의 복잡한 블랙박스 보상 함수에 직접 적용 가능.

Limitation & Further Study

한계:

• 트리 너비 w(t)와 빔 폭 b(t)의 설정이 휴리스틱하며, 최적값 선택에 대한 이론적 지침이 부족. 다양한 도메인에서의 하이퍼파라미터 민감도 분석 필요.
• 룩어헤드 스텝 K의 증가에 따른 계산 비용 증가로 인한 트레이드오프 미분석. 실제 적용 시 계산 예산(computational budget) 제약에서의 최적 K 선택 방법 미제시.
• 샘플 다양성(diversity)과 자연스러움(naturalness) 사이의 균형에 대한 이론적 분석 부족. 온도 파라미터 α의 설정과 성과 간의 정량적 관계 미제시.

후속 연구:

• 동적 빔 폭 조정의 최적성에 대한 이론적 분석 (예: regret bounds 도출)
• 다양한 도메인/모델에 대한 자동 하이퍼파라미터 선택 전략 개발
• 대규모 생성 모델(예: 텍스트 생성, 동영상 생성)으로의 확장 및 확장성 검증
• 여러 보상 함수의 다목적 최적화(multi-objective optimization)로의 확대

Evaluation

총평: 본 논문은 확산 모델의 추론 시간 정렬을 체계적인 트리 탐색 문제로 재정의하고 동적 빔 폭 조정으로 비미분 보상 함수를 효율적으로 최적화하는 실용적이고 참신한 방법을 제시한다. 다만 하이퍼파라미터 설정의 휴리스틱함과 동적 조정의 이론적 최적성 분석이 부족하며, 계산 비용과 성능의 트레이드오프에 대한 더 상세한 분석이 필요하다.