마스크 기반 확산 언어 모델을 위한 어텐션 할인 적응 샘플러

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

마스크 기반 확산 언어 모델은 디노이징 반복당 여러 토큰을 공개하여 추론 단계를 줄일 수 있지만, 이러한 병렬성은 불안정합니다. 개별적으로 높은 신뢰도를 보이는 위치라도, 예측이 결합될 때 함께 선택하는 것은 위험할 수 있습니다. Top-k, Fast-dLLM 및 EB-Sampler와 같은 기존의 학습이 필요 없는 샘플러는 주로 공개할 토큰의 수를 제어하지만, 종종 선택된 집합 내의 상호 작용을 무시하는 토큰별 점수를 사용하여 후보를 순위화합니다. 우리는 병렬 마스크 기반 확산 디코딩을 위한 학습이 필요 없는 재순위 규칙인 ADAS를 제안합니다. ADAS는 기본 샘플러의 중단 규칙은 그대로 유지하고, 서브셋 구성 방식만 수정합니다. 즉, 이미 선택된 위치에 강하게 어텐션을 갖는 후보는 불확실한 예측을 가진 위치와 관련된 경우, 탐욕적으로 할인합니다. 그래프 기반 제약 방법과 달리, ADAS는 어텐션을 하드(Hard)한 호환성 제약으로 변환하는 대신, 어텐션을 연속적인 상태로 유지하고 소프트(Soft) 마진 페널티로 활용합니다. LLaDA-8B-Base 및 Dream-7B-Base 모델을 사용하여 GSM8K, MATH500, HumanEval 및 MBPP 데이터셋에서 Top-k, Fast-dLLM 및 EB-Sampler에 ADAS를 적용한 결과, 디노이저 평가 성능이 동일하게 유지되는 상태에서 평균적으로 각각 9.11% 및 10.46%의 낮은 NFE(Number of Forward Evaluations) 성능 향상을 보였습니다. 또한, 추론 시간 오버헤드는 약 3.1% 증가했습니다. 이러한 결과는 소프트 어텐션 할인을 통한 재순위화가 마스크 기반 확산 언어 모델의 고도로 병렬적인 디코딩에서 품질을 향상시키는 간단하고 모듈화된 방법임을 보여줍니다.

Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration, but this parallelism is fragile: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers such as Top-\(k\), Fast-dLLM, and EB-Sampler mainly control how many tokens to reveal, while often ranking candidates by token-wise scores that ignore interactions within the selected set. We propose ADAS, a training-free reranking rule for parallel masked diffusion decoding. ADAS leaves the base sampler's stopping rule unchanged and modifies only subset construction: it greedily discounts a candidate when it attends strongly to already selected positions whose predictions remain uncertain. Unlike graph-constrained methods that turn attention into hard compatibility constraints, ADAS keeps attention continuous and uses it as a soft marginal penalty. Across LLaDA-8B-Base and Dream-7B-Base on GSM8K, MATH500, HumanEval, and MBPP, plugging ADAS into Top-\(k\), Fast-dLLM, and EB-Sampler improves low-NFE performance at matched denoiser evaluations by \(9.11\) and \(10.46\) percentage points on average, respectively, with \(3.1\%\) per-forward runtime overhead. These results show that soft attention-discounted reranking is a simple and modular way to improve quality in highly parallel decoding for masked diffusion language models.