Jae-Kyung Cho LLM Developers who was a Robotics engineer

Diary - Diffusion language model과 LLM reversal curse

24 Nov 2025

몇 달 전에 Gemini Diffusion 모델이라는 게 나왔다는 뉴스를 봤다. 일단 첫 번째로 Diffusion으로 어떻게 language model을 만든거지? 가 궁금했고, 두 번째로는 데모 비디오를 봤는데 미친듯이 빠른 속도가 눈에 띄었다. (꼭 데모 비디오 보시길!)

그래서 한번 근간이 되는 논문을 읽어보기로 했다. Alibaba Ant group 에서 써낸 Large Language Diffusion Models, 일명 LLaDA 라는 모델을 제안한 논문이 되겠다 (근데 마지막 A가 제목에는 없음…Large Language diffusion with mAsking 라고 함)

꽤 재미있게 읽은데다가, 중간에 LLM reversal curse 라는 개념이 나오는데, 굉장히 흥미로웠다. 그리고 관련하여 팀 동료분과 나눴던 토론도 재미있었기 때문에 글로 남겨 본다.

LLaDA란 무엇인가?

그럼 LLaDA가 어떤 모델인지, 어떤 개념으로 diffusion을 이용해 language model을 학습시켰는지 간단히 알아보자. 우선 논문의 저자들은 대형 언어 모델의 핵심 능력(scalability, in-context learning, instruction following)이 auto-regressive 구조에만 국한되지 않으며, 디퓨전 모델을 통해서도 달성 가능하다고 주장한다.

LLM 하던 사람들은 아래와 같은 auto-regressive 형태의 확률 계산식에 익숙하다.
Image
이전 토큰 sequence의 확률을 가지고 다음 토큰의 확률을 예측한다는 건데, 이걸 가지고 negative log likelihood (NLL) 를 최소화시키면 언어모델이 학습된다.
Image

LLaDA는 아래의 NLL 최소화 식을 만족시키기 위해서 아래와 같은 surrogate loss function 을 제안한다.
Image

이게 무슨 뜻이냐면 \(x_t\) 는 t의 비율로 masked 된 text sequence를 말한다. \(x_0\)는 아예 masking 되지 않은 text sequence, 즉 ground-truth가 되겠다. 결국 위 loss function은 masked text sequence 가 주어졌을 때, 전체 ground-truth sequence 를 예측하는 것이고, masked token 에 대해서만 loss 를 계산하겠다는 것이다. (\(1\left[x_t^i=M\right]\) 이 의미하는 것)

저자들은 이게 diffusion 의 개념과 같다고 본 것이다. text sequence 에서 token들을 masking 하는 것을 noise 를 추가하는 것, masked 된 token sequence 에서 text sequence 를 만들어 내는 것을 de-noising 으로 본 것이다.

이 때 하나 짚고 넘어가야 할 점이 있다. 이 LLaDA loss를 사용해서 NLL loss 를 최소화시킬 수 있냐는 건데, 이건 DeepMind에서 이미 증명을 해줬다. 근데 좀 어려우니까 원하는 사람만 읽어보자 (나도 읽다 말았다 ㅎㅎ)

그림으로 살펴보는 LLaDA

지금까지 수식으로 LLaDA에 대해 살펴봤는데, 당연하게도 이해가 잘 안된다. 그래서 뭘 어떻게 학습하고 어떻게 쓴다는건지? 친절하게도 저자들이 그림을 매우 잘 그려놨다.

LLaDA pre-training

아무튼 이것도 language model이다 보니 pre-training과 post-training 단계가 있다. Post-training 단계부터 살펴보도록 하자.

Image

LLaDA pre-training 방법

  1. t ~ U(0,1) 로 샘플링
  2. t 만큼의 token을 masking
  3. masked token prediction loss 적용 (위에 LLADA loss function)

매우 단순하죠? 이 그림에서 알 수 있듯이 LLaDA는 기존 auto-regressive LM에 비해 엄청나게 큰 장점을 하나 가진다. 바로 전체 sequence 에 대한 attention 을 얻을 수 있다는 것이다. Masked token predictor 는 전체 sequence의 token으로부터 attention을 계산하여 토큰을 예측하도록 구성된다. 즉, 나보다 뒤에 있는 토큰으로부터 정보를 얻을 수 있다는 것이다.

(🤪읽거나말거나) 논문의 디테일한 점들:

  • LLaMA3와 비슷한 아키텍쳐를 사용함. 다만 편의를 위해 GQA 안쓰고 MHA 사용하였음.
  • Sequence length 는 4K를 사용. 그러나 1% 데이터 정도는 [1,4096] 범위에서 랜덤하게 사용해 various sequence length에 대한 일반화 성능 확보
  • 2.3T token으로 8B급 모델 학습

LLaDA post-training

비슷한 방식을 이용해 post-training (instruction tuning)까지 수행해주면, 실제로 사용할 수 있는 모델이 된다.
Image
Instruct tuning은 그냥 LLM 학습하던거랑 비슷하게 prompt는 건드리지 않고 response 쪽만 masking 해서 학습한다. 또 chat template을 적용해서 포맷팅을 맞춰주면 된다.

