FAST-DLLM V2 Efficient Block-Diffusion LLM

대규모 언어 모델(LLM)은 자연어 처리 분야에서 놀라운 성능을 보여주고 있지만, 기존의 자기회귀(Autoregressive, AR) 방식은 토큰을 순차적으로 하나씩 생성하기 때문에 추론 속도에 제약이 있습니다. 이러한 한계를 극복하기 위해 확산 기반 언어 모델(Diffusion Language Model, dLLM)이 제안되었지만, 여전히 실용적인 배포에는 어려움이 있었습니다.

이번 논문에서 소개하는 Fast-dLLM v2는 사전 학습된 AR 모델을 블록 확산 모델로 효율적으로 변환하여 병렬 텍스트 생성을 가능하게 합니다. 특히 단 1B 토큰만으로 파인튜닝이 가능하다는 점에서, 기존 Dream 모델이 580B 토큰을 필요로 했던 것과 비교하면 500배나 적은 학습 데이터로 동일한 성능을 달성할 수 있습니다.

C. Wu, H. Zhang, S. Xue, S. Diao, Y. Fu, Z. Liu, P. Molchanov, P. Luo, S. Han and E. Xie, "Fast-dLLM v2: Efficient Block-Diffusion LLM", arXiv preprint arXiv:2509.26328, 2025.

들어가는 글

요약

아키텍처: 블록 확산 언어 모델(Block Diffusion LLM) 구조를 기반으로 하며, Qwen 2.5-Instruct 모델(1.5B, 7B)을 파인튜닝하여 구현됩니다.

핵심 메커니즘:

• 블록 단위 확산(Block-wise Diffusion): 시퀀스를 블록으로 나누고, 각 블록 내에서는 확산 방식으로, 블록 간에는 자기회귀 방식으로 생성
• 상호 보완 마스킹(Complementary Masking): 모든 토큰이 학습될 수 있도록 보완적인 마스크 전략 사용
• 계층적 캐싱(Hierarchical Caching): 블록 레벨 캐시와 서브 블록 캐시를 활용한 효율적인 추론

학습 방법:

• LLaMA-Nemotron 후처리 데이터셋 사용
• 1.5B 모델: 학습률 \(2 \times 10^{-5}\), 6,000 스텝 (약 3.15B 토큰)
• 7B 모델: 학습률 \(1 \times 10^{-5}\), 2,500 스텝 (약 1.31B 토큰)
• 64개의 NVIDIA A100 GPU 사용

데이터셋: LLaMA-Nemotron 후처리 데이터셋, 배치 크기 256, 컨텍스트 길이 2048

평가 메트릭: HumanEval, MBPP(코드 생성), GSM8K, MATH(수학 추론), MMLU, GPQA(지식 집약적 QA), IFEval(명령 수행) 등 다양한 벤치마크

성능:

• 표준 AR 디코딩 대비 최대 2.5배 속도 향상
• 7B 모델에서 평균 60.3점 달성, Qwen2.5-7B-Nemo-FT(59.6점), Dream(57.6점) 등 모든 베이스라인 능가
• GSM8K에서 임계값 0.9 설정 시 정확도 소폭 감소만으로 2.6배 속도 향상 달성

논문 상세

배경 및 동기

기존 AR 언어 모델은 토큰을 순차적으로 생성하기 때문에 디코딩 과정에서 완전한 병렬 처리를 활용할 수 없습니다. 반면 확산 기반 언어 모델(dLLM)은 여러 토큰을 동시에 예측하거나 개선할 수 있어 높은 디코딩 병렬성을 제공할 수 있습니다.

하지만 기존 dLLM들은 다음과 같은 한계가 있었습니다:

• 양방향 어텐션으로 인해 KV 캐시를 효과적으로 사용할 수 없음
• 추론 지연 시간이 AR 모델보다 길 수 있음
• 고정된 시퀀스 길이가 필요하거나 생성 길이에 제약이 있음

블록 확산 언어 모델로의 적응

Fast-dLLM v2는 사전 학습된 Qwen2.5-Instruct 모델을 블록 확산 프레임워크로 변환합니다. 핵심 아이디어는 시퀀스를 블록 단위로 나누고, 각 블록 내에서는 확산 방식으로 토큰을 생성하되, 블록 간에는 자기회귀 방식을 유지하는 것입니다.

블록 단위 구성: 각 시퀀스를 블록 크기 \(D\)의 배수로 패딩하고, 컨텍스트 길이 \(L\)로 패킹합니다. 이를 통해 \(B = L/D\)개의 겹치지 않는 블록으로 자연스럽게 분할됩니다.

