본문 바로가기
언어모델

[논문 리뷰 | 정리] Large Language Diffusion Models

by 인공지능과 함께 2025. 2. 23.

 Large Language Models, 줄여서 LLMs는 마치 사람이 생각하는 것처럼, 우리의 질문에 대해 다양한 답변을 내놓습니다. 이런 LLMs는 대체로 대량의 텍스트 데이터로 학습이 됩니다. 이 LLMs를 학습하는 방법에는 다양한 방법들이 존재하지만, 공통적으로는 질문을 입력으로 넣고, 출력으로 답변을 도출하도록 학습이 됩니다.

 

 BERT이나 Gemini 같은 모델의 경우, Mask And Reconstruction이라고 하여 일부러 입력 텍스트 일부를 지우고, 그 지워진 텍스트를 복원하는 방식으로 학습이 진행됩니다. 인간처럼 빈 칸을 뚫어 놓고, 그 빈 칸에 들어갈 단어를 찾는 방식으로 학습시키는 방법입니다. 또 다른 방법으로는 Next Token Prediction 방식이 있습니다. 단어 그대로, 주어진 문장 다음에 올 단어(토큰)를 예측하는 방식입니다. 대표적으로 GPT 같은 Autoregressive 모델이 Next Token Prediction 방식으로 학습되었습니다. 사람이 글을 써 내려가는 것 같은 방식과 유사합니다. 

 

 이렇듯 언어 모델들은 사람이 글을 배우는 방식과 유사하게 텍스트를 학습합니다. 각각의 학습 방법에 영향을 받아 추론 단계에서도 위의 모델들은 다른 방식으로 출력을 생성해 냅니다.

 

 다른 방식으로 답변을 생성하는 방법은?

 

(이번 포스트는 조금 길어질 것 같다는 것을 미리 알려드립니다.)

 


 

LLaDA, Diffusion 방식의 언어 모델?

 Preprint: https://arxiv.org/pdf/2502.09992

 

 인공지능 분야를 조금 공부하신 분들이라면, 생성이라는 단어에 적합한 모델을 떠올리실 겁니다. 바로 Diffusion 기반의 이미지 생성 모델인 Stable Diffusion을 말이죠. GAN 이후로 Diffusion 기반 모델은 뛰어난 이품질의 이미지를 생성할 수 있어 지금까지도 각광 받고 있는 구조이며, 그 변형들이 많이 나와 있습니다. 

 

 

 

 

Diffusion이 뭔지부터 알고 가자...

 LLaDA는 이런 이미지 분야의 Diffusion 방식을 언어 분야로 끌어들인 모델이라고 할 수 있습니다. 이 Diffusion 방식은, 미지의 값에서 출력 1을 생성하고, 그 출력 1을 다시 모델에 넣어 출력 2를 생성하고, ... , 출력 N-1을 모델에 다시 넣어 최종 출력인 출력 N을 만들어 냅니다. 그림으로 보면 다음과 같습니다.

 

Diffusion Denoising Model

최초의 렌덤 값(정규분포)에서 차근차근 출력을 다시 입력으로 넣어 최종 출력을 만들어 내게 되는 방식이죠. 보통 Diffusion Denoising Model이라고 부릅니다.

 

 

알아야 할 게 또 있어... 바로 Autoregressive...

LLaDA를 본격적으로 설명하기 전, preprint에 Autoregressive 방식이 있어 잠만 수식만 알아보고 가겠습니다. 이 Autoregressive 방식 혹은 Next Token Prediction 방식으로 알려진 것은 GPT의 방식으로 유명합니다. 이전 토큰들을 기반으로, 다음 토큰으로 뭐가 올지에 대한 확률을 구합니다. 그리고 그 확률이 가장 높은 단어(토큰)을 다음 토큰으로 지정하는 방식으로 전체 문장을 완성합니다. 수식으로는 다음과 같습니다.

 

$$p_\theta(x) = p_\theta(x^1)\prod_{i=2}^L p_\theta(x^i | x^1, \ldots, x^{i-1})$$

 

$x^i$는 $i$번째 토큰 값.

$L$은 예측하려는 토큰의 길이

$p_\theta$는 인공지능 모델. (파라미터가 $\theta$)