여기서 한 가지 짚고 넘어가야 할 것이 있는데, LLaDA는 치명적인 단점이 하나 있다. 바로 inference시 sequence length를 사전에 고정해야 한다는 것이다. 그래야 denoising 이 가능하다. Pre-training이야 corpus니까 내가 원하는 길이대로 자른다고 치는데, post-training은 prompt-response pair의 길이가 천차만별이다. 따라서 padding을 적용해야하는데, LLaDA는 여기서 한 가지 트릭을 사용한다. 바로 padding token 대신 EOS (End-of-Sequence) token을 사용하는 것.

기존 auto-regressive LM은 padding token에 대해서는 loss를 계산하지 않지만 LLaDA는 EOS 토큰 전체에 loss 를 계산한다. 즉 마지막에 EOS를 얼마나 생성하느냐가 response length를 결정짓는 것이다.

LLaDA post-training 방법

  1. Response 부분만 masking 을 적용
  2. Prompt-response 의 길이에다가 <|EOS|> 를 붙여 sequence length 를 맞춰줌
  3. 나머지는 pre-training과 동일

(🤪읽거나말거나) 논문의 디테일한 점들:

  • 450만개 데이터 샘플 사용

LLaDA inference

이제 가장 중요한 부분이다. 그래서 대체 어떻게 쓰는데?
Image

Basic Inference

  1. Response 가 전부 masked 된 상태에서 시작. De-masking (sampling) step을 지정
  2. Response 전체 token prediction
  3. 랜덤하게 일부 token을 다시 re-masking (이전보다는 적게)
  4. 2~3 프로세스 반복

놀랍게도 토큰을 하나씩 생성하는게 아니라, 전체 토큰을 한 번에 다 생성하는 것이다. 그러고 나서 일부를 다시 masking 한다. 그러니까 매 step 마다 response 영역의 모든 토큰이 계속 재생성되기를 반복한다.
이 방식은 de-masking step 을 sequence length 와 동일하게 두면 토큰을 하나씩 생성하는 것처럼 보이게 된다. de-masking step 을 적절히 높게 설정하면 퀄리티가 올라가고, 낮게 설정하면 inference throughput이 올라간다.

이제 핵심은 어떤 토큰을 다시 masking 할 것인지 정하는 일이다. 기본적으로는 랜덤하게 골라서 masking을 하지만 더 좋은 방법들을 저자들은 제시하고 있다. 우선은 확률값이 낮았던 토큰들을 다시 masking 하는 전략이다.

Low-confidence inference

  1. Response 가 전부 masked 된 상태에서 시작. De-masking (sampling) step을 지정
  2. Response 전체 token prediction
  3. Low-confidence (=low-probability) token을 다시 re-masking (이전보다는 적게)
  4. 2~3 프로세스 반복

그런데 이렇게 하면 재밌게도 EOS 들이 가장 먼저 생성된다. 왜냐면 post-training에서 가장 많이 본 토큰이 EOS이기 때문이다 (하나의 샘플에 여러개가 들어가 있으니 절대적인 학습량이 많다) 그래서 제대로 결과가 안나온다. 그래서 아래와 같이 새로운 방법을 제시했다.

Image

Semi-autoregressive inference

  1. Response 가 전부 masked 된 상태에서 시작. De-masking (sampling) step을 지정
  2. Response 를 sequence block 으로 나눔
  3. 앞쪽 block 부터 전체 token prediction
  4. 해당 블록에 대해 Low-confidence token을 다시 re-masking (이전보다는 적게)
  5. block 의 token prediction 이 끝나면 다음 block 으로 넘어가 3~4 반복

response 전체를 한방에 생성하는게 아니라, 블럭을 나눠서 앞쪽 블럭부터 생성하겠다는 것이다. 이러면 EOS가 우선 생성되는 현상을 막을 수 있다.

LLaDA가 그래서 얼마나 좋은데? (feat. reversal curse)

뭐 당연하게도 논문에는 LLaDA가 좋다는 벤치마크들을 잔뜩 써놨다. 엄청 눈에 띄는 성능은 아니었기 때문에 따로 기록하지는 않겠다 (궁금한 사람들은 직접 찾아보길 바란다) 대신 엄청 재미있는 현상인 LLM reversal curse를 LLaDA로 해결했다는 주장이 있어, 이 내용을 다뤄보겠다.

Reversal curse

저자들은 Poem completion 이라는 task를 제안한다. 496개의 2문장으로 이루어진 아주 유명한 중국 poem 데이터가 있는데, forward task는 앞 문장을 주고 뒷 문장을 맞추는 문제이고, reversal task는 뒷 문장을 주고 앞 문장을 맞추는 문제다. 이 496개의 시들은 진짜 너무너무 유명한 것들이어서, 이 모델의 학습데이터에 무조건 포함되어 있다. 즉, 이 496개의 시에 대한 지식을 이미 학습한 모델에 대해서 forward/reversal task의 성능을 측정한 것이다.

결과는 아래와 같다.
Image

