[논문리뷰] Large Language Diffusion Models

선정이유

arxiv 논문이지만, Autoregressive라는 LLM의 틀을 깨는 발상이 담긴 논문이라고 생각해서 선정하게 되었습니다. 아직 여러 limitation이 많고 납득 가능할 정도의 성능이 나오지 않았다는 점을 감안하더라도 기존 LLM에서의 한계를 정면돌파하려는 시도를 담은 것 자체로 유의미한 논문이라고 생각해 소개드리고 싶었습니다.

---

Background

• GPT, LLaMA와 같은 Autoregressive model(ARM) 기반의 LLM들은 거의 언어 모델링의 표준이 됨.
  • 심지어 방대한 데이터와 함께 확장되면서 in-context learning, instruction-following 등 인상적인 능력을 보여왔음.

이에 대해, 저자들은 다음과 같은 의문점을 제시

1. LLM의 핵심 능력들(확장성, in-context learning, instruction-following등)이 ARM이라는 특정 구조에 본질적으로 종속된 것일까? 아니면 더 넓은 의미의 generative modeling 원리에서 비롯된 것일까?
2. ARM 방식 자체에 내재적인 한계는 없을까?

• 순차적인 생성으로 인한 계산 비용, 왼쪽에서 오른쪽으로의 모델링 한계로 인한 reversal reasoning 문제... 등등

Related Works

• LM에 Diffusion Model을 적용:

  • 텍스트 데이터를 연속적인(continuous) 공간에 임베딩하거나, 이산 분포의 연속적인 매개변수를 모델링하여 diffusion model을 적용하는 방법들이 있었음.
  • 하지만 이 방법들은 scalability 측면에서 어려움을 겪었으며, ARM과 비슷한 성능을 내기 위해 훨씬 많은 계산 자원이 필요했음.

• Discrete Diffusion Process를 사용하는 방법들:

  • 연속적인 diffusion 대신 이산적인 공간에서 새로운 순방향 및 역방향 프로세스를 정의하는 접근 방식.
  • 이 분야에는 다양한 변형들이 제안되었는데, 특히 Masked Diffusion Model (MDM)은 이산 diffusion의 한 특별한 경우로, 마스크 예측을 통해 언어를 모델링함.
  • Lou et al. (2023)은 MDM이 GPT-2 스케일에서 ARM과 유사하거나 더 나은 perplexity를 달성할 수 있음을 보여줌.
  • Nie et al. (2024)은 GPT-2 스케일에서 MDM을 Q&A 같은 언어 작업에 활용하는 것을 탐색.
  • Gong et al. (2024)은 기존 ARM을 MDM 형태로 fine-tune하는 방식을 시도.

• 그렇다면 LLaDA는?:

  • LLaDA는 이전 MDM 연구들을 기반으로 함.

  • 그러나 가장 중요한 차별점은, LLaDA가 MDM을 scratch부터 8B 매개변수까지 확장했다는 점.

  • 이를 통해 LLaDA는 LLaMA 3와 비교 가능한 성능을 달성함.

    → MDM 기반의 첫 번째 경쟁력 있는 non-autoregressive model

Purpose

궁극적으로, 1) LLM의 핵심 능력들이 ARM에 본질적으로 연결되어 있다는 통념에 도전하고, 2) diffusion model이 ARM의 대안이 될 수 있음을 LLaDA를 통해 입증하려 함.

---

Method

[Probabilistic Formulation]

• ARM: 다음 토큰을 순차적으로 예측.

• LLaDA: diffusion model처럼 순방향 및 역방향 프로세스를 통해 데이터 분포 p_theta(x_0)를 모델링.

  • 순방향 프로세스: 원본 데이터 x_0의 각 토큰이 시간 t에 따라 독립적으로 마스크될 확률이 선형적으로 증가.
    → t=0 에는 마스크가 없고, t=1에는 모든 토큰이 마스크됨.
  • 역방향 프로세스: 마스크된 x_t 로부터 원래 데이터 x_0 를 복원하는 과정
    → t=1 에서 t=0으로 이동하면서 마스크된 토큰을 반복적으로 예측
  • mask predictor 인 p_theta(·|x_t) 는 x_t를 입력받아 마스크된 모든 토큰을 동시에 예측.
  • 즉, ARM과 달리 causal mask를 사용하지 않고 양방향으로 문맥을 봄.

