[논문리뷰] Simple and Effective Masked Diffusion Language Models

Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J., ... & Kuleshov, V. (2024). Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37, 130136-130184.

https://s-sahoo.com/mdlm/

MDLM Blog post

 

한동안 글을 작성하지 않다가 폭풍처럼 Diffusion Language Model (DLM)을 살펴 보고 있는데요.

수많은 수식들 때문에 머리가 터지기 일보 직전입니다.

그래서 관련 논문 리뷰들은 그런 수식 없이 비교적 단순하게 그냥 제가 이해한 내용을 작성하고 있어요.

(절대 수식을 옮겨 적기 귀찮아서 그런게.. 그런 것도 맞습니다 😎)

 

원래 저는 인트로부터 차례 차례 리뷰하는걸 선호하긴 하는데,

이번에는 본론으로 바로 넘어가볼게요.

여튼 이 논문은 continuous가 아닌 discrete과 관련된 확산 모델입니다!

 

3 Simple Masked Diffusion Models

확산 모델은 이미 여러번 언급했지만 정방향 (노이즈를 추가하는 과정)과 역방향 (노이즈를 제거하는 과정)을 통해 학습이 이루어집니다.

근데 원래 확산 모델은 연속형 데이터를 다루기 위해 설계된 모델입니다.

확산 모델이 성공적인 성능을 보여줬던 이미지, 오디오 분야는 모두 연속형 데이터라는 특성을 가지고 있지요.

근데 우리 저자들은 굳이 굳이 이산형 데이터에 이 모델을 적용해보고 싶은 겁니다. 

그래서 연속형 데이터를 요하는 모델을 어떻게 이산형 데이터 (텍스트)에 적용할까?

이런 고민들 그리고 그걸 해결하는 과정에서 정말 다양한 논문들이 생산되는 거죠.

 

 

3.1 Interpolating Discrete Diffusion

정방향 과정에 대해 설명합니다. 즉, 깨끗한 데이터에 노이즈를 입히는 과정이죠. 앞서 말했듯이 연속형 데이터라면 단순하게 가우시안 분포를 활용해서 랜덤 노이즈를 각 데이터에 더해주면 되겠지만, 텍스트는 이산형 분포이기 때문에 가우시안 분포를 활용할 수 없습니다. 대신 확률 분포의 혼합을 사용해서 노이즈 상태를 구현합니다. (원본 데이터의 확률과 노이즈 분포의 확률을 가중합으로 결합하는 방식이에요!)

 

이렇게 얘기하면 전 절대로 이해하지 못합니다. 

머리에 물음표가 한 백만개 뜬 것 같아요.

 

예시로 "과일 가게에서 (토마토)를 하나 샀어." 라는 문장이 있다고 해봅시다.

여기서 우리가 예측해야 하는 단어가 토마토이고, 선택지는 다음 5개입니다.

[당근, 토마토, 토끼, 꿀, 벌]

 

그럼 이때 확률을 극단적으로 [0, 1, 0, 0, 0]이라고 가정합시다.

이런 벡터가 주어지면 100% 확률로 토마토를 선택하게 되겠죠.

 

노이즈 분포는 어떤 단어를 고르지 못하게 에러항을 더하는 개념이니, 모든 단어가 동일한 확률을 가진다고 생각할 수 있습니다.

여기선 [0.2, 0.2, 0.2, 0.2, 0.2]으로 표현할 수 있어요. (어떤 단어를 선택해야 하는지 알 수 없는 상태)

 

여기서 노이즈를 한 번 더해볼까요?

논문에서 정의하는 식에는 시간 개념도 함께 활용됩니다.

노이즈를 점진적으로 추가하는 모델이니, 시간이 지남에 따라 더 많은 노이즈가 낀다고 생각하면 됩니다.

 

t = 0.3

p = 0.7 * [0, 1, 0, 0, 0] + 0.3 * [0.2, 0.2, 0.2, 0.2, 0.2] = [0.2, 0.9, 0.2, 0.2, 0.2]

 

결국, 시간이 지남에 따라 원본 데이터의 정보 비율이 점차 줄어드는 것을 의미해요.

 

이렇게 열심히 설명하긴 했지만, 이건 기존의 다른 연구들에서 제안된 방법에 가깝고 결국 본 논문은 이 방법 대신 마스킹 토큰으로 노이즈를 대체하는 방식을 취하고 있습니다.

 

3.2 Masked Diffusion

정뱡향 과정에서는 깨끗한 데이터에서 점진적으로 마스킹 상태로 전이됩니다. 한 번 마스킹된 토큰은 다시 복구되지 않습니다. (논문에서는 이를 Absorbing state = 흡수 상태로 표현합니다.) 반대로 역방향 과정에서는 노이즈를 제거하며 복원합니다. 이때도 역시 마스킹 상태가 아닌 토큰은 그대로 유지하면서 복원합니다. (이 과정에서 SUBS 매개변수화가 사용되는데, 마스킹 상태와 비마스킹 상태를 명확히 구분하여 복원하기 위해 사용되는 기술로 이해했습니다.)

 

 

 