Auto-regressive LM은 Reversal 성능이 매우 떨어지는데 비해, LLaDA는 두 task 성능이 비슷하다는 것이다. 같은 데이터를 넣더라도, causal mask를 적용해 omni-directional 하게 텍스트를 학습한 LLM은 reversal task를 제대로 풀지 못한다는 것이다. 반면 bi-directional 하게 텍스트를 학습한 LLaDA는 forward task와 reversal task의 성능 격차가 작다!! (물론 애초에 성능이 떨어지긴 하지만 저자들은 해당 부분에 대해서는 언급하지 않았다…)

이 부분에 대해서 팀 동료분과 나눴던 이야기가 굉장히 흥미롭다.
나의 주장은

LLM이 지식을 습득했는데, 방향에 따라서 답변을 할 수 있고 없고가 나뉘는 것은 진정한 지능이 아니다. 따라서 reversal curse 실험을 통해서, auto-regressive LM은 그냥 지능을 가지고 있는 척하는 것일 뿐이고 사실은 통계적 앵무새에 불과하다는 것을 알 수 있다.

팀 동료분의 주장은

우리는 영어 알파벳을 다 알고 있다. 그런데 거꾸로 외워보라고 하면 바로 말 할 수 있는 사람이 몇이나 될까? 지식의 정의에 따라 다르겠지만, 무언가를 안다는 것에 방향성이 있는 것은 이상하지 않다. 애초에 거꾸로 외우는 것은 해본적이 없기 때문에 어려운 것이고, 알고 있기는 하기 때문에 적절한 추론을 해서 이끌어내면 된다.

결국 설득되어 버렸다!!

📌 사실 이건 Code 도메인의 FIM (Fill-in-the-middle) task와 유사하다

Efficient training of language models to fill in the middle 논문에서 제안
fim.png

  • 코드의 앞부분과 뒷부분을 주고 가운데를 채우는 문제를 일반적인 LLM은 잘 풀지 못함
  • Code 쪽에서도 FIM 문제를 해결하기 위해서 별도 데이터를 부어서 Coder 모델을 만들었던 상황
  • FIM 데이터를 넣어서 해결할 수는 있는데, 엄청나게 많이 넣어야 함
  • 따라서 LLADA는 코드 특화 모델로서 가능성이 있음 (Gemini diffusion 이 주목하고 있는 파트이기도 함)

LLaDA의 한계

LLaDA는 사실 아직 단점이 더 많다. 그렇기 때문에 Gemini diffusion도 웹페이지에 벤치마크랑 데모 비디오만 올려놓고, 아직 공개를 못한게 아닐까 싶다.

추론의 비효율성

우선 auto-regressive LM은 causal mask 에 기반한 다양한 알고리즘과 kernel들이 개발되어 있다. KV-cache도 그 중 하나다. KV-cache는 연산을 sequence length에 linear 한 스케일로 만들어준 아주 중요한 개념이다. 하지만 LLaDA는 매번 전체 sequence 에 대한 attention을 다시 계산해야 해서 KV-cache를 사용할 수 없다.

간단하게 time-complexity를 계산해 보면 그 차이를 알 수 있다.

  • C : prompt 길이, L : response 길이, d : hidden dimension, T : sampling step
    • LLADA: \(O(T\cdot[(C+L)^2d + (C+L)d])\)
    • ARM (w/o KV-cache): \(O(\sum_{i=1}^{L}[(C+i)^2d + (C+i)d])\)
    • ARM (w/ KV-cache): \(O(L\cdot[d^2+Cd])\)
  • C=1024, L=256, d=4096, T=32 넣어보면
    • LLADA: 9E+11
    • ARM (w/o KV-cache): 6.35E+12
    • ARM (w/ KV-cache): 5E+9

아쉽게도 LLaDA는 효율적인 attention 알고리즘이 아직 개발되지 않았다. 개발된다면 auto-regressive LM과 대적해볼만 할텐데 아쉽다.

대형 모델 검증

대부분의 아카데미 페이퍼가 그렇듯 대형모델에서 LLaDA는 아직 검증되지 않았다. 논문에서 사용한 가장 큰 모델이 7B인데, 이건 SOTA 모델들과 비교하면 너무 작다. 얼른 Google이 Gemini diffusion을 공개해주길 기다릴 뿐이다 https://deepmind.google/models/gemini-diffusion/

구글이 적어둔 속도만 보자면 어마무시하기는 하다.

  • Gemini diffusion 은 1479 tokens/sec (32k length generation 기준)
  • Gemini 2.0 flash lite 는 181 tokens/sec

참고로 Qwen3-0.6B 가 generation throughput이 414.17 tokens/secs (sglang 기준) 이다. 이정도면 그냥 넣자마자 inference 끝나는 수준이다.

실사용 소감 끄적끄적

  • Sampling step 을 너무 줄이면 난리법석임. 특히 semi-autoregressive mode일 때.
  • 근데, 아무튼 속도를 사용자가 조절할 수 있다는 점은 꽤 흥미로운 것
  • 생각보다 답변의 퀄리티가 좋아서 놀랐음.
  • 코드 쪽에서는 진짜 유용하게 쓰일지도 모르겠음
Comments

Related posts