• 학습 목표: 마스크된 토큰에 대한 cross-entropy loss를 최적화. 이 손실 함수는 모델 분포의 negative log-likelihood에 대한 upper bound임이 이론적으로 증명되어, principled generative approach를 가능하게 함.

  • ARM의 목표: Chain Rule로 분해된 각 스텝의 다음 토큰 예측 확률을 최대화함으로써, 문장 전체가 나올 확률 (log-likelihood)을 직접적으로 최대화하는 것
  • but, LLaDA의 목표: 마스크된 토큰을 정확하게 예측하는 proxy task를 통해, 데이터 분포의 negative log-likelihood의 upper bound를 최소화함으로써, 간접적으로 데이터 분포를 잘 모델링하는 것
    +) 마스크된 토큰 예측 loss를 최소화하는 과정이 실제 데이터 분포가 모델 분포 하에서 나올 확률의 음수 값 (negative log-likelihood)을 이론적으로 넘어가지 않는(upper bound) 성질을 가지기 때문.
    ++) 마스크 예측을 잘 한다 → 이론적으로 실제 데이터의 언어 분포를 잘 학습했을 가능성이 높음.

• 기존 Masked Language Model (MLM)이 고정된 마스크 비율을 사용하는 반면, LLaDA는 t ∼ U[0, 1]에서 샘플링된 random masking ratio를 사용함.

  → 이 차이가 in-context learning 및 scalability에 중요하다고 주장

[Pre-training]

• 모델 아키텍처: Transformer를 사용하며, RMSNorm, SwiGLU, RoPE 등을 적용.
  • KV caching을 사용하지 않으므로 multi-head attention을 사용하며, 파라미터 수를 맞추기 위해 FFN dimension을 조정하는 등 일부 수정을 가함.
  • =) LLaDA는 ARM처럼 문장을 단어 하나씩 순서대로 만드는 게 아니어서, 이전 단어들의 계산 결과를 저장했다가 재활용하는 KV caching이라는 속도 기술을 쓰지 않았음. 이 때문에 가장 기본적인 attention 계산 방식을 사용할 수밖에 없음.
    → 이 기본적인 방식을 쓰면 어텐션 계산 부분에 들어가는 모델의 파라미터가 늘어나는데, 이를 다른 LLM들과 비슷하게 맞추려고 FFN의 크기를 약간 줄여서 전체 용량을 조절함.
• 데이터: 2.3조 개의 토큰으로 구성된 대규모 데이터셋(일반 텍스트, 코드, 수학, 다국어 데이터 포함)을 사용. 기존 LLM과 유사한 데이터 전처리 및 혼합 방식을 따름.
• 8B 모델 기준 0.13백만 H800 GPU hours 소요

[Supervised Fine-tuning(SFT)]

• instruction-following 능력 향상을 위한 단계
• 데이터: 450만 개의 (prompt, response) 쌍으로 구성된 데이터셋
• 프롬프트 p_0는 그대로 두고 응답 r_0의 토큰만 독립적으로 mask하여 r_t를 얻음.
  → (p_0, r_t)를 모델에 입력하고 mask된 응답 토큰 r_0을 예측하도록 학습 (pre-training과의 일관성 유지)
  • 짧은 응답쌍에는 EOS 토큰을 패딩하여 variable length를 처리하고, 모델이 EOS 생성을 통해 응답 길이를 제어하도록 함.
• Multi-turn dialogue는 이전 대화를 프롬프트에 포함하는 방식으로 처리.

[Inference]

• Sampling: t=1 (완전히 마스크됨)에서 시작하여 t=0으로 가는 reverse process를 시뮬레이션하여 텍스트 생성.
  • 주어진 프롬프트 p_0와 현재 마스크된 응답 r_t를 모델에 입력하여 마스크된 토큰에 대한 예측을 얻고. 다음 시간 단계 s에 대해, 예측된 토큰 중 일부(s/t 비율)를 다시 마스크하여 r_s를 만듬. 이 과정을 반복하여 t=0까지 진행함.
  • Remasking Strategies
    • random remasking
    • low-confidence remasking: 예측 신뢰도가 가장 낮은 토큰부터 일정 비율을 다시 마스크
    • semi-autoregressive remasking: 시퀀스를 여러 블록으로 나누고 왼쪽에서 오른쪽으로 블록 단위로 생성하며, 각 블록 내에서는 diffusion 샘플링을 적용
• Conditional Likelihood Evaluation
  • 응답 r_0에서 균일하게 토큰 l개를 마스크한 r_l을 만들고, (p_0, r_l)이 주어졌을 때 마스크된 토큰을 예측하는 모델의 로그 확률을 계산

Experiment

[Scalability]

