dLLM: 단순화된 디퓨전 언어 모델링 프레임워크

디퓨전 언어 모델(DLM) 연구는 파편화된 코드베이스로 인해 재현과 확장이 어려웠다. dLLM은 학습부터 평가까지 표준화된 파이프라인을 제공하며, 특히 기존 BERT나 GPT 모델을 적은 비용으로 DLM으로 변환할 수 있는 경로를 제시하여 연구 진입 장벽을 대폭 낮춘다.

기존 언어 모델은 토큰을 왼쪽에서 오른쪽으로 하나씩 생성하는 자기회귀(Autoregressive) 방식을 주로 사용하지만, 이는 생성 순서가 고정되어 수정이 어렵다는 한계가 있다. 디퓨전 언어 모델(DLM)은 문장 전체에 노이즈를 섞은 뒤 이를 단계적으로 복원하는 방식을 취한다. dLLM은 이 복잡한 과정을 '마스킹(Masking)'과 '복원(Denoising)'이라는 기초 개념으로 표준화한다.

입력 텍스트의 일부를 마스크 토큰으로 가리고, 모델이 주변 문맥을 참조하여 가려진 부분을 예측하도록 학습시킨다. 이때 dLLM은 학습 로직을 아키텍처와 분리하여, 개발자가 복잡한 수식 구현 없이도 기존의 Transformer 모델을 디퓨전 엔진으로 활용할 수 있게 돕는다. 이는 마치 완성된 퍼즐 조각 일부를 빼낸 뒤 다시 채워 넣는 과정을 반복하며 전체 그림을 완성하는 것과 같다.

결과적으로 연구자들은 dLLM을 통해 기존의 BERT나 GPT 모델을 가져와 디퓨전 방식으로 동작하게 만들 수 있다. 이는 단순히 생성 방식의 변화를 넘어, 병렬 디코딩을 통한 속도 향상과 더 정교한 문장 편집 기능을 기존 모델에 부여할 수 있음을 의미한다. 특히 ModernBERT와 같은 최신 인코더 모델을 활용하면 양방향 문맥 이해 능력을 극대화한 강력한 디퓨전 챗봇을 구축할 수 있다.

dLLM은 Trainer, Sampler, Evaluation의 세 가지 핵심 컴포넌트로 구성된다. Trainer는 Masked Diffusion(MDLM)과 Block Diffusion(BD3LM) 목적 함수를 지원한다. MDLM은 각 토큰을 독립적으로 마스킹하는 방식이며, BD3LM은 텍스트를 블록 단위로 나누어 블록 내에서는 디퓨전으로 병렬 생성하고 블록 간에는 자기회귀적으로 연결하는 하이브리드 방식이다.

Sampler는 Sampler(model).sample() 형태의 추상화 인터페이스를 제공한다. 이는 모델의 내부 구조를 건드리지 않고도 추론 알고리즘을 교체할 수 있게 한다. 예를 들어, 표준적인 반복 디코딩 대신 Fast-dLLM 알고리즘을 적용하면, 신뢰도가 높은 토큰들을 한 번에 확정하는 병렬 업데이트를 통해 추론 속도를 대폭 개선할 수 있다.

모델 변환 과정에서는 ModernBERT와 같은 인코더 모델이나 Qwen과 같은 자기회귀 모델을 활용한다. 인코더 모델의 경우 양방향 문맥 이해 능력을 활용해 마스크를 채우도록 SFT(Supervised Fine-Tuning)를 진행하며, 자기회귀 모델은 기존의 다음 토큰 예측 능력을 유지하면서도 전체 문맥을 참조해 마스크를 복원하도록 어댑테이션 과정을 거친다. 학습 시에는 마스킹된 위치의 음의 로그 가능도(Negative Log-Likelihood)를 계산하고, 노이즈 수준 t에 반비례하는 1/t 가중치를 곱해 손실 함수를 구성함으로써 다양한 노이즈 수준에서의 학습 균형을 맞춘다.

