DEVOCEAN

안녕하세요. SK텔레콤의 정지헌입니다.

Language model의 새로운 패러다임 Large Language Diffusion Model, 일명 LLaDA라고 불리는 논문을 소개하려고 합니다.

개인적으로 Diffusion model에 관심이 있고 Language model에 사용한 것이 재밌어서 유심히 읽어보게 되었고 여러 가지 활용 방법들을 구상 중에 있습니다.

우선 왜 Vision쪽에서 활용되던 Diffusion model이 Language에 사용되고, 적용되고 있는 지 알아보고자 합니다.

Large Language Diffusion Model의 논문 리뷰 시작합니다.

현재의 Autoregressive LLM의 한계

지금의 대형 언어모델의 성능에서는 너무 뛰어나지만, 대부분 Autoregressive(AR) 방식 (Left-to-Right)으로 학습 및 생성합니다.

약점이 전혀 없을 것 같지만 이 구조에는 세 가지 고질적 약점이 있습니다.

1. Exposure bias: 추론 때 모델이 스스로 낸 오류를 계속 입력으로 받아 악순환이 벌어집니다.

2. 길이 편향: 로그확률 합을 극대화하려다 보니 응답을 지나치게 짧게 끝내거나 불필요하게 길게 늘어뜨리곤 합니다.

3. 장기 의존 실패: 순차 생성 특성상 먼 거리 토큰 간 일관성을 유지하기가 어렵습니다.

이를 고치려고 scheduled sampling, Prefix LM, 고급 디코딩 기법 등을 시도했지만, 모두 AR 틀 안에서 부분 수정을 가한 정도에 머물렀습니다.

하이퍼파라미터가 민감하거나, 짧은 문장엔 효과가 있어도 길고 복잡한 context에서는 오류 누적을 막지 못합니다.

Diffusion 이 무엇인가요?

이미지 생성 분야에서 두각을 나타낸 Diffusion(확산) 모델은 무작위 노이즈를 넣었다가 조금씩 걷어 내어 데이터를 복원하는 방식을 자연어 영역으로 확장했습니다.

이미지에서 성공한 Diffusion기법 즉, 노이즈를 넣고 단계별로 걷어내며 복원 하는 방식을 그대로 자연어에 적용한 것이 Diffusion LM입니다.

AR 모델은 토큰을 왼쪽에서 오른쪽으로 하나씩 예측하기 때문에, 앞서 틀린 토큰이 계속 영향을 미치는 exposure bias와 긴 문장에서 뒤쪽 품질이 떨어지는 길이 편향 문제가 있었습니다.

반면에, Diffusion LM은 문장 전체(또는 넓은 블록)를 여러 단계에 걸쳐 반복 정제합니다.

한 단계에서 잘못 생성된 토큰도 다음 단계에서 다시 고칠 수 있으므로 exposure bias가 줄어들고, 주어진 컨텍스트 윈도 안에서 전역적 일관성을 확보하기 쉽다고 합니다.

훈련(Forward) 단계에서는 일정 비율의 토큰을 [MASK]로 가려 노이즈를 만듭니다.

그 이후 추론(Reverse) 단계에서는 모든 토큰이 [MASK] 상태에서 시작해 마스크를 단계적으로 지우며 디노이징을 수행합니다.

일부 구현에서는 중간에 다시 re-masking을 넣어 세밀하게 수정하는 과정을 반복해 “한꺼번에 문장을 그려 넣는다”는 이상적인 생성 방식으로 문장을 생성하게 됩니다.

즉, AR은 한 단계에서 틀린 토큰을 다음 단계에서 되돌릴 수 없지만, Diffusion LM은 문장 전체를 반복 정제하기 때문에,

exposure bias가 줄어들고, context window 내에서 일관성을 확보하기 쉽습니다.

위 내용을 그림으로 좀 더 직관적으로 살펴보겠습니다.

LM에서의 Diffusion

생성 모델에서의 확산 모델은 데이터를 점진적으로 노이즈화한 뒤, 역방향으로 노이즈를 제거하면서 원본을 복원하도록 학습되는 모델입니다.