상호 보완 마스킹: 각 블록에 대해 이진 마스크 \(m \in {0, 1}^D\)를 샘플링하고, 모든 토큰이 학습되도록 보완 마스크 \(\bar{m} = 1 - m\)도 함께 사용합니다.

토큰 시프트 예측: AR 모델의 표현 품질을 유지하기 위해, 마스킹된 위치 \(i\)의 토큰을 예측할 때 이전 위치 \(i-1\)의 로짓을 사용합니다.

학습 목표: 마스킹된 토큰에 대해서만 교차 엔트로피 손실을 계산합니다:

\[\mathcal{L}_{block}(\theta) = -\mathbb{E}_{x,m}\left[\sum_{i=1}^{L} \mathbb{1}[x_i^t = \text{[MASK]}] \log p_\theta(x_i^0 | x_{<i}, x_{block(i)})\right]\]

여기서 \(x_{block(i)}\)는 위치 \(i\)를 포함하는 블록의 모든 토큰을 의미합니다.

추론 파이프라인

Fast-dLLM v2는 추론 시 다음과 같은 전략을 사용합니다:

블록 단위 자기회귀 디코딩: 각 블록이 인과적 순서로 디코딩되므로, 디코딩된 블록은 이후 블록의 읽기 전용 컨텍스트로 캐싱됩니다. 이를 통해 블록 레벨 KV 캐시 재사용이 가능해집니다.

블록 내 병렬 개선: Fast-dLLM의 신뢰도 기반 병렬 디코딩 전략을 채택하여, 모델 신뢰도가 높은 토큰들을 병렬로 디코딩하고, 불확실한 위치는 추가 개선을 위해 마스킹 상태로 유지합니다.

DualCache: 부분적으로 디코딩된 블록에 대해 prefix와 suffix KV 캐시를 모두 유지하여, 추가 토큰이 드러날 때 효율적인 재계산을 지원합니다.

배치 디코딩: 다양한 길이의 시퀀스를 배치로 생성하기 위해, 각 시퀀스를 블록 크기의 배수가 되도록 [MASK] 토큰으로 패딩합니다.

실험 결과

성능 비교: 7B 모델 기준으로 Fast-dLLM v2는 평균 60.3점을 달성하여, Qwen2.5-7B-Nemo-FT(59.6점)와 Dream(57.6점)을 포함한 모든 베이스라인을 능가했습니다. 1.5B 모델에서도 45.0점으로 동일 규모의 AR 및 확산 기반 모델 중 최고 성능을 기록했습니다.

속도 향상: GSM8K에서 임계값 0.9로 설정했을 때, 정확도는 소폭만 감소하면서 처리량이 39.1에서 101.7 토큰/초로 증가하여 2.6배 속도 향상을 달성했습니다.

하드웨어 확장성: A100과 H100 GPU 모두에서 배치 크기가 증가할수록 AR 베이스라인 대비 우수한 처리량을 보였습니다. H100에서는 최대 1.8배 속도 향상을 달성했습니다.

분석 실험

상호 보완 마스킹의 효과: 패딩 전략과 상호 보완 마스킹을 모두 적용했을 때 평균 정확도가 기본 전략 대비 +3.7점 향상되었습니다.

서브 블록 크기: 서브 블록 크기 8이 대부분의 작업에서 최적 성능을 보였으며, 작업에 따라 최적 크기가 다를 수 있음을 확인했습니다.

블록 크기 일치: 학습 시와 추론 시의 블록 크기가 일치하지 않으면 성능이 크게 저하됩니다. 서브 블록 디코딩 전략을 도입하여 이러한 일관성을 유지하면서도 추론 세밀도를 유연하게 조절할 수 있습니다.

캐싱 효과: 서브 블록 캐시는 배치 크기가 작을 때는 미미한 효과를 보이지만, 배치 크기가 큰 계산 집약적 환경에서는 상당한 속도 향상을 제공합니다. 중요한 점은 캐싱이 모델 정확도에는 전혀 영향을 미치지 않는다는 것입니다.

결론

Fast-dLLM v2는 사전 학습된 AR 모델을 효율적인 블록 확산 프레임워크로 변환하여, 단 1B 토큰의 파인튜닝만으로 표준 AR 디코딩 대비 최대 2.5배 속도 향상을 달성했습니다. 블록 단위 확산 메커니즘과 상호 보완 마스킹, 계층적 캐싱 전략을 통해 생성 품질을 유지하면서도 추론 효율성을 크게 개선했습니다.

광범위한 벤치마크 실험을 통해 Fast-dLLM v2가 AR 베이스라인과 동등하거나 더 나은 정확도를 유지하면서도, 기존 dLLM들 중 최고 수준의 효율성을 달성함을 입증했습니다. 이는 고품질, 저지연 LLM의 실용적 배포를 향한 중요한 진전을 의미합니다.