확산 언어 모델 구축: 아키텍처, 샘플링, 확장성

실행 요약

확산 언어 모델(dLLMs)은 텍스트 생성에 대한 표준 순차적 생성(AR) 접근 방식에서 패러다임 전환을 나타냅니다. AR 모델은 왼쪽에서 오른쪽으로 토큰을 순차적으로 생성하는 반면, 확산 모델은 전체 시퀀스를 병렬로 생성하고, 반복적으로 원시 초안을 정교한 출력으로 개선합니다. 이 접근 방식은 AR 모델의 세 가지 핵심 한계를 해결합니다: 오류 수정의 부재, 느린 순차적 생성, 그리고 인과적(과거만 참조하는) 어텐션에 대한 제약입니다.

이산 확산의 핵심 메커니즘

텍스트와 같은 이산 데이터에 확산을 적용하려면 '노이즈'의 정의를 재정의해야 합니다. 이미지 확산은 가우시안 노이즈를 사용하지만, 언어 확산은 일반적으로 마스킹 또는 상태 교체 전략을 사용합니다.

마스킹 확산 모델(MDLM)

마스킹 확산은 "생성형 BERT"와 같습니다. 이 과정은 두 단계로 구성됩니다:

전진 과정: 깨끗한 시퀀스가 특정 스케줄에 따라 토큰을 [MASK] 토큰으로 점차 교체하면서 점차 손상됩니다. 결국 시퀀스는 완전히 마스킹됩니다.
역방향 과정: 양방향 트랜스포머는 마스킹된 시퀀스에서 원래 토큰을 예측하도록 훈련됩니다. 생성 과정에서는 완전히 마스킹된 시퀀스에서 시작하여 빈칸을 반복적으로 채우고, 때때로 토큰을 다시 마스킹하여 개선을 허용합니다.

균일 상태 확산(UDLM)

마스크 토큰 대신, UDLM은 토큰을 전체 어휘에서 무작위로 선택된 값으로 교체합니다. 이는 중간 시퀀스가 마스크 없이 유지되며, 어떤 단계에서든 어떤 토큰도 수정할 수 있도록 해, 더 나은 오류 수정과 제어 가능성을 가능하게 합니다.

고급 아키텍처 및 샘플링

이론적 모델에서 생산 가능한 LLM으로 전환하기 위해 여러 아키텍처적 및 알고리즘적 확장이 필요합니다.

블록 확산 및 가변 길이

표준 MDLM은 고정 길이 시퀀스에 한정됩니다. 블록 확산은 이 문제를 해결하기 위해 텍스트를 블록(예: 256 토큰) 단위로 생성하고, 이전에 생성된 블록에 조건을 붙입니다. 이 방법은 AR 모델과 유사하게 KV 캐싱을 지원하여 효율적인 가변 길이 생성을 가능하게 합니다.

인코더-디코더 프레임워크

현대의 dLLM, 예를 들어 Gemma Diffusion과 Nemotron Diffusion은 인코더-디코더 아키텍처를 사용합니다. 무거운 인코더는 깨끗한 컨텍스트를 한 번만 처리하고, 가벼운 디코더는 토큰의 "캔버스"를 반복적으로 노이즈 제거합니다. 이 분리 구조는 노이즈 제거 단계의 계산 비용을 줄이고 훈련 속도를 가속화합니다.

반복적 개선 및 오류 수정

표준 마스킹은 한 번 언마스킹된 토큰을 다시 마스킹하지 않기 때문에, 초기 실수는 영구적으로 남습니다. 두 가지 주요 해결책이 있습니다:

재마스킹: 예측된 토큰의 일부를 무작위로 다시 마스킹하는 플러그인 샘플러로, 더 많은 컨텍스트가 채워질수록 문법적 또는 논리적 오류를 수정할 수 있도록 합니다.
균일 상태 확산: 모든 토큰이 잠재적으로 노이즈가 될 수 있도록 하여, 모델이 시퀀스의 어떤 위치든 자연스럽게 수정할 수 있습니다.