$x$는 길이 $L$ 만큼의 문장. 즉, $x^1, x^2, \ldots, x^L$의 긴 토큰들을 $x$라고 간단하게 표현한 것.

 

수식이 많이 어려워 보입니다. 설명을 자세히 해드리고 싶지만, 이 포스트는 Autoregressive 방식을 설명하는 포스트는 아니기 때문에 아래의 더보기를 통해 수식을 이해하시면 될 것 같습니다.

 

더보기

"나는 아침을 먹는다"이라는 문장은 [나, 는, 아침, 을, 먹, 는다] 총 6개로 분할할 수 있습니다. 이를 토큰으로 표현하면, $x^1$는 "나"에 해당하는 토큰이 되고, $x^6$는 "는다"에 해당되겠죠. 그러면, "나는 아침을 먹는다"는 $x^1x^2x^3x^4x^5x^6$로 표현됩니다.

 

이 문장을 확률로 표현한다면, $p(x^1, x^2, x^3, x^4, x^5, x^6)$와 같은 결합확률분포로 표현할 수 있습니다. Autoregressive 방식 혹은 Next Token Prediction 방식은, $x^1$부터 $x^i$까지 주어졌을 때, $x^{i+1}$을 예측하는 방식입니다. 즉, "나는 아침"이라는 문장이 주어지면, "을" 을 예측하여 "나는 아침을"을 만들어내는 방식이죠. 

 

'주어지면'이라는 것을 수학적으로 표현하면, "조건부"라는 단어로 생각할 수 있습니다. 여기서 바로 "조건부확률"을 떠올리셨다면, 훌륭한 분이신 것 같습니다. "나는 아침"이라는 문장이 주어졌을 때, "을"의 확률분포는 다음과 같습니다.

 

$$p(x^4 | x^1, x^2, x^3)$$

 

비슷한 방식으로 다음을 생각해 볼 수 있습니다.

"나"가 주어졌을 때 "는"의 확률분포: $p(x^2 | x^1)$
"나는"이 주어졌을 때 "아침"의 확률분포: $p(x^3 | x^1,x^2)$
"나는 아침"가 주어졌을 때 "을"의 확률분포: $ p(x^4 | x^1, x^2, x^3) $
"나는 아침을"가 주어졌을 때 "먹"의 확률분포: $p(x^5 | x^1, x^2, x^3, x^4) $
"나는 아침을 먹"이 주어졌을 때 "는다"의 확률분포: $p(x^6 | x^1, x^2, x^3, x^4, x^5) $

 

자 그럼 여기까지는 다음 한 개의 토큰의 확률분포를 조건부 확률로 구하는 방법을 알아보았습니다. 그러면, 하나의 토큰이 아닌, 문장(여러 토큰)의 확률분포는 어떻게 구할까요? 즉, "나는"이라는 "나"와 "는"으로 이뤄진 문장의 확률은 다음과 표현할 수 있습니다.

 

$$p(x^1, x^2)$$

 

그리고 이건 조건부확률의 특성을 이용하여 $p(x^1)$이 주어지면, 다음과 같이 구할 수 있습니다.


$$p(x^1, x^2) = p(x^1) \times p(x^2 | x^1)$$

 

마찬가지로, "나는 아침"이라는 문장도 비슷한 방식으로 구할 수 있습니다.

$$\begin{align} p(x^1, x^2, x^3) &= p(x^1, x^2) \times p(x^3 | x^1, x^2) \\&= p(x^1) \times p(x^2 | x^1)\times p(x^3 | x^1, x^2)\end{align}$$

 

전부 [ 결합확률분포 $(_1, \ldots, _i)$ = 결합확률분포 $(_1, \ldots, _{i-1})$ x 조건부확률 ] 형식인 것을 알 수 있습니다.

 

저런 방식으로 "나는 아침을 먹는다"는 $p(x^1, x^2, x^3, x^4, x^5, x^6)$로 표현되고, 이것을 곱 연산인 $\prod$로 표현하면, 가장 윗줄에서 본 수식이 나오는 것입니다. ($x = [x^1, x^2, x^3, x^4, x^5, x^6]$)

 

 

 

 