ModernBERT-large를 기반으로 한 BERT-Chat 모델은 GPT-2 Medium 대비 MMLU, MATH 등 대부분의 벤치마크에서 우수한 성능을 보였다. 특히 아키텍처 수정 없이 SFT만으로도 인코더 모델이 훌륭한 대화형 디퓨전 모델로 변환될 수 있음을 입증했다.

Qwen3-0.6B 모델을 BD3LM 방식으로 변환한 결과, 코드 생성 벤치마크인 HumanEval에서 46.3점을 기록하여 원본 베이스 모델(32.3점)보다 높은 성능을 나타냈다. 이는 디퓨전 방식의 반복적 정교화 과정이 논리적 추론이 필요한 작업에서 효과적일 수 있음을 시사한다.

추론 효율성 측면에서 Fast-dLLM을 적용했을 때, GSM8K 벤치마크 기준 공식 구현체와 유사한 정확도를 유지하면서도 기본 샘플러 대비 최대 8~11배의 토큰 처리 속도(Tok/s) 향상을 달성했다. 또한 터미널 시각화 도구를 통해 디퓨전 모델의 비순차적 토큰 생성 과정을 실시간으로 확인할 수 있게 하여 모델의 해석 가능성을 높였다.

dLLM의 아키텍처는 Hugging Face의 transformers 및 accelerate 생태계 위에 구축되어 FSDP, DeepSpeed와 같은 분산 학습 기술을 즉시 지원한다. Trainer 모듈은 transformers.Trainer를 경량 래핑하여 파라미터 효율적 미세 조정(PEFT)과 같은 최신 학습 기법과의 호환성을 유지한다.

수학적 기반으로 MDLM은 흡수 상태(Absorbing-state) 마스킹 과정을 따르며, 시간 t에 따른 선형 스케줄을 사용하여 노이즈 수준을 조절한다. 학습 목적 함수는 마스킹된 토큰의 음의 로그 가능도를 최소화하는 방향으로 설정되며, 노이즈 수준에 따른 기여도를 균형 있게 조절하기 위해 1/t 가중치를 적용한다. [입력값: 마스킹된 시퀀스 xt와 시간 t] → [연산: 모델 pθ가 예측한 원본 토큰 x0의 확률에 로그를 취하고 1/t를 곱함] → [출력: 가중치가 적용된 손실값] → [의미: 모델이 다양한 노이즈 수준에서 원본 데이터를 얼마나 정확히 복원하는지 나타냄]

BD3LM은 결합 확률을 블록 단위로 분해하여 pθ(x) = Π pθ(xBk | x<Bk)로 정의한다. 각 블록 Bk 내에서는 이전 블록의 깨끗한 이력을 조건부로 하여 디퓨전 프로세스가 진행되므로, KV 캐시를 활용한 효율적인 추론과 블록 내 병렬 생성이 동시에 가능하다. 이 방식은 자기회귀의 긴 문맥 처리 능력과 디퓨전의 병렬성을 결합한 형태이다.

평가 파이프라인은 lm-evaluation-harness를 확장하여 디퓨전 모델 특유의 추론 하이퍼파라미터(최대 토큰 수, 샘플링 온도 등)에 대한 민감도를 정확히 반영하도록 설계되었다. 이는 서로 다른 DLM 간의 공정한 성능 비교를 가능하게 하는 핵심 요소이며, 공식 벤치마크 결과를 높은 정확도로 재현할 수 있음을 확인했다.

DLM 평가는 추론 하이퍼파라미터(예: 최대 신규 토큰 수)에 매우 민감하며, 최적의 설정에서 조금만 벗어나도 성능이 급격히 저하되는 특성이 있다. 또한 변환된 모델들이 지식 및 추론 벤치마크에서 여전히 대규모 자기회귀 모델(ARLM)의 성능에는 미치지 못하는 격차가 존재한다.