다중 블록 확산 언어 모델(Multi-Block Diffusion Language Models, MBD-LMs)

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

다중 블록 확산 추론(MultiBD)은 블록 수준의 병렬화를 통해 텍스트 생성의 실효 처리량을 높일 수 있다. 기존 BD-LM들은 Teacher Forcing이나 D2F 같은 학습 상태와 실제 MultiBD 추론 상태 간 불일치 때문에 병렬화 이득을 실 서비스 수준의 속도로 전환하지 못했다. 본 논문은 학습 상태를 MultiBD 추론과 정합시키는 MultiTF와 추론 엔진 측의 Block Buffer를 결합해 토큰당 유용한 성능 개선과 실제 TPS 향상을 동시에 달성했다.

왜 중요한가

다중 블록 확산 추론(MultiBD)은 블록 수준의 병렬화를 통해 텍스트 생성의 실효 처리량을 높일 수 있다. 기존 BD-LM들은 Teacher Forcing이나 D2F 같은 학습 상태와 실제 MultiBD 추론 상태 간 불일치 때문에 병렬화 이득을 실 서비스 수준의 속도로 전환하지 못했다. 본 논문은 학습 상태를 MultiBD 추론과 정합시키는 MultiTF와 추론 엔진 측의 Block Buffer를 결합해 토큰당 유용한 성능 개선과 실제 TPS 향상을 동시에 달성했다.

핵심 기여

MBD-LM으로의 통합적 정식화

Multi-Block Diffusion Language Models(MBD-LMs)를 러닝셋(running-set) 크기라는 관점에서 통합적으로 정식화해 Teacher Forcing과 D2F를 특수한 극한 사례로 위치시켰고, 유용한 MultiBD를 유한 크기의 러닝셋으로 규정해 학습·추론 요구사항을 명확히 했다.

Multi-block Teacher Forcing (MultiTF) 포스트트레이닝

MultiTF는 시스템적으로 생성된 그룹 레이아웃과 무작위 레이아웃을 결합하고 각 노이즈 그룹 내에 chain-uniform 노이즈 스케줄러와 Group-Aware Dual-Stream Mask를 적용해 BD-LM을 MultiBD 유사 상태로 포스트트레이닝하여 학습·추론 불일치를 줄였다.

Block Buffer 기반의 정적-입력 추론 엔진 설계

Block Buffer 메커니즘은 고정된 슬롯 수 N_buf를 유지해 CUDA Graph 캡처·재생을 가능하게 하며, 슬롯 상태(dummy/active/to-cache/in-cache) 전이를 통해 KV 캐시 재사용과 디코딩-스토어 중첩을 보존해 TPF 증분을 실제 TPS 개선으로 전환했다.

실험적 검증과 구성요소 절단검사

여러 BD-LM 백본과 수학·코드 벤치마크에서 MultiTF 적용 모델이 TPF·AUP를 크게 향상시키는 것을 확인했으며, chain-uniform 스케줄러 및 체계적+무작위 그룹 레이아웃의 조합이 최적 결과를 낳는 근거를 ablation으로 제시했다.

핵심 아이디어 이해하기

블록 확산(Block Diffusion) 접근은 시퀀스를 블록으로 나누고 각 블록을 블록-인과적(block-causal)하게 모델링해 KV 캐싱을 가능하게 한다. 기존 Teacher Forcing 기반 학습은 한 번에 하나의 노이즈 블록만 관찰하도록 학습시키므로 여러 블록을 함께 디코딩하는 MultiBD 추론과 학습 상태가 일치하지 않아 품질 저하가 발생했다. 이 불일치가 병렬화 이득을 실제 처리량으로 전환하는 데 방해 요소로 작용했다.

방법론

MultiTF는 학습 측면에서 러닝셋의 유한 크기와 슬롯별 이질적 노이즈 패턴을 모사하는 데 초점을 맞춘다. 구체적으로 입력 시퀀스를 블록 단위로 분할한 뒤 G_max로 정의된 최대 노이즈 그룹 크기 범위에서 체계적(systematic)과 무작위(random) 레이아웃을 생성해 다양한 그룹 상대 위치와 경계 패턴을 커버했다. 각 노이즈 그룹 내부에서는 chain-uniform 스케줄러를 사용해 그룹 레벨 floor ℓ를 샘플링하고 이후 각 블록의 마스크 비율 t_j를 ℓ에서 t_eff 사이에서 누적적으로 샘플링해 슬롯 간에 이질적이지만 단조적인 노이즈 분포를 생성했다.