LLaDA

 LLaDA는 바로 이 Diffusion 방식을 차용하여 출력을 만드는 모델입니다. 구체적으로는 Mask and Reconstruction 방식과 Diffusion 방식을 결합한 Transformer 구조 기반의 Masked Diffusion Model (MDM)입니다. (이 부분에서부터 LLaDA가 특별한 모델이라고 생각되지는 않았습니다.)

 

 여기서 이미지 계열의 Diffusion 모델과 차이점이 존재합니다. 바로, noise의 유무입니다. 이미지 계열의 Diffusion 모델을 아시는 분들은, Diffusion 모델이 $noise \rightarrow \cdots \rightarrow x_0$의 형태로 최종 출력을 만든다는 것을 알고 계실 겁니다. 즉, noise에서 순차적으로 출력을 만들고 그걸 다시 입력으로 받으며 최종 출력을 만들게 되는 것이죠.

 

 하지만, LLaDA는 random noise가 아닌, [MASK]라는 임의의 $L$개의 토큰으로부터 순차적으로 일부 토큰들을 복원(출력)하고, 복원된 토큰들과 그렇지 않은 [MASK] 토큰들을 다시 모델에 넣고 또 일부 토큰들을 복원합니다. 이러한 방식으로 최종 토큰들을 복원하여 완성된 문장을 만들게 됩니다.

 

 


 

좀 더 구체적인 학습 방식으로.

Pretraining

 앞에서는 LLaDA의 전체적인 설명을 했다면, 여기선 좀 더 구체적으로 학습이 어떻게 진행됐는지 알아보겠습니다. 학습 방식을 알기 위해서 일부 몇 가지 수학 기호들에 대해 정의하고 가겠습니다.

 

$t$: $0$에서 $1$ 사이의 값을 가지는 값으로, Diffusion step을 의미합니다. $t$가 0이면, 온전한 문장일 때, $t$가 1이면, 모든 토큰들이 [MASK]될 때입니다.

$L$: 처리하려는 문장의 길이.

$i$: 토큰의 위치. 예) $i$가 1이면, 첫 번째 토큰의 의미.

$x_t^i$: $t$ 시점의 Diffusion step의 $l$ 번째 토큰.

$x_t$: $t$ 시점의 Diffusion step의 출력 토큰들.

$p_{\theta}$: 인공지능 모델 (파라미터가 $\theta$)

$M$: [MASK] 토큰의 수학적 표현

$N$: 총 Diffusion step의 수.

 

좀 많고 어렵네요. 앞으로 나올 수식들이 더 어려울 수 있다고 생각해서 아래와 같이 그림을 그려보았습니다. 

 

 

Masked Diffusion Model (MDM) 방식

 

그리고 이게 LLaDA의 pretraining 방식입니다. 화살표 방향으로 모델이 학습하며, 그림에서 직관적으로 알 수 있듯이, 모델은 $M$이라고 되어있는 일부분을 추론합니다. (그림으로는 한 step에 하나씩 추론하는 것처럼 그려져 있지만, 실제로는 한 번에 여러 $M$을 추론합니다.) 이런 방식으로, LLaDA는 이전 step에 있는 문장을 통째로 받고, 그 다음 문장을 추론하도록 학습됩니다.  

 

 

 

 

