기술 요약: 멀티 블록 확산 언어 모델 (MBD-LMs)

1. 문제 정의

확산 언어 모델(Diffusion Language Models, DLMs)은 자기회귀적 생성의 대안으로서 네이티브 병렬 디코딩을 제공하지만, KV 캐싱 및 동적 길이 생성 지원의 부재로 인해 역사적으로 효율적인 서빙에 어려움을 겪어왔다. 최근의 블록 확산 언어 모델(Block Diffusion Language Models, BD-LMs)은 블록 인과적(block-causal) 생성을 도입하여 Single-Block Diffusion (SingleBD)을 가능하게 함으로써 이 문제를 해결했다. SingleBD에서 모델은 깨끗한(clean) 접두사(prefix)를 조건으로 하나의 노이즈가 섞인 블록을 디코딩하며, 이를 통해 KV 캐싱과 블록 내 병렬성을 허용한다. 그러나 블록들은 여전히 순차적으로 처리되며, 이 과정에서 시스템이 다음 블록을 디코딩하기 전에 완료된 블록을 캐싱하기 위해 대기하는 "스토어링 버블(storing bubbles)"이 발생하여 블록 간 병렬성을 제한한다.

Discrete Diffusion Forcing (D2F)는 다중 노이즈 블록 간의 가시성(visibility)을 도입하여 Multi-Block Diffusion (MultiBD) 및 블록 간 병렬성을 가능하게 했으나, 심각한 **훈련-추론 불일치(train–inference mismatch)**가 남아 있다. D2F는 길고 단조로운(monotonic) 노이즈 접미사(suffix) 상에서 훈련되는 반면, 실제 MultiBD 추론은 경계가 정해진 러닝 셋(bounded running-set) 내의 연속된 블록들과 이질적인 슬롯별 노이즈 패턴(heterogeneous slot-wise noise patterns) 상에서 동작한다. Teacher Forcing (TF) 또는 D2F로 훈련된 기존 BD-LM들은 이러한 특정 추론 상태와 일치하지 않으며, 이는 여러 블록을 동시에 디코딩하려고 할 때 최적이 아닌 성능을 초래한다. 또한, 나이브한 MultiBD 구현은 종종 동적 입력 형상(dynamic input shapes)에 의존하는데, 이는 CUDA Graph 캡처와 같은 정적 형상 실행 메커니즘과 호환되지 않아 실질적인 처리량 이득을 저해한다.

2. 방법론

본 논문은 기존 BD-LM과 효율적인 MultiBD 추론 사이의 간극을 메우기 위한 통합 프레임워크이자 훈련 전략인 **멀티 블록 확산 언어 모델 (MBD-LMs)**을 제안한다.

2.1 통합 정식화: MBD-LMs

저자들은 MBD-LM을 연속된 블록의 **러닝 셋($R_s$)**이 동시에 디코딩되는 BD-LM의 일반화된 형태로 정의한다. $|R_s|=1$인 SingleBD나 전체 시퀀스 D2F와 달리, 실제적인 MBD-LM은 병렬성과 계산 효율성의 균형을 맞추기 위해 $|R_s|$가 작은(예: 2~4개 블록) 유한한 중간 단계에서 동작한다.

2.2 멀티 블록 티처 포싱 (Multi-block Teacher Forcing, MultiTF)

훈련 상태를 경계가 정해진 러닝 셋 추론 체제와 일치시키기 위해, 저자들은 기존 BD-LM을 위한 사후 훈련(post-training) 방법인 **Multi-block Teacher Forcing (MultiTF)**을 도입한다.

• 노이즈 그룹 구성: 단일 블록을 오염시키는 방식(TF)이나 긴 단조적 접미사를 사용하는 방식(D2F) 대신, MultiTF는 깨끗한 접두사를 조건으로 하는 경계가 있는 연속적인 노이즈 블록 그룹인 **노이즈 그룹(noise-groups)**을 구성한다.
• 그룹 레이아웃: 훈련 데이터는 체계적 레이아웃(최대 $G_{max}$까지의 모든 이동 및 그룹 크기를 열거)과 무작위 레이아웃을 모두 사용하여 파티션되며, 이를 통해 러닝 셋 크기와 블록 상대적 위치에 대한 다양한 커버리지를 보장한다.
• 체인-유니폼 노이즈 스케줄러 (Chain-Uniform Noise Scheduler): 각 노이즈 그룹 내에서 MultiTF는 체인-유니폼 노이즈 스케줄러를 사용한다. D2F의 단조적 스케줄과 달리, 이 방법은 그룹 레벨의 하한값(floor)을 샘플링한 후, 이전 블록의 비율에 따라 각 후속 블록의 노이즈 비율을 샘플링한다. 이는 MultiBD 추론 중에 관찰되는 경험적 노이즈 패턴과 더 잘 일치하는 **이질적인 슬롯별 노이즈 갭(heterogeneous slot-wise noise gaps)**을 생성한다.
• 그룹 인지 듀얼 스트림 마스크 (Group-Aware Dual-Stream Mask): 특수화된 어텐션 마스크는 그룹 내의 노이즈 블록들이 서로를 참조할 수 있게 하면서도, 깨끗한 접두사에 대한 블록 인과적 가시성을 유지하여 미래의 노이즈 토큰으로부터 깨끗한 접두사로 정보가 누출되는 것을 방지한다.

2.3 최적화된 추론 엔진: 블록 버퍼 (Block Buffer)

실용적이고 높은 처리량의 MultiBD 추론을 가능하게 하기 위해, 저자들은 블록 버퍼(Block Buffer) 메커니즘에 기반한 최적화된 파이프라인을 설계한다.

