Diffusion Language Model Parallel Decoding via Product-of-Experts Bridge

기술 요약: Product-of-Experts Bridge를 통한 확산 언어 모델의 병렬 디코딩

문제 정의
확산 언어 모델(Diffusion Language Models, DLMs)은 자기회귀(Autoregressive, AR) 모델의 대안으로서 유망한 가능성을 제시하며, 병렬 디코딩을 가능하게 하여 이론적으로 순차적인 AR 모델이 소요되는 시간의 극히 일부만으로 전체 시퀀스를 생성할 수 있게 한다. 그러나 현재 DLM은 강력한 AR 모델에 비해 상당한 품질 격차를 겪고 있다. 이러한 성능 저하는 병렬 디코딩을 위해 요구되는 조건부 독립성 가정에서 기인하며, 여기서 동일한 단계 내에서 생성되는 토큰들은 결합(jointly)되어 모델링되는 대신 독립적으로 모델링된다.

이 격차를 메우기 위한 몬테카를로 방법(특히 거부 샘플링(Rejection Sampling, RS) 및 중요도 샘플링(Importance Sampling, IS))을 이용한 최근의 시도들은 근본적인 한계에 직면해 왔다. 이러한 설정에서 DLM은 제안 분포($p_D$) 역할을 하고, 강력한 AR 모델은 타겟($p_{AR}$) 역할을 한다. 병렬 DLM 제안 분포와 순차적 AR 타겟 사이의 거대한 분포 불일치로 인해 다음과 같은 문제가 발생한다:

1. **거부 샘플링(Rejection Sampling)**은 빈번한 거부를 초래하여, 디코딩 과정이 낮은 처리량과 함께 거의 순차적 생성에 가까운 상태로 붕괴된다.
2. **중요도 샘플링(Importance Sampling)**은 고품질의 샘플을 찾기 위해 비현실적으로 많은 수의 후보 입자(particle)를 필요로 하며, 제한된 예산 하에서는 진정으로 충실한 AR 샘플을 선택하기보다 "가장 덜 그럴듯한" DLM 연속형을 선택하게 된다.

방법론: PoE-Bridge
본 논문은 DLM 제안과 AR 타겟 사이의 불일치를 완화하기 위해, DLM 제안과 AR 타겟의 Product-of-Experts(PoE)로 구성된 중간 "브릿지(bridge)" 분포를 도입하는 디코딩 프레임워크인 PoE-Bridge를 소개한다.

1. PoE 브릿지 분포
PoE-Bridge는 DLM 초안을 AR 타겟으로 직접 교정하는 대신, 중간 분포 $p_{PoE}$를 통해 교정을 두 단계로 분해한다. 브릿지는 다음과 같은 토큰 수준의 보간(interpolation)으로 정의된다:
$$p_{PoE}(x_i | x_{<i}) \propto \tilde{p}_D(x_i | x_{<i})^w \cdot p_{AR}(x_i | x_{<i})^{1-w}$$
여기서 $w \in [0, 1]$은 브릿지의 위치를 제어한다. $w$가 1에 가까울수록 브릿지는 DLM 제안을 닮게 되며, 0에 가까울수록 AR 타겟에 접근한다. 이러한 토큰 수준의 정의는 샘플링과 가능도(likelihood) 평가 모두에 대해 계산 가능한 구조를 보장한다.

2. 2단계 디코딩 프로세스
프레임워크는 두 개의 병렬화 가능한 단계로 작동한다:

• 1단계: 추측적 거부 샘플링 (DLM $\to$ PoE). DLM은 여러 토큰 연속형을 병렬로 초안 작성한다. 이 초안들은 AR 타겟이 아닌 PoE 분포를 기준으로 검증된다. PoE는 AR 타겟보다 DLM 제안에 더 가깝기 때문에 수락률(acceptance rate)이 현저히 높다. 수락된 토큰은 브릿지 분포 $p_{PoE}$로부터 샘플링된다.
• 2단계: 중요도 샘플링 (PoE $\to$ AR). 1단계에서 수락된 후보들은 최종 AR 타겟 $p_{AR}$로 교정하기 위해 중요도 샘플링을 통해 재가중치(re-weighting)가 부여된다. 후보들이 이미 (PoE를 거쳐 지나왔기에) AR 분포에 더 가까워진 상태이므로, 중요도 가중치가 더 안정적이며, DLM으로부터 직접 IS를 수행할 때보다 고충실도 샘플을 얻기 위해 더 적은 수의 후보만 필요하다.