3.3 Rao-Blackwellized Likelihood Bounds

Rao-Blackwellization을 통해 본래 복잡했던 손실 함수를 단순화하여 계산 효율성을 높입니다. 이게 가능했던 이유는 앞선 설명한 SUBS 매개변수화 과정에서 특정 항이 자연스럽게 0으로 처리됐기 때문입니다. 손실 함수가 단순하면 모델 학습 과정에서 분산이 줄어들고 계산이 쉬워져 학습 시 안정적인 수렴이 가능해집니다.

 

3.5 Masked Diffusion Language Models

다른 부분은 이전에 설명했던 부분과 같은 맥락이고 목적 함수를 이해하는게 중요합니다.

하지만, 오히려 다른 부분보다 더 쉽게 이해할 수도 있습니다. 왜냐면 목적함수가 BERT의 Masked Language Modeling (MLM)과 매우 유사한 형태를 띄고 있기 때문입니다. 기존 BERT 모델은 고정된 마스킹 비율을 사용하지만, MDLM은 확률적으로 다양한 마스킹 비율 학습이 가능합니다. 또한, 기존 BERT는 인코더 모델이기 때문에 생성 능력이 부족하지만, MDLM은 생성 확산 모델로 자연스러운 문장 생성이 가능합니다. 학습시 몇 가지 고려사항이 있는데, 첫째 작은 어휘 집합 (8,000개)은 긴 의존성을 유발해서 성능 저하를 초래할 수 있습니다. 따라서, 적절한 토큰화 기법과 충분히 큰 어휘 집합이 필요합니다. 또, 유연한 학습을 위해서 모델 구조를 일부 개선했습니다. Diffusion Transformer (DiT)를 사용하고 Rotary Positional Embedding을 사용합니다.

 

... 그냥 우리 이거 이거 이거 사용했따! 이러고 넘어가는데 Diffusion Transformer, Rotary Positional Embedding 이런 기술들은 사실 하나 하나 다 중요한 내용이거든요. DiT는 살짝 맛보기만 하고 넘어가 볼게요.

 

4 Inference and Sampling in Masked Diffusion Language Models

본 논문에서는 Semi-Autoregressive (SAR)를 활용하는데, 이는 자기회귀 모델과 비자기회귀 모델의 장점을 결합하여 효율성과 유연성을 모두 확보한 방식입니다. SAR은 일부 토큰을 미리 고정하고 복원 과정에서 고정된 부분은 그대로 유지합니다. 따라서 복원이 지속될수록 새로운 토큰만 추가 복원하면 되기 때문에 빠른 복원이 가능합니다. (여기서 복원은 생성과 같은 맥락입니다.) 따라서, 추가 문장을 생성 시 유연한 확장이 가능합니다.

저는 뒷부분이 조금 헷갈렸어요. 복원된 문장을 prefix로 사용해서 추가로 문장을 복원(생성)한다고 하면, 결국 앞을 보고 뒤를 생성하는 auto-regressive와 동일한게 아닌가? 하는 생각이 들었거든요. 생성 단위가 블록이라고 이해하면 편합니다. 하나의 블록을 생성할 때는 diffusion 모델링을 활용하지만, 전체적으로 블록간 처리 방식을 보면 auto-regressive 성격을 띄고 있는거죠.

 

쉬운 이해를 위해서 예시를 들어보겠습니다.

편지를 쓴다고 가정하고, 처음에 안녕? 이라는 문장을 생성합니다.

그 다음 추가로 글을 작성할 때, 00아 안녕? 나는 00이야. 이렇게 생성하는게 아니라 (안녕?) 나는 00이야! (뒷 부분만 추가 생성합니다.)

대신 나는 00이야!를 생성하는 방식은 "나는 → 00이야" (이렇게 순차적으로 생성하는게 아니라) "나는 00이야!" (한번에 생성)을 택하는 거죠.

 

그럼 블록간 연결성도 강화되고 또 한 블록을 생성할 때는 순차적으로 토큰을 생성하지 않고 동시 다발적으로 생성하게 되니깐 효율성도 잡을 수 있겠죠?

흐음 그럼 또 multi-token prediction과는 어떤 차이가 있는건지 문득 궁금해지네요.

지금 이 궁금증을 해결하려면 또 한세월일거라 미래의 저에게 넘기겠습니다.

5 Experiments

실험에서는 DNA 언어 모델링도 언급했는데, 이 부분은 생략하고 리뷰할게요.

 

Likelihood Evaluation.

 

MDLM은 이전 확산 모델 (SEDD)보다 우수하며, AR 모델과의 성능 차이도 크게 줄였습니다.

