검색 또는 가속화: 확산 언어 모델을 위한 신뢰도 기반 위치 빔 검색 (Confidence-Switched Position Beam Search)

Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language Models

확산 언어 모델(DLM)은 마스크된 시퀀스를 반복적으로 노이즈 제거하면서 텍스트를 생성하며, 각 단계마다 어떤 위치를 우선적으로 노출할지 결정합니다. 일반적인 디코딩 방식은 가장 신뢰도가 높은 위치를 먼저 노출하는 탐욕적인 방법을 사용하지만, 이러한 지역적인 선택은 모델을 최적의 노출 순서에서 벗어나게 만들 수 있으며, 특히 추론이 필요한 프롬프트에서 이러한 문제가 더욱 두드러집니다. 본 논문에서는 모델의 불확실성에 따라 동작을 조정하는 훈련이 필요 없는 디코딩 알고리즘인 SOAR을 제시합니다. SOAR은 신뢰도가 낮을 때, 조기에 확정하는 것을 방지하기 위해 대체 노출 결정에 대한 검색 범위를 일시적으로 넓히고, 신뢰도가 높을 때에는 검색 범위를 좁히고 여러 위치를 동시에 디코딩하여 노이즈 제거 반복 횟수를 줄입니다. Dream-7B 및 LLaDA-8B 모델에서 GSM8K, MBPP, HumanEval과 같은 수학적 추론 및 코드 생성 벤치마크를 통해 SOAR은 생성 품질을 향상시키면서도 경쟁력 있는 추론 속도를 유지하며, DLM 디코딩에서 품질과 효율성의 균형을 맞추는 실용적인 방법을 제공합니다. 저희의 코드는 https://github.com/duterscmy/SOAR 에서 확인하실 수 있습니다.

Diffusion Language Models (DLMs) generate text by iteratively denoising a masked sequence, repeatedly deciding which positions to commit at each step. Standard decoding follows a greedy rule: unmask the most confident positions, yet this local choice can lock the model into a suboptimal unmasking order, especially on reasoning-heavy prompts. We present SOAR, a training-free decoding algorithm that adapts its behavior to the model's uncertainty. When confidence is low, SOAR briefly widens the search over alternative unmasking decisions to avoid premature commitments; when confidence is high, it collapses the search and decodes many positions in parallel to reduce the number of denoising iterations. Across mathematical reasoning and code generation benchmarks (GSM8K, MBPP, HumanEval) on Dream-7B and LLaDA-8B, SOAR improves generation quality while maintaining competitive inference speed, offering a practical way to balance quality and efficiency in DLM decoding. Our Code is available at https://github.com/duterscmy/SOAR