노이즈 혹은 어떠한 방법으로 원본 데이터를 훼손을 시키고, 다시 훼손시키는 것을 복구하는 방법을 학습해 완전 노이즈인 상태에서 부터 생성을 시작하는 아이디어를 갖고 있고,

두 가지 주요 process가 존재합니다.

Forward process(노이즈 주입, Data to Noise)

• 이미지 같은 연속 데이터에서는 단계 t마다 가우시안 노이즈를 더합니다.

• 텍스트처럼 이산 토큰으로 이루어진 데이터에서는 단계마다 일부 토큰을 [MASK]로 바꿔 노이즈를 만들어줍니다.

• 이 process에서 노이즈를 어떻게 복원시킬 것인가를 학습하게 됩니다. 연속형 데이터의 경우 epsilon을 학습하거나 (stein) score 를 학습합니다.

  사실 epsilon과 score는 수식적으로 큰 차이는 없습니다. 이산 토큰에서는 [Mask]을 다시 어떤 Token으로 예측할 지 학습하게 됩니다.

Reverse process(디노이징, Noise to Data)

• 학습된 모델이 가장 강한 노이즈 상태에서 출발해 단계 번호를 하나씩 줄여 가며 노이즈를 걷어 냈습니다.

• 연속 도메인에서는 스코어 매칭을 통해 노이즈 자체를 예측했고,

  이산 도메인에서는 Masked-LM 손실이나 Rao-Blackwellized Cross-Entropy(RB-CE) 손실을 이용해 [MASK] 위치의 정답 토큰 분포를 맞히도록 학습했습니다.

Rao-Blackwellized Cross-Entropy(RB-CE)

텍스트 Diffusion LM의 학습에서는 [MASK] 토큰이 한 번 마스크되면 다시 원본으로 돌아가지 않는 흡수 상태가 됩니다.

(흡수 상태란, 현재의 상태가 [MASK]가 되면 값이 [MASK]에서 변경되지 않음을 의미합니다.)

이렇게 두 가지 상태(마스크 / 비-마스크)만 고려하면,

원래 Diffusion ELBO에서 요구하던 복잡한 KL 계산 대신 마스크 위치에서 올바른 토큰을 맞히도록 하는 가중 Cross-Entropy 하나로 목적 함수를 단순화할 수 있었습니다.

즉, Diffusion ELBO인

에서 실제 데이터 분포 q를 살펴보면 Mask/Non-mask로 구분될 수 있습니다. 따라서,

으로 정의될 수 있고, [Non-Masking] 부분에서는 t번째 z가 원본 토큰이라면, posterior가 Dirac-delta로 deterministic하게 됩니다.

Rao-Blackwellization 후에는 “현재 토큰이 [MASK]일 때만 1, 아닐 때는 0”인 indicator가 KL 앞 가중치에 곱해지기 때문에, [Non-Masking] 자리의 KL 기여도는 실제로 0이 됩니다.

그러므로 [Masking] 부분만 고려하게 된다면,

Rao-Blackwellized CE으로 변경될 수 있고, 이 손실 함수를 쓰면

• Monte-Carlo 샘플링이 줄어들어 분산이 낮아졌고, (Rao-Blackwell theorem)

• 기존 Masked-LM 코드베이스를 거의 그대로 활용할 수 있고,

• 다양한 노이즈 수준을 고려해 학습했기 때문에 다양하고 일관된 문장을 생성할 수 있습니다.

Pretraining

autoregressive모델이 next token을 예측하는 것이라면, diffusion 모델은 쉽게 말해서 Mask token predictor라고 생각하시면 됩니다.

그림 (a) 에서 전체 Token들 중에서 임의의 확률로 masking을 한 후에, [Mask] 되어진 부분만 Mask Predictor를 통해서 예측하게 됩니다.

[Non-Masking]일 때,

그림 (a)에서 1, 3, 4, 5, 6, 8번째 Token [non-masking]은 흡수 상태이기 때문에 mask token이 아닐 시에 mask token으로 변경되지 않습니다.