(하지만 Autoregressive 모델을 보면, 여전히 perplexity 면에서 떨어지는 수치를 보입니다.)

 

Zero-Shot Likelihood Evaluation.

 

흥미로운 결과가 나왔어요. 일반화할 수 있는지 확신은 없지만 Lambada, Punmed, Arxiv 데이터셋에서 MDLM 모델이 AR 모델에 비해 더 우수한 성능을 보였습니다. 이 결과를 두고 AR 모델보다 MDLM 모델이 일반화 성능이 더 좋다고 논문에서는 밝히고 있습니다. AR 모델은 도메인 차이가 클수록 (학습 데이터와 실제 사용 측면에서) 성능이 급격히 떨어지는 특성이 있는데, MDLM은 이 부분에서 좀 더 강건한 모습을 보인다는 의미입니다. 그 이유는 MDLM은 학습 시 마스킹 토큰을 복원하는 방식을 사용하기 때문에 노이즈에 굉장히 강건한 모델이 되는거죠. (BERT 모델이 GPT 계열 모델보다 NLU 성능이 더 좋을거라 개인적으로 믿고 있기 때문에 이 부분도 어느 정도 납득이 되는 내용이었습니다.)

 

Downstream Task Evaluation.

 

이 부분은 지금도 조금 헷갈리는 부분이긴 합니다.

BERT를 MDLM 방식으로 Fine-tuning 하여 향상된 결과를 보여줬는데 그렇다면 왜 MDLM 모델까지 함께 비교해주지 않은건지 모르겠습니다.

 

Semi-Autoregressive Modeling.

MDLM의 SAR 성능 평가를 위해 추가로 SSD-LM 모델과 성능 비교를 했습니다. SSD-LM은 블록 단위로 문장을 생성하는 자기회귀 확산 모델입니다. 논문에서는 단순히 SSD-LM보다 퍼플렉시티 측면에서도 우수하며 디코딩 속도 역시 약 25~30배 빠르다고 언급하고 있습니다.

 

하지만, 제가 궁금한건 단순 결과 보다 모델 자체에 어떤 차이가 있냐는 겁니다. 블록 단위로 문장을 생성하는 자기회귀 확산 모델? 이게 뭔가요?

분명 아까 MDLM도 블록 단위로 prefix로 설정해서 복원 시 고정하면서 비슷한 뉘앙스를 풍겼거든요.

굉장히 헷갈리는 개념이었지만, 결국 가장 큰 차이는 한 블록을 처리할 수 있는 토큰 개수의 차이입니다. SSD-LM은 25토큰, MDLM은 512토큰이거든요.

SSD-LM은 짧은 블록으로 인해 AR - DLM간 전환이 잦아 속도가 느릴 수 밖에 없습니다. 또 MDLM의 경우, 복원 시 prefix로 사용하기 때문에 AR 느낌은 나지만 실상 전체적으로 diffusion 모델링 기법을 활용하는거라 병렬 처리면에서 훨씬 효율적일 수 밖에 없다는 생각이 들었습니다. (선 복원된 문장은 단순히 복원 생략일뿐)

 

6 Related Work

이 부분은 사실 진한 감동이었어요. 이렇게 많은 모델들과 비교하다니, 괜히 NeurIPS가 아니네요.

(표는 개인적으로 제가 작성했습니다. 줄 글로 읽으니 조금 힘들더라고요.)

	MDLM (Ours)	D3PM	CTMC	SEDD	BERT	Latent DLM
모델 유형	마스킹 확산 모델	이산 확산 모델	연속 확산 모델	점수 기반 확산 모델	인코더	연속 공간 확산 모델
노이즈 처리 방식	마스킹	이산 노이즈 처리	무한 행렬 곱 - 연속 노이즈 처리	점수 계산 기반 노이즈 처리	마스킹	가우시안 확산
학습 속도	중간	느림	매우 느림	느림	빠름	느림
추론 속도	매우 빠름 (병렬적 복원)	느림	매우 느림	느림	빠름	느림
생성 품질 (PPL)	우수 (PPL 낮음)	우수	이론적으론 우수, 실제 구현 어려움	복잡성 높음	낮음	성능이 우수할 수록 복잡성 높음
학습 복잡성	낮음	복잡	매우 복잡	매우 복잡	간단	복잡
학습 안전성	높음	불안정	불안정	불안정	안정	불안정

---

 

연속으로 NeurIPS 논문들을 읽고 있는데, 쉽지 않네요.

Diffusion 모델이 텍스트 분야에서 이렇게 다양한 연구가 이뤄졌는지 몰랐습니다.

자고 일어나면 논문이 뚝딱 생긴다는 말은 정말 온 분야를 통틀어서 일어나는 일이었나 봅니다.

매번 잠은 나만 자지...

 

 

하지만 졸립죠?