[논문 리뷰] Scaling Beyond Masked Diffusion Language Models

TL;DR

디퓨전 언어 모델은 자기회귀(AR) 모델의 순차적 생성 한계를 극복할 잠재력을 가집니다. 이 논문은 기존의 마스크드 디퓨전 모델을 넘어, 유니폼-스테이트(Uniform-State) 및 인터폴레이팅(Interpolating) 디퓨전 아키텍처의 스케일링 법칙을 최초로 탐구했습니다. 연구 결과, 퍼플렉시티가 전부는 아니며, 유니폼-스테이트 모델은 추론 능력에서, 인터폴레이팅 모델은 생성 속도에서 강점을 보임을 입증했습니다. 이는 디퓨전 언어 모델의 실용적인 가능성을 한 단계 끌어올린 중요한 연구입니다.

연구 배경 및 동기

대규모 언어 모델(LLM)의 주류는 GPT와 같은 자기회귀(Autoregressive, AR) 모델입니다. AR 모델은 이전 토큰을 바탕으로 다음 토큰을 순차적으로 예측하는 방식으로 높은 품질의 텍스트를 생성하지만, 토큰을 하나씩 생성해야 하므로 느린 추론 속도라는 본질적인 한계를 가집니다.

이에 대한 대안으로 디퓨전(Diffusion) 모델이 주목받고 있습니다. 이미지 생성 분야에서 뛰어난 성능을 보인 디퓨전 모델은 전체 시퀀스를 병렬적으로 처리하여 한 번에 생성할 수 있으므로, 이론적으로 훨씬 빠른 추론이 가능합니다. 하지만 언어 모델링 분야에서는 아직 초기 단계이며, 대부분의 연구가 BERT의 마스크드 언어 모델링(MLM)을 차용한 **마스크드 디퓨전 모델(MDLM)**에 집중되어 있었습니다.

이 연구는 MDLM이라는 단일 아키텍처에서 벗어나, 다양한 디퓨전 모델 아키텍처의 스케일링 법칙을 분석하여 각 모델의 잠재력과 한계를 규명하고자 합니다. 이를 통해 디퓨전 모델이 AR 모델의 강력한 경쟁자가 될 수 있는 길을 모색합니다.

관련 연구

디퓨전 모델은 점진적으로 노이즈를 제거(denoising)하며 데이터를 생성하는 모델입니다. 언어 모델링 분야에서는 주로 MDLM이 연구되었으며, 이는 퍼플렉시티(Perplexity, PPL) 지표에서 매우 뛰어난 성능을 보여주었습니다. 그러나 선행 연구들은 PPL 지표에 과도하게 의존하는 경향이 있었고, 이는 실제 다운스트림 태스크에서의 성능을 완벽히 대변하지 못하는 한계가 있었습니다.

본 연구는 이러한 한계를 지적하며, PPL 외에 실제 추론 능력 및 생성 속도와 같은 실용적인 지표를 통해 모델을 다각적으로 평가합니다. 특히, 유니폼-스테이트와 인터폴레이팅 디퓨전 모델의 스케일링 법칙을 분석한 최초의 연구라는 점에서 차별점을 가집니다.

핵심 기여

1. 최초의 스케일링 법칙 연구: 유니폼-스테이트 및 인터폴레이팅 디퓨전 모델의 스케일링 법칙을 최초로 정립하여, 모델 크기에 따른 성능 변화를 예측하고 이해할 수 있는 기반을 마련했습니다.
2. 학습 효율성 개선: 마스크드 디퓨전 모델의 학습 목표를 기존의 복잡한 방식에서 단순한 크로스 엔트로피(Cross-Entropy)로 변경하여, 약 12%의 연산량 절감을 달성하면서도 성능을 유지했습니다.
3. 퍼플렉시티의 한계 지적: PPL이 동일 아키텍처 내에서는 유용한 지표이지만, 서로 다른 아키텍처(예: 디퓨전 vs. AR)를 비교할 때는 실제 성능을 오해하게 만들 수 있음을 실험적으로 증명했습니다.
4. 실험적 검증: 다양한 벤치마크 실험을 통해, 유니폼-스테이트 디퓨전 모델이 수학 및 추론 능력에서 다른 모델들을 능가하는 성능을 보임을 입증했습니다.

핵심 방법론: 세 가지 디퓨전 아키텍처

디퓨전 모델은 깨끗한 데이터 $x_0$에 점진적으로 노이즈를 추가하는 **순방향 프로세스(Forward Process)**와, 노이즈가 낀 데이터 $x_t$로부터 원본 데이터 $x_0$를 복원하는 **역방향 프로세스(Reverse Process)**를 학습합니다. 이 연구에서 비교하는 세 가지 아키텍처는 순방향 프로세스, 즉 노이즈를 정의하는 방식에서 차이가 있습니다.

1. 마스크드 디퓨전 (Masked Diffusion)

