기술 요약: 마스크된 확산 모델을 위한 VGB (VGB for Masked Diffusion Model)

문제 정의

많은 시퀀스 생성 작업은 전역적인 구조적 제약 조건을 충족하거나 다운스트림 보상을 최적화하는 출력을 요구합니다 (예: 수학적 증명, 코드, 스도쿠, 단백질 설계). 완전한 구성(complete configurations)은 결과 검증기(outcome verifiers, 예: 컴파일러, 제약 조건 체크 도구)를 통해 검증할 수 있는 경우가 많지만, 부분적인 구성을 보고 이것이 유효한 전체 시퀀스로 확장될 수 있는지 판단하는 것은 훨씬 더 어렵습니다.

기존의 추론 시간 스케일링(inference-time scaling) 방법들은 다음과 같은 트레이드오프에 직면해 있습니다:

1. Best-of-N (BoN): $N$개의 완전한 시퀀스를 샘플링한 후 가장 좋은 것을 선택합니다. 이는 시퀀스 길이에 대해 지수적 복잡도를 초래하는데, 이는 초기 오류를 수정할 수 없기 때문입니다.
2. 고정 순서 백트래킹 (Fixed-Order Backtracking, AR-VGB): Jerrum-Sinclair 마르코프 체인에서 영감을 받은 이 방법은 백트래킹을 허용하지만, 고정된 순서(예: 좌에서 우로)로 제한됩니다. 이는 가장 최근에 생성된 토큰만을 삭제할 수 있으므로, 구조적 생성이나 편집 작업에서 발생하는 "치명적인" 초기 오류를 복구하는 데 비효로적입니다.
3. 프로세스 검증기 노이즈 (Process Verifier Noise): 프로세스 검증기에 의존하여 생성을 가이드하는 방법들은 부분적인 구성의 가치를 추정하는 과정에서 발생하는 미세한 오류가 기하급수적으로 누적되어 최종 결과가 저하되는 "지평선의 저주(curse of horizon)" 문제를 겪는 경우가 많습니다.

본 논문은 임의 순서 백트래킹(임의 위치의 토큰을 언마스킹 및 재마스킹)을 지원하면서도, 노이즈가 있는 프로세스 검증기가 존재하는 상황에서도 보상 기울기 분포(reward-tilted distribution)로 수렴한다는 이론적 보장을 유지하는 샘플러의 필요성을 다룹니다.

방법론: MDM-VGB

저자들은 고정 순서 접두사 트리(fixed-order prefix tree)가 아닌 마스크 상태 그래프(masked-state graph) 상에서 작동하는 이산 확산 샘플러인 **MDM-VGB (Masked Diffusion Model Value-Guided Backtracking)**를 소개합니다.

핵심 메커니즘

• 상태 공간 (State Space): 샘플러는 일부 토큰은 드러나 있고 나머지는 마스크 처리된 부분적 구성(partial configurations) 상에서 작동합니다.
• 그래프 구조: AR-VGB에서 사용하는 접두사 트리와 달리, 상태 공간은 하나의 좌표를 드러내거나(전진 이동) 다시 마스킹하는(후진 이동) 단일 변화로 연결되는 무방향 그래프입니다.
• 전이 확률 (Transition Probabilities): 샘플러는 이 그래프 위에서 랜덤 워크를 수행합니다. 전이 확률은 추정된 가치가 높은 상태 쪽으로 편향됩니다.
  • 전진 이동 (Reveal): 참조 모델의 조건부 확률과 결과로 나타나는 자식 상태의 추정 가치에 의해 가중치가 부여됩니다.
  • 후진 이동 (Re-mask): 현재 상태의 가치와 결과로 나타나는 부모 상태의 가치에 의해 가중치가 부여됩니다. 결정적으로, 이 방법은 $\lambda$ 파라미터에 의해 제어되는 기하학적 게이팅(geometric gating) 메커니즘을 도입합니다. 이를 통해 샘플러는 제거했을 때 더 높은 추정 가치를 가진 부모 상태로 이어지는 토큰을 우선적으로 재마스킹함으로써, "신뢰할 수 없는" 토큰을 효과적으로 식별하고 수정할 수 있습니다.
• 가치 함수 (Value Function): 프로세스 검증기 $\hat{V}$는 주어진 부분 상태에 대해 호환 가능한 완전한 구성들의 기대 보상을 근사합니다. 이 방법은 검증기의 승법적 오차(multiplicative errors)에 대해 강건합니다.

변형 모델