증류를 통한 샘플링 가속화

확산 모델은 병렬 생성 덕분에 AR 모델보다 5~10배 빠를 수 있습니다. 더 빠른 속도를 위해 진행형 증류(progressive distillation)가 사용되며, 학생 모델이 교사 모델의 두 단계의 노이즈 제거를 단일 단계로 모방하도록 훈련하여, 필요한 샘플링 단계를 반복적으로 절반으로 줄입니다.

제어 가능한 생성 및 안내

확산 모델은 전역 시퀀스를 개선하기 때문에, 국소적이고 되돌릴 수 없는 편집을 고정하는 AR 모델보다 제어 가능한 생성에 더 적합합니다. 이는 샘플의 자연스러움과 속성 만족도 사이의 파레토 최적화로 관리됩니다.

분류기 기반 안내(CBG): 별도의 예측 모델을 사용하여 노이즈 제거 과정을 특정 속성(예: 분자에서의 결합 친화도)으로 유도합니다.
분류기 없는 안내(CFG): 조건부 및 무조건적 생성을 동시에 처리할 수 있도록 조건 신호를 무작위로 제거하여 단일 모델을 훈련하고, 두 출력을 결합하여 결과를 유도합니다.

실제 응용 및 확장성

생물학 및 과학 분야

확산은 생물학 분야에서 초기 성공을 거두었는데, 생물학적 시퀀스는 인간 언어의 왼쪽에서 오른쪽으로의 편향이 없으며 높은 제어 가능성이 필요하기 때문입니다.

ESM3: 단백질 시퀀스, 구조, 기능 모델링에 사용되는 100B 파라미터의 MDLM.
Nucleotide Transformer v3 (NT-v3): DNA 시퀀스 생성을 위한 수십억 파라미터 모델로, CFG와 재마스킹을 활용하여 특정 유전자 발현 수준을 가진 조절 DNA 시퀀스를 설계합니다.

일반 목적 확산 LLM

여러 선도 연구소에서 확장된 확산 모델을 공개했습니다:

LLaDA: 8B 파라미터의 오픈웨이트 모델로, GSM8K 및 MMLU와 같은 벤치마크에서 AR과 유사한 확장성을 보여줍니다.
Mercury 2: 상용 모델로, 일반 GPU에서 초당 약 1,200 토큰의 속도를 달성하여, Claude Haiku와 같은 속도 최적화 AR 모델을 크게 능가합니다.
Gemma Diffusion: Google의 오픈웨이트 모델로, UDLM 백본과 인코더-디코더 아키텍처를 사용합니다.
Nemotron Diffusion: NVIDIA의 모델 시리즈(최대 35B 파라미터)로, AR-확산 병합 목적을 사용하여 유사한 AR 모델 대비 2~8배의 처리량을 제공합니다.

비판적 분석 및 미래 전망

확산은 추론 속도와 하드웨어 활용도에서 막대한 성과를 제공하지만, 일부 기술적 과제는 여전히 남아 있습니다. 커뮤니티 피드백은 특정 약점, 즉 토큰 간 조율 문제를 지적합니다: 두 위치가 모두 두 가지 유효한 상태 사이에서 흔들릴 때, 확산 모델은 이들을 일치하는 쌍으로 합쳐내지 못해 일관성 없는 결과를 초래할 수 있습니다.

확장 관점에서 볼 때, 확산은 추론 시 계산에 대해 트랜스포머가 사전 훈련 계산에 기여한 것과 유사한 역할을 할 수 있습니다. AR 추론의 순차적 병목을 병렬 처리로 전환함으로써, 확산은 초당 FLOPs를 높일 수 있으며, 확장된 후 훈련 및 추론 시 검색을 통해 새로운 수준의 지능을 열어줄 수 있습니다.