BERT의 마스크드 언어 모델링(MLM)과 가장 유사한 방식입니다. 원본 시퀀스의 일부 토큰을 [MASK]라는 특수 토큰으로 교체하여 노이즈를 만듭니다. 모델은 이 [MASK] 토큰의 원래 단어를 예측하도록 학습됩니다.

• 예시: The quick brown fox → The [MASK] brown [MASK]
• 특징: 기존 언어 모델링 벤치마크에서 매우 낮은 퍼플렉시티를 달성하며 강력한 성능을 보입니다.

2. 유니폼-스테이트 디퓨전 (Uniform-State Diffusion)

[MASK] 토큰 대신, 어휘(vocabulary) 전체에 대한 **균등 분포(Uniform Distribution)**를 초기 노이즈 상태로 사용합니다. 이는 특정 토큰을 가리는 것이 아니라, 해당 위치에 어떤 단어가 올지 전혀 모르는 '완전한 불확실성' 상태에서 시작하여 점차 정답 토큰으로 좁혀나가는 과정을 학습하는 것입니다.

• 예시: The quick brown fox → The [UNIFORM] brown [UNIFORM]
• 특징: 퍼플렉시티는 상대적으로 높지만, 샘플링 효율성이 뛰어나 더 적은 스텝으로도 고품질 텍스트 생성이 가능합니다. 이는 모델이 특정 [MASK] 토큰에 의존하지 않고, 일반적인 불확실성을 해결하는 방법을 배우기 때문입니다.

3. 인터폴레이팅 디퓨전 (Interpolating Diffusion)

이전 상태와 다음 상태를 **보간(interpolate)**하는 방식으로 상태를 전환합니다. 이 구조의 가장 큰 장점은 자기회귀 모델의 핵심 기술인 **KV 캐싱(Key-Value Caching)**을 자연스럽게 지원한다는 점입니다. 각 디코딩 스텝에서 이전 상태의 계산 결과를 재사용할 수 있어, 다른 디퓨전 모델보다 훨씬 빠른 추론 속도를 달성할 수 있습니다.

• 특징: AR 모델과 유사한 구조적 이점을 가져 추론 속도가 매우 빠릅니다. 실시간 응용에 가장 적합한 디퓨전 방식입니다.

디퓨전 모델의 학습 원리 (수식)

디퓨전 모델은 일반적으로 NELBO(Negative Evidence Lower Bound)를 최소화하도록 학습되며, 이는 전체 시퀀스 $x_0$에 대한 로그 우도(log-likelihood)를 최대화하는 것과 같습니다. 모델의 손실 함수 $L$은 모든 디퓨전 스텝 $t$에 대한 디노이징(denoising) 능력의 기댓값으로 표현됩니다.

• $x_0$: 원본 시퀀스 (예: "The quick brown fox")
• $t$: 디퓨전 스텝 (시간, $1$부터 $T$까지)
• $x_t$: 스텝 $t$에서 노이즈가 추가된 시퀀스
• $q(x_t|x_0)$: 원본 $x_0$에서 노이즈 버전 $x_t$를 만드는 순방향 프로세스
• $p_\theta(x_{t-1} | x_t)$: 모델이 현재 상태 $x_t$를 보고 한 단계 전의 상태 $x_{t-1}$을 예측할 확률 (역방향 프로세스)

이 수식의 핵심은 모델($p_\theta$)이 현재의 노이즈 상태($x_t$)에서 이전의 덜 노이즈 낀 상태($x_{t-1}$)를 복원하는 능력을 학습하는 것입니다. 각 디퓨전 아키텍처는 $q(x_t|x_0)$를 정의하는 방식, 즉 노이즈를 만드는 방식에서 차이가 발생하며, 이것이 모델의 고유한 특성을 결정합니다.

실험 설정

연구진은 세 가지 디퓨전 아키텍처와 자기회귀 모델을 최대 1.7B 파라미터까지 확장하며 성능을 비교했습니다. 모든 모델은 The Pile 데이터셋으로 학습되었습니다.

• 데이터셋: The Pile (사전 학습), GSM8K (수학 추론), MMLU (상식 추론) 등 다양한 벤치마크 사용
• 평가 지표: 퍼플렉시티(PPL), 제로샷/퓨샷 정확도, 추론 속도
• 베이스라인: LLaMA 아키텍처 기반의 자기회귀(AR) 모델과 기존 마스크드 디퓨전(MDLM) 모델

실험 결과 분석

실험 결과는 각 디퓨전 아키텍처가 서로 다른 강점을 가지고 있음을 명확히 보여줍니다.

주요 결과 요약

결과 해석