• 정적 형상 실행 (Static-Shape Execution): 입력을 동적으로 추가하는 방식(이는 CUDA Graph 호환성을 깨뜨림) 대신, 엔진은 고정된 크기의 물리적 버퍼($N_{buf}$)를 유지한다.
• 상태 전이: 버퍼 내의 블록은 dummy(비활성), active(디코딩 중), to-cache(완료되었으나 아직 캐싱되지 않음), in-cache(KV 캐시에 커밋됨)의 네 가지 상태를 거친다.
• 메커니즘: 미래의 블록들은 입력 시퀀스를 확장하는 대신 기존의 dummy 슬롯을 활성화함으로써 버퍼에 진입한다. 완료된 블록은 KV 캐시에 커밋되고, 꼬리 부분의 새로운 dummy 슬로트로 대체된다. 이는 효율적인 GPU 실행에 필요한 정적 입력 형상을 유지하면서도, 블록 간 병렬성을 위한 논리적 러닝 셋을 유지한다.
• 접두사 캐시 보존: 이 설계는 불변의 캐시된 접두사와 활성 버퍼를 엄격히 분리하여, 네이티브 D2F가 위반하는 경우가 많은 표준 BD-LM 접두사 캐싱 의미론과의 호환성을 보장한다.

3. 주요 기여

1. 통합 MBD-LM 정식화: 본 논문은 MBD-LM을 경계가 정해진 러닝 셋으로 매개변수화된 통합 프레임워크로 정식화하여, 실질적인 MultiBD가 SingleBD와 전체 시퀀스 D2F 사이의 중간 단계임을 식별하였다.
2. MultiTF 사후 훈련: 이질적인 노이즈 패턴을 가진 경계가 있는 노이즈 그룹을 통해 훈련함으로써 BD-LM을 MultiBD 추론과 정렬하는 새로운 훈련 전략을 제시하여, 훈련-추론 불일치를 크게 줄였다.
3. 최적화된 추론 엔진: CUDA Graph를 위한 정적 형상 실행을 가능하게 하고, 접두사 KV 캐싱을 보존하며, 디코딩과 캐시 저장을 오버랩하여 벽시계 시간(wall-clock time) 속도 향상을 달성하는 블록 버퍼 기반 파이프라인을 구축하였다.

4. 실험 결과

저자들은 LLaDA2 및 SDAR 모델 제품군을 사용하여 수학적 추론(GSM8K, MATH500) 및 코드 생성(MBPP+, HumanEval+) 벤치마크에서 MBD-LM을 평가하였다.

• 병렬성 및 정확도: MBD-LLaDA2-Mini는 평균 **전방 패스당 토큰 수(Tokens Per Forward pass, TPF)**를 3.47(SingleBD)에서 6.19(+78.4%)로 증가시켰으며, 동시에 평균 정확도를 79.95%에서 81.03%로 개선하였다.
• DMax와의 결래: DMax 디코딩 전략과 결합했을 때, MBD-LLaDA2-Mini-DMax는 평균 TPF 9.34를 달달성하였으며(SingleBD DMax 대비 +47.1% 개선), 정확도 저하는 1.02%에 불과했다.
• 처리량 (TPS): 최적화된 추론 엔진은 TPF의 이득을 벽시계 시간 처리량으로 전환했다. MBD-LLaDA2-Mini-DMax는 평균 926.67 TPS를 달성하였으며, 이는 SingleBD DMax 베이스라인의 779.49 TPS와 비교된다.
• 절제 연구 (Ablation Studies):
  • 훈련-추론 정렬: 훈련되지 않은 모델에 MultiBD 추론을 적용한 훈련 없는(Training-free) MultiBD는 TPF는 개선했으나 정확도는 저하되었다. MultiTF 사후 훈련은 정확도를 회복하고 오히려 향상시켰으며, 이는 상태 정렬의 필요성을 확인시켜 주었다.
  • 스케줄러 설계: 제안된 체인-유니폼 스케줄러는 D2F 스타일의 단조적 스케줄러 및 기타 유니폼 변형보다 우수한 성능을 보였으며, 이는 이질적인 노이즈 갭의 중요성을 강조한다.
  • 레이아웃 다양성: 체계적 레이아웃과 무작위 레이아웃을 결합했을 때 TPF와 정확도 사이의 최적의 균형을 얻을 수 있었다.

5. 의의

본 논문은 신뢰할 수 있고 효율적인 멀티 블록 확산(Multi-Block Diffusion)을 위해서는 **훈련 시점의 상태 정렬(state alignment)**과 추론 시점의 시스템 지원을 결합한 총체적인 접근 방식이 필요하다고 주장한다.

• 단순히 기존 모델에 MultiBD 추론을 적용하는 것만으로는 노이즈 패턴과 러닝 셋 크기의 분포 불일치 때문에 불충분하다는 것을 입증하였다.
• MBD-LM이 BD-LM의 효율성을 만드는 핵심인 KV 캐싱과 접두사 재사용의 이점을 희생하지 않으면서도 블록 간 병렬성을 달성할 수 있음을 증명하였다.
• 본 연구는 "경계가 정해진 러닝 셋"이 실용적인 최적 지점임을 확립하였으며, 이는 SingleBD의 순차적 병목 현상을 넘어, 전체 시퀀스 확산의 복잡성과 캐시 불호환성을 피하면서도 확산 기반 생성의 속도를 높일 수 있는 실행 가능한 경로를 제시한다.