주요 결과

메인 결과에서 MBD-LLaDA2-Mini는 평균 Tokens Per Forward pass(TPF)를 3.47에서 6.19로 증가시키며 평균 정확도를 79.95%에서 81.03%로 향상시켰다. DMax와 결합된 MBD-LLaDA2-Mini-DMax는 평균 TPF 9.34를 달성하면서 수학·코드 벤치마크에서 평균 정확도 손실은 1.02퍼센트포인트에 그쳤다. 추론 엔진 측면에서는 MBD-LLaDA2-Mini-DMax가 평균 TPS 951.41을 기록해 LLaDA2-Mini-DMax의 781.50 대비 실질적 처리량 증가를 보였다.

기술 상세

전체 아키텍처 관점에서 MBD-LM은 러닝셋 ℛ_s라는 연속된 블록 집합을 동시 복원 대상으로 정의하고, 각 러닝셋 내부에서 블록별 마스크 비율 t_{k,s}을 가진 상태를 조건으로 블록 단위 확률 분해를 수행한다. MultiTF는 학습 데이터 생성 단계에서 입력을 [noisy_part ; clean_part]로 쌓아 Group-Aware Dual-Stream Mask A_lambda를 만들며, 상단의 M_GD는 그룹 내부의 노이즈 블록 가시성을, 우상단의 M_GOC는 노이즈 그룹이 깨끗한 접두사에 조건하는 가시성을, 하단의 M_BC는 깨끗한 부분에 대한 표준 블록-인과적 가시성을 각각 제공한다. 추론 엔진은 Block Buffer로 고정된 물리적 버퍼 크기 N_buf를 유지하며 각 슬롯을 dummy→active→to-cache→in-cache로 전이시켜 새로운 블록은 기존의 dummy 슬롯을 활성화함으로써 입력 형태를 정적 상태로 유지하고 KV 캐시 재사용을 보장한다.

한계점

논문이 명시한 한계로는 정적 Block Buffer를 사용함에 따라 물리적으로 처리되는 토큰 수 Q_s가 증가하고, 이로 인해 forward step의 비용이 늘어나 토큰 효율 η_tok이 감소할 수 있다는 점이 있다. Block Buffer 내의 dummy 슬롯과 이미 거주 중인 블록들은 즉시 커밋되는 유용 토큰이 아니므로 이상적인 토큰 효율보다 낮아질 여지가 있으며, 최종 TPS 증가는 TPF 증가 대비 per-step 비용 증가의 균형에 따라 제한된다. 또한 MultiTF는 포스트트레이닝 단계가 필요하며 G_max 및 스케줄러·레이아웃 하이퍼파라미터 선택이 성능과 품질에 민감하게 작용한다.

실무 활용

MultiTF와 Block Buffer 조합은 기존 BD-LM을 사후 학습(post-training)만으로 MultiBD 친화적으로 전환해 디코딩 병렬성을 늘리고 실제 TPS까지 개선할 수 있게 한다. 이는 서비스 환경에서 KV 캐시 재사용과 정적 입력 형태(CUDA Graph 캡처)를 유지하면서 병렬 디코딩을 적용하려는 사례에 적합하다. GitHub에 실험 코드와 도구가 공개되어 있어 실무 적용을 위한 재현성이 확보되어 있다.

대규모 배치 또는 저지연 서버 환경에서 블록 기반 텍스트 생성량과 처리량을 높여 응답 처리량(TPS)을 개선하는 용도
수학·코드 생성처럼 품질 유지가 중요한 생성 작업에서 디코딩 병렬성을 늘려 처리시간을 단축하는 용도
기존 BD-LM에 사후 학습으로 MultiBD 추론을 도입해 추론 파이프라인을 점진적으로 개선하는 마이그레이션 시나리오

코드 공개 여부: 공개

코드 저장소 보기

키워드