수학적으로 표현하면, 임의이 Diffusion step인 $t$에서의 문장 $x_{t+\frac{1}{N}$를 받았을 때, $x_{t}$을 예측합니다(여기서 $N$은 총 Diffusion step 개수입니다). 학습 관점에서 말하자면, $x_{t+\frac{1}{N}$ 문장을 받고 바로 그 다음 스탭의 문장인 $x_t$를 예측한다고 할 때, 이 문장 속 $M$을, 최종 정답인 $x_0$에서 동일한 위치의 토큰과 일치하는 방향으로 pretrain이 진행됩니다. 

 

Preprint에는 다음과 같이 로스의 형태로 정의돼 있습니다. (수식 3) 이 로스를 사용하기 위해서는 우선, 최종 정답지인 $x_0$와 $t$ 시점 그리고 그 $t$ 시점에 대한 모델 예측 값 

 

$$\mathcal{L(\theta)}\triangleq -\mathbb{E}_{t,x_0,x_t}\left[\dfrac{1}{t}\sum_{i=1}^L\mathbb{1}[x_t^i=M]\text{log} p_{\theta}(x_0^i | x_t)\right]$$

 

수식이 참 어려워 보이네요. 차근차근 뜯어보면, 제가 했던 말과 일치하는 수식입니다.

 

  • 우선, $i=1$부터 $L$까지 $\sum$기호의 의미는, $i$ 번째 토큰에서부터 $L$ 번째 토큰까지 일단 모두 검사하겠다는 의미입니다.
  • 다음으로, $\mathbb{1}[x_t^i=M]$의 의미는, "$t$ step의 $i$ 번째 토큰인 $x_t^i$가 $M$ ([MASK]일 때!"라는 조건부 수식입니다.
  • $p_\theta(x_0^i|x_t)$는 $x_t$라는 문장이 주어졌을 때, masking된 토큰($x_t^i$)와 동일한 위치($i$)에 있는 정답 $x_0^i$ 값의 확률이라는 의미입니다.
  • 여기에 Expectation을 취하고, 음수를 취함으로써, Negative Log Likelihood (NLL) loss가 완성되는 것입니다.

 

사실 중요한 수식은 저게 다인 것 같습니다. 나머지는 masking된 문장을 만들 때 (forward pass) $M$의 비율에 대한 수식입니다. 이 부분은 preprint에도 설명 돼 있으니 쉽게 알 수 있을 것입니다. 간단한 특징으로는, $t=1$에 가까울 수록 $M$의 비율을 올리고, $t=0$에 가까울 수록 낮추는 것이네요. 추가로 저 수식은 Supervised Fine-Tuning (SFT)에서도 거의 동일한 형태로 나오게 됩니다.

 

 

Supervised Fine-Tuning

 Pretrain 단계에서는, 단순히 한 문장 내에서 토큰들을 복원(reconstruction)하는 방법을 배웠다면, 이제는 실제 환경처럼 [질문] -> [답변] 형태로 학습하게 됩니다. 이런 관계 떄문에 새로운 수학 기호들이 더 등장합니다.

 

$p$: 질문 (prompt)

$r$: 답변 또는 예측값 (response)

 

Supervised Fine-Tuning도 Pretrain과 마찬가지로 그림으로 그려보았습니다.

Superised Fine-Tuning Mask and Reconstruction

 

학습 과정은 pretrain 때와 유사하며, 차이점은 [질문] 또한 모델의 입력으로 들어간다는 점입니다. 수식으로 로스를 표현하면 다음과 같습니다. 수식 (6)

 

$$\mathcal{L(\theta)}\triangleq -\mathbb{E}_{t,p_0,r0,r_t}\left[\dfrac{1}{t}\sum_{i=1}^{L'}\mathbb{1}[r_t^i=M]\text{log} p_{\theta}(r_0^i | p_0, r_t)\right]$$  

 

다를 바 없습니다. 단지 수학 표기법이 바뀌고, $p_{\theta}$ 부분에 [질문]에 해당되는 $p_0$가 조건부로 들어가 있습니다. Supervised Fine-Tuning 끝.   

 

 

 

Inference    

 LLaDA가 추론하는 과정은 사실 Supervised Fine-Tuning과 유사합니다. LLaDA는 이미지 계열에서 사용하는 Diffusion 모델과 비슷하게 한 스텝 한 스텝씩 추론 하면서 최종 출력을 만들어 냅니다. 보통은 많은 스텝을 거치면 거칠 수록 최종 출력물의 질이 높아집니다(Trade-off). 하지만, LLaDA는 기존 Diffusion과는 조금 다른 방식으로 한 스텝 한 스텝 추론합니다. 보통의 Diffusion 모델들은 바로 다음 스텝의 출력물을 출력하도록 돼 있는 반면, LLaDA는 한 번에 모든 토큰을 바로 추론합니다.

 

 하지만 이러면, LLaDA가 학습 과정에서 사용한 방식과 추론 과정이 달라지게 됩니다. 그래서 이를 맞춰주기 위해 remasking이라는 기법을 활용합니다. 즉, 한 스템에서 추론된 토큰들 중, 일부를 다시 마스킹하고, 이걸 기반으로 다시 다음 스텝의 출력을 만들어 내는 방식입니다. Remasking의 비율은 스텝이 지나면 지날 수록 작아지며, 결국 최종 스텝에 가까워질 수록 remasking의 비율은 0에 수렴합니다. 수식으로 쓰면 다음과 같습니다.

 

$$\dfrac{\text{남은 스텝}}{\text{현재까지의 스텝}}$$

 

여기서 더 나아가 LLaDA는 semi-autoregressive라는 방식으로도 실험을 진행합니다. semi-autoregressive 방식은, remasking과 같이 사용되는 방식으로 autoregressive 방식으로 모델이 출력을 생성하도록 마스크를 씌우는 방식입니다. 전체 토큰를 block이라는 임의의 단위로 묶은 뒤, 아래와 같이 block 단위로 remasking을 진행합니다.

 

 

 

Inference 단계의 수식은 Supervised Fine-Tuning과 유사합니다. 하지만, 현재까지 공개된 Preprint에는 오타인지 아니면 오류인지는 모르겠으나 수식에 이상한 부분이 있어 inference 수식은 넘어가도록 하겠습니다. (수식 6)

 

 


 

 

실험

Scalability?

 LLaDA가 autoregressive 모델보다 scalability, 즉 모델이 커질 수록 성능도 같이 올라가는 성질이 더 좋다고는 합니다. Preprint의 그림 3이 scalability에 관한 것인데, GSM8K 데이터셋 외에는 autoregressive 모델과 유사한 것으로 보입니다. (주관적인 생각입니다.)

 

특징

Reversal curse

 Reversal curse 문제는 LLaDA preprint 첫 페이지에 등장할 정도록 강조되는 LLaDA의 특징입니다. 출처에 따르면, Reversal curse 문제는 "A는 B다"라는 문장을 모델이 학습했을 때, "B는 A다"라는 문장은 이해하지 못하는 현상입니다. Autoregressive 모델의 경우, 순차적으로 토큰들을 처리하기 때문에 이러한 reversal curse 문제가 발생한다고 합니다.

 

 LLaDA는 autoregressive 모델보다 성능이 좋은 것으로 나옵니다. Preprint의 Table 3를 기준으로, Qwen2.5 7B 모델과 GPT-4o(2024-08-06)과 비교했을 때, 다음과 같습니다.

  Forward Reversal
GPT-4o 82.7 34.3
Qwen2.5 7B 75.9 38.0
LLaDA 8B 48.8 42.4

 

Reversal 문제에서는 다른 두 모델보다 성능이 좋은 것으로 나왔지만, forward에서는 성능이 현저히 낮다는 큰 단점이 존재합니다. 이게 LLaDA의 큰 단점 중 하나라고 생각할 수도 있습니다.

 

성능

 LLaDA의 경우, LLaMA2, Qwen2, Mistral, 그리고 Deepseek 모델과 비교하여 성능을 보였습니다. 다들 눈치채셨겠지만, 여러 데이터셋 (benchmarks)에서 네 개의 모델보다 성능이 좋다고 나왔습니다. 하지만, 중간 중간 다른 모델들에 대한 성능이 안 나왔다는 점에서 조금 아쉬움이 남았습니다. 또한 (공정한 비교는 아니겠지만) reversal curse 문제에서 언급된 GPT-4o 모델과의 비교도 없다는 것이 아쉬운 부분이었습니다. 

 

 

 

 

 

 

 


 

 

 

 

마무리

 Preprint라서 그런지 아직 부분 부분 완벽하지 못한 것도 있습니다. LLaDA는 Diffusion 방식을 언어에도 적용시켜본 점, reversal curse 문제를 어느정도 완화한 점에서 좋은 시도라고 생각됩니다. 하지만, LLaDA는 reversal 문제 외에 Forward 문제에서 단점을 보인다는 점, 일반적으로 Diffusion 모델들은 많은 양의 데이터와 컴퓨팅 자원이 필요하다는 점에서 아직 더 연구할 부분이 남는다고 생각됩니다.

 

 끝으로 Appendix쪽에 재밌는 그림이 있어 그림에 대한 설명과 함께 마무리합니다.

LLaDA가 각 step에서 생성한 토큰들. 진항 색상일 수록 나중에 생성되었다는 의미. 윗쪽은 semi-autoregressive 방식으로, 문장에 앞단에 있는 토큰들이 상재적으로 먼저 생성됐다는 것을 확인할 수 있다. 이 색상들은 remasking 시점에 영향을 받는다.