• LLaDA는 자체 구축한 ARM baseline과 동일한 데이터로 학습했을 때, FLOPs 증가에 따라 성능이 인상적으로 확장됨을 보여줌.
• MMLU나 GSM8K 같은 특정 태스크에서는 ARM baseline보다 더 강한 확장성을 보이기도 함

→ diffusion model 방식도 ARM만큼, 혹은 특정 측면에서는 더 효과적으로 확장될 수 있음!

[In-Context Learning 및 Benchmark Results]

• pre-trained LLaDA 8B Base 모델은 LLaMA2 7B Base를 거의 모든 태스크에서 능가, LLaMA3 8B Base(15조 토큰 학습)와 전반적으로 경쟁력 있는 성능을 보임.

• 특히 수학(GSM8K, Math) 및 중국어(CMMLU, C-Eval) 태스크에서 LLaDA의 강점이 두드러졌음.

  +) 데이터셋 차이 등 여러 요인이 있을 수 있지만, 일단은 LLaDA가 특정 영역에서 우수함...

[Instruction-Following 및 SFT Results]

• SFT을 거친 LLaDA 8B Instruct 모델은 대부분의 downstream 태스크에서 성능이 향상.
• LLaMA3 8B Instruct가 SFT 외에 reinforcement learning (RL) 기반의 alignment를 추가로 수행한 것에 비해, LLaDA는 SFT만으로도 비슷한 수준의 명령어 이해 능력을 보여줌.
• 그러나, RL alignment를 거친 모델들보다는 뒤처짐.

→ 어느정도의 instruction following도 보장됨.

[Reversal Reasoning 및 Reversal Curse 극복]

• ARM 모델의 한계로 지적되었던 Reverse Task에서의 성능 비교
  Task 예시
  • Prompt (Reversal): 但闻人语响的上一句是什么？直接输出句子即可。
    (단문인어향의 이전 행은 무엇인가요? 문장만 바로 출력하세요.)
  • Poem Context: 이 시구는 왕유(王維)의 '녹채(鹿柴)'라는 시의 일부입니다. 시는 다음과 같습니다:
    空山不见人，(공산불견인)
    但闻人语响。(단문인어향)
    返景入深林，(반경입심림)
    复照青苔上。(부조청태상)
  • Expected Reversal Answer: 空山不见人， (공산불견인)
    LLaDA 8B Instruct Answer: 空山不见人 (정확함, Tab 14 Q1)
    GPT-4o Answer: 无可奈何花落去 (오답, 이 시와 전혀 상관없는 이청조의 시구임)
• LLaDA는 시(poem) 완성 태스크(정방향 vs 반전)에서 GPT-4o 및 Qwen2.5와 달리 정방향과 반전 태스크에서 일관된 성능을 보이며 reversal curse 현상을 효과적으로 해결했음을 입증.
• 특히 반전 태스크에서는 GPT-4o보다 월등히 우수한 성능을 달성.

→ LLaDA의 bidirectional modeling 능력을 통해 Reverse Task의 한계를 극복!

[Case Studies]

• LLaDA 8B Instruct 모델의 생성 샘플을 통해 유창한 텍스트 생성 능력을 볼 수 있음.
  • 위의 예시: 복잡한 지시를 이해하고 다단계 추론을 통해 수학 문제를 해결하는 능력
    * 이미지 속 토큰의 색상 변화는 LLaDA의 샘플링 과정에서 토큰이 예측되는 단계를 시각화한 것으로, 어두운 색일수록 나중에 예측된 토큰임을 나타냄.
    +) 그래도 답이 가장 나중에 나올 줄 알았는데 그건 아니었음
  • 아래 예시: 대화의 맥락을 유지하며 복잡하고 다양한 유형의 지시를 이해하고 따르는 모습(유명한 시 제공, 번역 후 창작 등)
• multi-turn dialogue에서 대화 맥락을 유지하며 다양한 언어로 적절하게 응답하는 능력을 보여주며, 기존 ARM과 차별화되는 chat 능력을 강조.

---

Key point

1) LLM의 핵심이 ARM이라는 특정 구조를 통해서만 가능한 것이 아니라, 더 넓은 generative modeling 원리에 기반하며 diffusion model로도 달성될 수 있다는 근본적인 주장을 펼침.

2) MDM 기반의 새로운 LLM 접근 방식인 LLaDA를 제안하고, 이를 표준적인 pre-training 및 SFT 파이프라인에 맞춰 from scratch to 8B 규모까지 학습시켰음.

3) LLaDA는 Transformer를 마스크 예측기로 사용하며, ARM과 달리 양방향 문맥을 활용.