그렇기 때문에 Rao-Blackwellization을 통해서 KL앞에 가중치를 0으로 설정해 ELBO에서 기여를 0으로 만들어 현재 Step에서는 학습이 이뤄지지 않게 합니다.

[Masking]일 때,

그림 (a)에서 2, 7번째 Token은 masking상태이기 때문에 현재 step에서 non-mask token을 예측해야 합니다.

결국 모델은 “시퀀스 전체를 인코딩하되, loss/gradient는 [MASK] 위치에서만 계산한다”는 Masked-LM 방식으로 학습됩니다.

SFT

Pretraining이 끝난 후에는 SFT를 진행하게 됩니다. SFT를 진행할 때에는 Prompt 부분은 남겨두고 Response부분만 Masking하여 다시 한 번 학습을 진행하게 됩니다.

학습을 진행할 때 데이터는 기존 autoregressive model과 큰 차이는 없지만, |EOS| token을 다룰 때 다르게 작동합니다.

두 모델 모두 학습을 진행할 때, 동적 시퀀스 길이를 사용합니다.

이 과정에서 AR-based LM은 |EOS| token 이후에 |PAD| token을 주어 Attention score에 0을 줌으로써 Loss/Gradient 계산을 진행하지 않습니다.

하지만, Diffusion-based LM에서는 |EOS| token은 padding의 역할도 함께 하고 있고 Loss 및 Gradient 계산도 함께 진행하게 됩니다.

즉, |EOS| Token도 패딩 및 학습 대상 Token으로 취급함으로써 모델이 스스로 길이를 가변적으로 생성할 수 있도록 설계되었습니다.

Inference

학습이 끝난 이후엔 Response부분만 생성하게 됩니다. 먼저 Prompt 영역을 제외한 모든 부분을 Masking을 하고, 전체 Mask prediction을 진행합니다.

그 이후에 일부 Token을 제외하고 Re-masking을 통해 다시 Mask prediction을 시도합니다. 이러한 Iteration을 통해 Response 전체를 생성하게 된답니다.

각 Iteration마다 정해둔 비율만큼 Remasking을 진행합니다. T=0으로 진행 될 수록, Remasking되는 비율이 작아집니다.

논문에서는 저자들이 Remasking하는 방법들을 크게 2가지를 제안했는데,

첫번째는 랜덤으로 remask되는 token을 정하는 방법이고 두번째는 mask predictor의 low-confidence를 갖는 token들로 정합니다.

하지만 low-confidence remasking 방법에서 instruct모델의 경우 |EOS| token을 과다하게 생성하는 경우가 존재합니다.

이는 큰 문제를 일으키게 되는데, |EOS| token이 만들어지면 대부분 하위 K 번째에서 선택되게 됩니다.

당연하게도, 우리가 답변을 할 때에도 항상 답변의 길이가 같게 답변을 하지는 않죠.

(첫번째 |EOS| 앞의) 길이가 동적이므로 |EOS| token은 low-confidence token에서 항상 선택되게 됩니다.

그렇기 때문에 Instruction model에서는 low-confidence token remasking 전략은 성공적이지 못했습니다.

그렇기 때문에 논문 저자들은 추가적으로 semi-autoregressive remasking방식을 추가적으로 소개하게 됩니다.

그림을 보시면 바로 알겠지만, Block 별로 autoregressive하게 생성하고 각 block 내에서는 diffusion 방식으로 생성되는 방법을 소개하게 됩니다.

결과적으로는 Low-confidence & semi-autoregressive remasking 방식이 가장 성능이 좋게 나온다고 합니다.

AR 방식에 문제점을 얘기하고 다시 AR로 회귀하는 아이러니... 이 방식은 후에 Block Diffusion 이라는 논문이 나오게 되는데,

이 내용은 추후에 다룰 수 있으면 다루도록 하겠습니다.

Result

LLaDA vs. LLaMA

대부분의 결과에서 성능이 AR모델과 비슷할 정도로 학습되었습니다. 그리고 특히 Math영역에서는 LLaDa가 좀 더 강세를 보입니다.

생성되었던 Token을 수정할 수 있으므로 Exposure bias가 줄어들어 수학과 같은 답변 퀄리티에서 강세가 보임으로 생각됩니다.

