DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

기술 요약: DSL-LLaDA: 8B 마스크 확산 언어 모델을 위한 연속적 디노이징 스케일링

문제 정의
LLaDA와 같은 이산 마스크 확산 언어 모델(Discrete Masked Diffusion Language Models, DLMs)은 반복적인 병렬 디코딩을 통해 텍스트를 생성합니다. 이러한 모델들이 수십억 개의 파라미터 규모로 확장되었음에도 불구하고, 낮은 스텝 예산(적은 단계의 생성) 환경에서는 추론 품질이 급격히 저하됩니다. 핵심 문제는 출력 길이와 품질 사이의 트레이드오프입니다. 고정된 수의 순방향 패스(forward passes)를 사용할 때, 표준적인 반복적 언마스킹(iterative unmasking)은 짧고 고품질인 출력을 생성하거나, 혹은 길고 반복적이며 일관성 없는 텍스트를 생성하게 됩니다. 이러한 실패는 역방향 전이(reverse transition)의 인수분해(factorization)에서 기인합니다. 즉, 각 토큰 위치가 컨텍스트에 조건화된 주변 분포로부터 독립적으로 디코딩되면서 토큰 간 상관관계를 무시하기 때문입니다. 스텝 예산이 낮을 때, 각 단계는 많은 위치를 동시에 결정해야 하며, 이는 오류를 증폭시켜 조기 종료나 반복 루프를 유발합니다.

임베딩 공간에서의 연속적 디노이징(continuous denoising)은 모든 위치를 공동으로 진화시켜 토큰 공간의 인수분해 문제를 우회함으로써 이론적인 해결책을 제시하지만, 임베딩 기하학, 노이즈 스케줄링, 학습 목적 함수의 복잡성으로 인해 이러한 모델을 처음부터 대규모로 구축하는 것은 여전히 미해결 과제로 남아 있습니다.

방법론
저자들은 기존의 사전 학습된 마스크 DLM(구체적으로 LLaDA-8B-Instruct)을 사용하여 연속적 임베딩 공간 디노이징을 지원하도록 가볍게 적응시키는 방법인 DSL-LLaDA를 제안합니다. 이는 연속 모델을 처음부터 학습시키는 대신 기존 모델을 활용하는 방식입니다.

1. 이산 확률적 국소화 (Discrete Stochastic Localization, DSL) 적응:

   • 이 방법은 표준 DLM의 이진(binary) clean-or-mask 입력을 연속적인 노이즈 상태로 대체합니다. 깨끗한 토큰 $x_i$에 대해, 신호 대 잡음비(SNR) $\gamma_i$를 샘플링하고, 노이즈가 섞인 상태 $z_i$를 토큰 임베딩과 가우시안 노이즈의 혼합물로 구성합니다.
   • **소프트맥스 컨버터(softmax converter)**는 이러한 연속적인 DSL 상태를 백본의 입력 임베딩 공간으로 다시 매핑합니다. 이 컨버터는 낮은 SNR은 마스크와 유사한 상태를, 중간 SNR은 신뢰할 수 없는 랜덤 토큰을, 높은 SNR은 명확한 토근 임베딩을 나타내는 "소프트 마스크"를 생성합니다.
   • 백본은 FineWeb-Edu 데이터셋을 통해 단 1,000 스텝 동안 **연속 사전 학습(continue-pretrained)**됩니다. 학습 목적은 표준 토큰 교차 엔트로피를 유지하지만, 입력 상태는 이제 이진 마스크가 아닌 연속적인 혼합물입니다. 결정적으로, 컨버터 덕분에 모델은 아키텍처 변경 없이 이러한 입력을 처리할 수 있습니다. 이는 최적의 디노이저가 갖는 SNR 불변성(SNR-invariance) 덕분에 명시적인 타임스텝 입력이 필요 없기 때문입니다.

2. 연속적 추론 (SDE 샘플러):

   • 추론 과정에서 모델은 모든 위치에 대해 연속적인 상태를 유지합니다. 모델은 낮은 SNR(높은 불확실성)에서 높은 SNR로 상태를 진화시키기 위해 SDE 기반 샘플러(Heun steps)를 사용합니다.
   • 모델은 토큰에 대한 확률 분포를 예측하며, 이는 (상위 $M$개의 유력한 토큰을 사용하는) 가중 평균 임베딩으로 변환됩니다. 이 임베딩 추정치가 연속 상태를 업데이트합니다.
   • 하드 토큰 결정(Hard token commitment)은 마지막 단계까지 유예됩니다. 이를 통해 모델은 모든 위치를 공동으로 정밀하게 조정하여 생성 과정 전반에 걸쳐 토큰 간 상관관계를 보존할 수 있습니다.