4) LLaDA 8B 모델이 기존의 강력한 ARM들(특히 LLaMA3 8B와 비교해)과 비슷하거나 특정 영역에서는 우수한 성능을 보임. 이는 diffusion model이 LLM 스케일에서 실제로 작동하며 경쟁력을 가질 수 있음을 처음으로 보여준 사례.

5) FLOPs 증가에 따라 LLaDA의 성능이 ARM baseline과 유사하거나 특정 태스크에서는 더 강하게 확장됨을 보여주며, diffusion model이 대규모 데이터와 모델에 대해 효과적으로 확장 가능함을 밝힘.

6) LLaDA의 양방향 모델링 특성이 ARM의 고유한 한계로 여겨졌던 reversal curse 현상을 효과적으로 해결했음. 반전 추론 태스크에서 GPT-4o보다 우수한 성능을 보인 것이 대표적.

---

Limitations and Suggestions

• LLaDA와 ARM baseline 간의 직접적인 비교(identical dataset 학습 등)는 컴퓨팅 자원 문제로 인해 10^23 FLOPs 이하의 스케일에서만 이루어졌기 때문에 LLaDA의 최대 규모(8B) baseline으로 ARM을 동일하게 학습시키지는 못했다.
• LLaDA를 위해 특별히 고안된 어텐션 메커니즘이나 위치 임베딩, 또는 시스템 수준의 최적화가 적용되지 못했다.
• 아직 어떤 guidance 기술이 LLaDA에 가장 잘 맞는지, 어떻게 써야 제일 좋은 결과를 내는지에 대한 연구가 충분히 안 되어 있다.
• LLaDA는 SFT만 수행했으며, RL 기반의 alignment를 거치지 않았다.
• LLaDA가 multi-modal data를 처리할 수 있는지 여부가 탐구되지 않았음.
  • diffusion model은 multi-modal 생성에 강점을 보이므로, LLaDA를 multi-modal LLM으로 확장하는 후속 연구가 있을지도..
• Distillation 등을 통한 샘플링 속도 개선 연구를 LLaDA에 적용하는 것도 중요해보임.

Personal Insight

• 가장 큰 의의는 기존 LLM 연구의 정석으로 여겨지던 ARM을 탈피했다는 점이라고 생각함!
  • LLM 능력의 근원이 ARM 구조 자체에 있는 것이 아니라, diffusion model과 같은 다른 principled generative modeling 방식에서도 충분히 발현될 수 있음을 증명했음.
• Diffusion Model은 이미지 생성 분야에서 대성공을 거두었지만 이산적이고 구조적인 텍스트 데이터에는 적용하기 어렵다는 인식이 강했는데, 이 논문은 MDM이라는 특정 형태의 Diffusion Model이 대규모 텍스트 데이터에서도 ARM과 견줄만한 성능을 낼 수 있음을 보여주면서, Diffusion Model의 적용 가능 범위를 크게 확장했음.
• ARM의 고질적인 문제였던 reversal curse를 LLaDA가 효과적으로 해결함!
  • 모델의 bidirectional modeling 능력이 단순한 마스크 예측을 넘어 관계와 논리를 더 잘 포착할 수 있음을 밝힘. ARM의 순차성이 일종의 정보 처리 병목이 될 수 있음을 보여주면서, LLaDA 방식의 잠재력을 부각함.
• 기존의 MDM 연구들이 주로 GPT-2 스케일(약 3억 파라미터) 이하에서 가능성을 탐색하거나, 이미 학습된 ARM을 활용한 것과 달리, LLaDA는 8B 파라미터 규모의 MDM을 "from scratch" 학습시키는 도전정신... 리스펙
• LLaMA3 8B나 GPT-4o와 같은 상위 LLM들과 성능을 비교하는 등 데이터 투명성 등의 한계가 있음에도 불구하고 자신들의 모델이 기존 모델들과 경쟁할 수 있음을 보여주려는 시도가 인상적이었음.

Overall...

저돌적인 시도인데다 말투도 뭔가 자신감 넘쳐서 ("We challenge this notion" 해놓고 시작.. 중간에 뜬금없이 등장하는 명언들..) 일단 hooking이 제대로였다. reversal curse를 극복한 부분은 diffusion model의 bidirectional 특성이 언어 이해 방식에 있어 ARM의 순차적 제약을 뛰어넘을 수도 있음을 보여주는 결정적인 증거라고 생각한다. 앞으로 Diffusion Model 기반 LLM이 ARM을 얼마나 흔들 수 있을지, 혹은 두 방식의 장점을 결합한 새로운 형태의 모델이 등장할지 기대가 된다!