AR모델의 경우 Left-to-Right 방식이 고정되어야 하기 때문에 이전에 생성된 결과가 뒤쪽에 강력한 bias를 주게 되므로, Step별로 생성하게되는 CoT나 Reasoning을 사용하게 되는데,

Diffusion의 경우 low-conf를 계속 수정해 생성하므로 중간 계산식을 수정하며 답변의 퀄리티와 정확도가 상승되는 것으로 추측됩니다.

LLaDA vs. Others

사실 LLaMA와 비교했을 때에는 성능이 비슷하다고 볼 수 있겠지만은, 실제로는 Qwen이나 Gemma와 비교하게 된다면 성능이 아직까지는 애매한 모습을 보입니다.

아직 초기 단계이기 때문에 연구를 더 진행해볼 여지가 충분히 남아있습니다.

실제로, 구글에서 공개한 Gemini Diffusion의 경우에서 HummanEval이 89.6%, MBPP가 76% 정도로 향상되어 있음을 알 수 있습니다.

Reversal curse

Language model은 왼쪽에서 오른쪽으로 생성되기 때문에, 뒷 문장을 보고 앞 문장을 추측하는 Reversal problem에 취약하게 됩니다.

이를 Reversal curse 라고 부르고, 평가를 위해 Forward (앞 문장을 보고 뒷 문장을 추측)과 Reversal (뒷 문장을 보고 앞 문장을 추측) 두 가지를 측정하게 됩니다.

GPT 마저도 Forward와 Reversal의 차이가 수십점이 나게 되어있습니다.

이는 AR 방식의 학습 방법이 갖는 치명적인 단점으로 지적되는데, LLaDA의 경우 큰 차이 없이 약간의 성능 감소만 있고 Reversal에서도 강인한 결과를 보였습니다.

(물론, Forward에서 성능은 그닥 좋지는 않지만...)

마찬가지로 이는 Diffusion model의 Token의 순서에 제한이 없는 생성 방식이 Reversal generation에서도 global context를 유지하며 생성되기 때문에

Reversal curse에서 해결책으로써 떠오르고 있습니다.

Limitation

아직 초기 논문이다 보니까 여러가지 Limitation이 존재합니다.

우선 컴퓨팅 제약 때문에 LLaDA와 정확히 같은 데이터·스케일로 AR 베이스라인을 학습할 수 없었고, 강화 학습을 통해 모델의 성능을 개선할 수 없었다고 합니다. (LLaDA 1.5)

또한, 모델과 하이퍼파라미터를 최적화 시키지 않았고 아직 멀티모달 혹은 에이전트 등에 통합되는 실험을 하지 않았다고 합니다. (LLaDa-V)

개인적으로는 “AR vs Diffusion”을 양 극단으로 놓기보다는, 두 패러다임의 장단점을 어디에 어떻게 배치할지가 관건이라고 봅니다.

• AR LLM

  • 추론이 빠르고 파인튜닝 자산이 풍부하지만, exposure bias·길이 편향·역순 추론 같은 구조적 약점을 갖습니다.

• Diffusion LM

  • 반복 정제(Self-Refine) 덕분에 장기 의존·논리 일관성에서 강점을 보이지만, 스텝 수가 latency에 직접 반영되고 하이퍼파라미터가 민감합니다.

LLaDA는 Vision뿐만이 아니라 NLP에서도 Diffusion 쪽으로도 접근할 수 있음을 보여줍니다.

물론 D3PM (Discrete Denoising Diffusion Probabilistic Model) 혹은 SEDD와 같이 이전에 사용했던 접근 방법이 있었지만,

처음으로 Large-scale의 모델을 만들어 냈기에 더 의미있는 연구 결과였다고 생각합니다.

이처럼 LLaDA는 AR 기반 언어모델의 구조적 한계를 정면으로 돌파하려는 흥미로운 시도이며, 앞으로의 후속 연구들이 더욱 기대됩니다.

읽어주셔서 감사합니다. 다음 글에서 더 흥미로운 주제로 찾아뵙겠습니다!