• 퍼플렉시티의 함정: 마스크드 디퓨전(MDLM)은 PPL이 가장 낮았지만, 수학 추론 능력에서는 다른 모델에 뒤처졌습니다. 이는 PPL이 언어의 통계적 패턴을 잘 포착하는 능력은 보여주지만, 복잡한 추론 능력을 직접적으로 대변하지는 못함을 시사합니다.
• 추론 능력의 승자, 유니폼-스테이트: 유니폼-스테이트 디퓨전(Duo)은 PPL이 가장 높았음에도 불구하고, 수학 추론 벤치마크인 GSM8K에서 66.0%로 가장 높은 정확도를 기록했습니다. 이는 AR 모델보다 약 5% 높은 수치로, 불확실한 상태에서 정보를 재구성하는 학습 방식이 추론 능력에 긍정적인 영향을 미쳤을 가능성을 보여줍니다.
• 속도의 승자, 인터폴레이팅: 인터폴레이팅 디퓨전(Eso-LM)은 KV 캐싱 덕분에 AR 모델보다 훨씬 빠른 추론 속도를 달성하여, 실시간 서비스에서의 활용 가능성을 입증했습니다.

Ablation Study

연구진은 학습 목표를 크로스 엔트로피로 단순화하는 등의 Ablation Study를 통해, MDLM의 학습 연산량을 12% 절감하면서도 성능 저하가 거의 없음을 확인했습니다. 이는 디퓨전 모델의 학습 효율을 높이는 중요한 발견입니다.

비판적 평가

강점

1. 포괄적인 아키텍처 비교: 단일 모델에 국한되지 않고, 여러 디퓨전 아키텍처의 스케일링 법칙을 분석하여 각 모델의 장단점을 명확히 규명했습니다.
2. 실용적 관점 제시: 퍼플렉시티라는 전통적인 지표의 한계를 지적하고, 추론 능력과 속도 같은 실용적인 측면을 강조하여 연구의 방향성을 제시했습니다.
3. 효율성 개선: 학습 목표 단순화를 통해 디퓨전 모델의 높은 학습 비용 문제를 일부 해결할 수 있는 실마리를 제공했습니다.

한계점과 개선 방향

1. 데이터셋 편향성: 실험이 주로 영어 중심의 The Pile 데이터셋으로 이루어져, 다국어 환경이나 특정 도메인에서의 일반화 성능은 추가 검증이 필요합니다.
2. 정량적 속도 비교 부족: 인터폴레이팅 모델의 빠른 속도를 강조했지만, 실제 하드웨어에서 측정한 구체적인 토큰/초(tokens/sec)와 같은 정량적 비교 데이터가 제시되었다면 더 설득력 있었을 것입니다.

재현성 평가

저자들은 학습된 모델 체크포인트와 코드를 모두 공개하여 연구의 재현성을 높였습니다. 상세한 실험 설정과 구현 코드를 통해 다른 연구자들이 결과를 검증하고 후속 연구를 진행하기 용이합니다.

향후 연구 방향

본 연구는 디퓨전 언어 모델의 새로운 가능성을 열었습니다. 향후 연구는 다음과 같은 방향으로 확장될 수 있습니다.

1. 아키텍처 융합: 각 모델의 장점을 결합한 하이브리드 모델 개발. 예를 들어, 유니폼-스테이트의 추론 능력과 인터폴레이팅의 빠른 속도를 결합하는 연구가 가능합니다.
2. 다양한 다운스트림 태스크 적용: 요약, 번역, 질의응답 등 더 넓은 범위의 다운스트림 태스크에서 각 디퓨전 모델의 성능을 평가하여 실용성을 검증해야 합니다.
3. 샘플링 전략 최적화: 디퓨전 모델의 생성 품질과 속도에 큰 영향을 미치는 샘플링 스케줄러 및 전략에 대한 심도 있는 연구가 필요합니다.

실무 적용 가이드

디퓨전 모델을 실무에 도입할 때, 다음 사항을 고려할 수 있습니다.

1. 목적에 맞는 모델 선택:
   • 높은 추론 능력이 필요하다면 (예: 코드 생성, 수학 문제 풀이): 유니폼-스테이트 디퓨전
   • 실시간 응답이 중요하다면 (예: 챗봇, 대화형 AI): 인터폴레이팅 디퓨전
   • 기존 LM 벤치마크와의 비교가 중요하다면: 마스크드 디퓨전
2. 추론 비용 고려: 디퓨전 모델은 병렬 처리의 이점이 있지만, 여전히 여러 번의 디노이징 스텝이 필요합니다. 서비스의 지연 시간(latency) 요구사항에 맞춰 적절한 스텝 수를 튜닝해야 합니다.

결론

이 논문은 디퓨전 언어 모델 연구를 마스크드 디퓨전이라는 틀에서 해방시키고, 다양한 아키텍처의 가능성을 탐색한 선구적인 연구입니다. 퍼플렉시티가 모델 성능의 절대적인 척도가 아니며, 목적에 따라 최적의 아키텍처가 다를 수 있음을 명확히 보여주었습니다. 유니폼-스테이트 모델의 뛰어난 추론 능력과 인터폴레이팅 모델의 빠른 속도는 디퓨전 모델이 머지않아 자기회귀 모델의 강력한 대안이 될 수 있음을 시사합니다.

참고 자료

• 논문 링크 (arXiv)
• 프로젝트 페이지 및 코드 저장소