Multi-Block Diffusion(다중 블록 확산)Block Buffer(블록 버퍼)MultiTF(다중 블록 Teacher Forcing)chain-uniform scheduler(체인-유니폼 노이즈 스케줄러)Tokens Per Forward (TPF)(순전달당 토큰 수)

용어 해설

Block Diffusion: — 문장을 고정 길이 블록 단위로 분할하여 각 블록을 블록-인과적(block-causal) 방식으로 순차적으로 생성하는 확산 기반 텍스트 생성 기법으로, 각 블록은 이전 블록의 캐시된 KV 상태를 조건으로 하여 병렬화와 캐싱을 동시에 지원한다.
Teacher Forcing: — 학습 시 현재 목표 블록만 노이즈화한 뒤 깨끗한 접두사(clean prefix)를 입력으로 주어 모델이 한 블록을 복원하도록 학습시키는 방법으로, 블록 캐시 사용에는 적합하지만 다중 블록 동시 복원(MultiBD) 상태와는 불일치가 발생한다.
Discrete Diffusion Forcing (D2F): — 접두사 조건 하에서 접미 블록들에 증가하는 마스크 비율을 부여해 여러 노이즈화된 블록을 가시화하는 학습 스케줄로, 다중 블록 가시성은 제공하나 실제 MultiBD 추론의 유한한 러닝셋 크기와 슬롯별 이질적 노이즈 분포와는 불일치가 존재한다.
Group-Aware Dual-Stream Mask: — 노이즈화된 블록 그룹과 깨끗한 전체 시퀀스를 연결한 입력에 대해, 그룹 내부의 노이즈 블록들끼리 상호 가시성을 허용하고 각 그룹이 자체 깨끗한 접두사에 조건하도록 행렬 형태의 블록-인과적 마스크를 구성하는 메커니즘이다.
Tokens Per Forward (TPF): — 한 번의 forward 연산에서 커밋되는(확정되는) 유용 토큰 수를 측정하는 지표로, 디코딩 병렬성의 실질적 이득을 정량화하는 데 사용되며 Block Buffer 크기와 연관되어 변화한다.

코드 예제

text

Algorithm 1 Multi-block Teacher Forcing
1: Clean sequence x0 ; block size B ; maximum noise-group size G_max ; noise bounds t_low, t_high ; margin ratio rho ; number of random layouts N_rand ; mask token [M].
2: Partition x0 into K blocks [b1,...,bK].
3: Generate systematic layouts by enumerating group sizes g in {2,...,G_max} and shifts.
4: Generate N_rand random layouts by sampling group sizes from {2,...,G_max} until all blocks are covered.
5: For each layout lambda in Lambda do
6:   Initialize x_t^lambda <- x0.
7:   For each noise-group H_m = (j1,...,j_nm) in lambda do
8:     Sample group floor l ~ U(t_low, t_eff).
9:     For i = 1 to n_m do
10:      Sample t_{j_i} ~ U(l, t_eff) and set l <- t_{j_i}.
11:      Mask floor(B * t_{j_i}) random positions in b_{j_i} as [M].
12:    end for
13:   end for
14:   Build X_lambda = [x_t^lambda ; x0] and construct Group-Aware Dual-Stream Mask A_lambda.
15:   Run model on (X_lambda, A_lambda) and compute masked CE over masked positions M_lambda.
16: end for
17: Return average loss J/|Lambda|.

이 코드는 MultiTF의 핵심 루프를 의사코드 형태로 보여주며, 입력 시퀀스를 블록 단위로 분할한 뒤 체계적·무작위 그룹 레이아웃을 구성하고 각 노이즈 그룹 내부에서 chain-uniform 방식으로 블록별 마스크 비율을 샘플링해 마스크된 토큰에 대해 masked CE를 계산하는 과정을 포함한다. 이 절차는 post-training 단계에서 BD-LM을 MultiBD 추론 상태와 유사한 노이즈-그룹 상태로 노출시켜 train–inference 정합을 향상시키는 목적을 가진다. 알고리즘의 주요 하이퍼파라미터로는 G_max, t_low, t_high, rho, N_rand가 있으며 이 값들이 노이즈 그룹 크기 분포와 슬롯별 노이즈 이질성을 결정한다.