1. AOAR-VGB (Any-Order Auto-Regressive): 싱글톤 블록 업데이트(한 번에 하나의 토큰씩)를 사용하는 가장 단순한 변형입니다. 이는 가장 깔끔한 이론적 보장을 제공합니다.
2. MDM-VGB-Momentum: 각 상태에 방향 변수(드러내기 vs 재마스킹 모드)를 추가하여 체인을 증강한 "모멘텀 리프트(momentum lift)" 버전입니다. 이는 확산적 진동(즉각적인 앞뒤 왕복 이동)을 줄이고, 특히 저오차 영역에서 첫 번째 리프(first-leaf) 도달 시간을 개선합니다.

주요 기여

1. 알고리즘 혁신: 고정 순서 트리에서 임의 순서 마스크 상태 그래프로 Jerrum-Sinclair 백트래킹 체인을 일반화한 MDM-VGB를 도입했습니다. 이를 통해 생성 및 편집 작업에서 초기 실수를 효율적으로 복구할 수 있습니다.
2. 이론적 보장:
   • 정상 분포 (Stationary Distribution): 저자들은 불완전한 프로세스 검증기가 있더라도, 리프(완전한 구성)에 도달하는 조건 하에서 MDM-VGB 체인의 정상 분포가 정확히 보상 기울기 타겟 분포 $\pi^*(y|x) \propto \pi_{ref}(y|x)\tau(x,y)$임을 증명했습니다.
   • 복잡도: 샘플러의 믹싱 타임(mixing time)은 시퀀스 길이 $n$에 대해 이차적($O(n^2)$)인 반면, Best-of-N은 지수적 비용이 듭니다.
   • 강건성: 샘포러는 프로세스 검증기의 상당한 승법적 노이즈에 대해 강건함이 증명되었습니다.
   • 모멘텀 가속: 모멘텀 변형의 경우, 근사 균형 조건 하에서 리프에 도달하는 기대 시간은 $n$에 대한 선형($O(n)$)으로 감소하여, 비모멘텀 버전 대비 이차적 속도 향상을 제공합니다.
3. 실험적 성능: MDM-VGB는 Best-of-N, MDM-VGR(순방향 전용 가치 가이드), AR-VGB와 같은 베이스라인에 비해 생성 및 편집 작업 모두에서 우수한 성능을 보여줍니다.

실험 결과

저자들은 여러 벤치마크에서 MDM-VGB와 MDM-VGB-Momentum을 평가했습니다:

• 제약 조건 충족: 스도쿠 및 Dyck 문법 편집.
• 과학적 설계: 분자 생성(QM9), DNA 인핸서 설계, 단백질 모티프 스캐폴딩.

주요 결과:

• 생성: MDM-VGB-Momentum은 동일한 추론 비용에서 Best-of-N 및 MDM-VGR에 비해 더 높은 보상 만족도(예: Pass@95)를 달착합니다. 특히 베이스 모델의 정확도가 낮은 어려운 작업에서 효과적입니다.
• 편집: 편집 작업(낮은 보상의 샘플을 복구하는 작업)에서 MDM-VGB는 고정 순서 백트래킹보다 크게 앞섭니다. Dyck 문법 편집의 경우, AR-VGB(정확도 25.64%, 이동 횟수 127.63)에 비해 훨씬 적은 편집 횟수(29.98)만으로도 거의 완벽한 정확도(99.41%)를 달성했습니다.
• 효율성: 모멘텀 변형은 고보상 구성을 달성하는 데 필요한 이동 횟수를 줄여 품질-비용 프런티어를 개선합니다.
• 검증기 강건성: 이 방법은 완벽하지 않은 학습된 검증기를 사용하더라도 잘 작동하여, 이론적 강건성 주장을 입증합니다.

의의 및 주장

본 논문은 MDM-VGB가 이산 확산 모델을 위한 원칙적이고 이론적으로 근거가 확실한 테스트 타임 스케일링 프레임워크를 제공한다고 주장합니다. 주요 의의는 다음과 같습니다:

• "지평선의 저주" 극복: 임의 순서 백트래킹을 허용함으로써, 고정 순서 방식의 오류 누적 문제와 비백트래킹 탐색의 비효율성을 피합니다.
• 통합된 생성 및 편집: 동일한 메커니즘이 고보상 시퀀스를 처음부터 생성하는 것과 기존의 저보상 시퀀스를 복구하는 것 모두를 지원하며, 특히 후자에서 강력한 이점을 가집니다.
• 이론과 실제의 일치: 조합 퍼즐부터 과학적 설계에 이르기까지 다양한 도메인에서 입증된 대로, 방법론은 강력한 이론적 보장(이차 믹싱 타임, 정확한 정상 분포)과 실험적 결과를 결합합니다.

저자들은 현재 작업이 프로세스 검증기 학습에 의존하고 있으며, 이는 긴 시퀀스와 큰 어휘 집합에 대해 비용이 많이 들 수 있다고 언급하며, 향-규모의 벤치마크로 확장하기 위한 검증기 증류(distillation) 및 계층적 업데이트에 대한 향후 연구를 제안합니다.