3. 조건부 생성:

   • 요약과 같은 작업의 경우, 컨텍스트 토큰은 깨끗한 임베딩으로 고정(clamping)되는 반면, 응답 영역은 연속적 디노이징을 거칩니다. 백본은 양방향으로 어텐션(attention)을 수행하여, 깨끗한 컨텍스트에 조건화되어 응답을 디노이징합니다.

주요 기여

1. 경량화된 변환: 기존의 8B 마스크 DLM을 DSL을 통해 최소 1,000 스텝의 적응만으로 연속적 공간 디노이징을 지원하도록 변환할 수 있음을 입증하였으며, 이를 통해 연속 모델을 처음부터 학습시켜야 하는 부담을 피했습니다.
2. 저효율 NFE 실패 모드 회피: 연속적 SDE 추론이 낮은 스텝 예산(NFE $\le$ 16)에서 발생하는 조기 종료 및 반복 루프 특성을 효과적으로 회피함을 보여주었습니다.
3. 선택적 노이즈 상태 강건성 식별: 적응 과정이 모델로 하여금 깨끗한 토큰은 보존하면서 오염된 토큰을 수정할 수 있는 "선택적 수정(selective correction)" 능력을 유도함을 확인했습니다. 이는 동일한 컴퓨팅 자원을 사용한 이산 베이스라인에서는 나타나지 않는 동작입니다.

결과

• 제로샷 요약: 네 가지 벤치마크(XSum, CNN/DailyMail, arXiv, PubMed)에서 DSL-LLaDA-SDE는 길이 제어 휴리스틱 없이도 낮은 스텝 예산($\le$ 16 forward passes)에서 가장 우수한 ROUGE-1 점수를 기록했습니다. 반면, 이산 베이스라인들은 출력을 조기에 잘라내거나 높은 반복률을 보였습니다.
• 반복 및 길이: 개방형 생성(open-ended generation)에서 DSL-LLaDA-SDE는 모든 스텝 예산에 대해 반복률을 9% 미만으로 유지했으나, 이산 베이스라인들은 낮은 NFE에서 60% 이상의 반복률을 보이거나 매우 짧은 길이로 잘리는 현상을 보였습니다.
• 오류 수정: 적응된 모델은 오염된 토큰을 복구하는 동시에 깨끗한 토큰의 98% 이상을 보존하는 "선택적 수정" 능력을 보여주었습니다. 이에 비해, 랜덤 토큰 오염(XDLM)을 적용한 이산 베이스라인은 깨끗한 토큰의 상당 부분을 손상시켰으며, 표준 마스크 확산(MDM-CPT)은 미미한 수정 능력만을 보였습니다.
• 대조 실험: 동일한 연산량을 사용한 표준 마스크 확산 학습(MDM-CPT) 및 랜덤 토큰 오염(XDLM) 실험은 연속적 디노이징의 이점이나 선택적 강건성을 재현하지 못했습니다. 이는 개선 사항이 단순히 추가적인 연산량 때문이 아니라 연속적 노이즈 노출에서 기인함을 확인시켜 줍니다.

의의 및 주장
본 논문은 DSL-LLaDA가 마스크 DLM을 위한 범용적인 변환 절차를 제공하여, 최소한의 학습 오버헤드로 연속적 디노이징(토큰의 공동 진화, 결정 유예)의 이점을 활용할 수 있게 한다고 주장합니다. 결과는 연속적 노이즈 적응이 낮은 스텝 예산에서의 반복적 언마스킹에 내재된 길이-품질 트레이드오프를 해결할 수 있음을 시사합니다.

저자들은 다음과 같은 한계점을 명시하며 신중한 태도를 유지합니다:

• 평가는 주로 앵커드 생성(anchored generation)(요약 및 지시 이행)에 국한되었습니다. 개방형 창의적 글쓰기에서는 여전히 더 높은 퇴보율(요약의 4% 미만 대비 24%)을 보입니다.
• 본 방법론은 현재 LLaDA-8B 규모에서만 검증되었습니다. 다른 아키텍처나 더 큰 규모에서의 성능은 아직 테스트되지 않았습니다.
• SDE가 낮은 NFE에서는 이산 방식보다 우수하지만, 긴 참조(long-reference) 작업의 높은 NFE에서는 베닐라 LLaDA가 SDE와 대등하거나 더 나은 성능을 보입니다(이 경우 조기 종료 문제가 크지 않기 때문입니다).
• 저자들은 동일 규모의 자기회귀(AR) 모델과 동등하다고 주장하지 않으며, 높은 스텝 예산의 긴 참조 작업에서는 여전히 AR 모델이 앞선다는 점을 언급했습니다.
• 인간 평가(human evaluation)는 향후 과제로 남겨두었습니다.