3. 실용적 개선 사항
현실적인 추론 환경에서 프레임워크를 효율적으로 만들기 위해, 저자들은 두 가지 핵심 기술을 제안한다:

• 혼합 온도 샘플링 (Mixed-Temperature Sampling): 제한된 수의 후보($K$)를 사용할 때 모드 붕괴(mode collapse)를 방로하기 위해, 프레임워크는 다양한 온도를 가진 템퍼드(tempered) PoE 분포 군집으로부터 후보를 샘플링한다. 이는 제안 엔트로피를 계층화하여 타겟 분포의 다양한 커버리지를 보장하고 유효 샘플 크기를 개선한다.
• 탄력적 거부 윈도우 (Elastic Rejection Window): 고정된, 잠재적으로 긴 접미사(suffix)를 검증하는 대신, 프레임워크는 관찰된 수락 길이에 따라 검증 범위를 적응적으로 조정한다. 이는 특히 여러 후보를 병렬로 처리할 때, 거부될 가능성이 높은 토큰에 대한 낭비되는 계산량을 줄여준다.

주요 기여

• 새로운 디코딩 프레임워크: PoE-Bridge는 어려운 교정 작업을 두 개의 관리 가능한 단계로 효과적으로 나누는 중간 PoE 브릿지를 도입함으로써, 병렬 DLM 디코딩과 AR 수준의 충실도를 성공적으로 결합한 최초의 방법이다.
• 계산 가능한 토큰 수준 PoE: 저자들은 시퀀스 수준의 PoE 샘플링의 난해함을 피하면서, 효율적인 추측적 검증과 중요도 가중치를 지원하는 토큰 수준의 PoE 보간을 정의했다.
• 효율성 메커니즘: 혼합 온도 샘플링과 탄력적 거부 윈도우의 도입은 다양성과 검증 오버헤드 측면의 실질적인 병목 현상을 해결한다.

실험 결과
도전적인 수학적 추론(GSM8K, MATH) 및 코딩(HumanEval, MBPP) 벤치마크에서 평가했을 때, PoE-Bridge는 다음과 같은 성과를 보여주었다:

• 정확도: 타겟 AR 모델 성능의 최소 **95%**를 회복하며, 표준 DLM 디코딩 및 나이브한 추측 샘플링(speculative sampling)을 유의미하게 능가한다.
• 속도: 표준 DLM 디코딩 방식 대비 최대 5배의 속도 향상을, AR 디코딩 대비 2배의 처리량을 달enc한다.
• 확장성: 혼합 온도 샘플링을 사용하면 정확도가 병렬 후보 수($K$)에 따라 일관되게 확장되는 반면, 균일 온도 샘플링은 수확 체감의 법칙을 보인다.
• 하이퍼파라미터 민감도: 수락률(처리량)과 AR 타겟에 대한 충실도 사이의 최적의 트레이드오프를 제공하는 중간 PoE 가중치($w \approx 0.3$)가 발견되었다.

의의 및 주장
본 논문은 PoE-Bridge가 기존의 품질-효율 트레이드오프를 깨뜨리고, DLM의 "최상의 조합(best-of-both-worlds)" 성능을 달성한다고 주장한다. 중간 브릿지 분포를 활용함으로써, 이 방법은 DLM이 지연 시간을 줄이는 병렬 디코딩의 이점을 유지하면서도 AR 디코딩만큼 충실한 생성을 할 수 있게 한다. 저자들은 이 접근 방식이 현실적인 제한된 컴퓨팅 예산 하에서도 몬테카를로 교정을 효과적으로 만들며, 확산 모델의 속도 이점을 희생하지 않고 고품질의 병렬 텍스트 생성을 위한 실용적인 솔루션을 제공한다고 단언한다. 본 연구는 모델을 재학습하거나 파라미터를 수정하지 않고, 빠른 DLM 제안과 강력한 AR 검증기 사이의 시너지를 활용하는 추론 시점(inference-time) 방법론으로 제시된다.