iLLaDA

현재 대규모 언어 모델은 자기회귀 패러다임이 지배한다. 최근 디퓨전 언어 모델은 언어 생성에 대한 다른 접근으로 주목받고 있으며, LLaDA는 마스크 디퓨전 정식화를 따라 완전 양방향 어텐션으로 언어 모델을 처음부터 학습했다.

LLaDA는 비자기회귀 모델도 문맥 내 학습, 지시 따르기 같은 핵심 LLM 능력을 획득할 수 있음을 보여, 언어 지능이 반드시 자기회귀 모델링에 의존해야 한다는 통념에 도전했다. 이 개념적 함의를 넘어, 양방향 디퓨전 언어 모델은 역방향·양방향 추론, 장기 지평 계획, 멀티모달·옴니 모델링에서 이점을 보였다. 최근 연구는 양방향 디퓨전 사전학습이 반복 학습하에서 제한된 데이터를 더 잘 활용해, 데이터 제약 환경에서 자기회귀 모델을 능가하게 함을 보였다.

iLLaDA의 주요 변경

이 변경들은 LLaDA를 실질적으로 개선한다. 처음부터 학습한 LLaDA와 Qwen2.5에서 미세조정한 Dream을 포함한 기존 양방향 디퓨전 언어 모델 대비, iLLaDA는 베이스·지시 평가 모두에서 최고 평균 성능을 얻는다. Qwen2.5 7B와 비교하면 iLLaDA-Base는 평균에서 약간 더 강하고, iLLaDA-Instruct는 아직 뒤처진다. 절제 연구는 신뢰도 기반 채점이 객관식 평가를 개선하고, iLLaDA가 다중 에폭 SFT에서 계속 이득을 얻음을 보인다.

LLaDA의 마스크 디퓨전 정식화를 유지하되, 확장·사후학습·평가에 중요한 몇 가지 실용적 변경을 가한다.

2.1 사전학습 masked diffusion objective

iLLaDA는 LLaDA와 동일한 사전학습 목적을 따른다. 길이 L의 깨끗한 시퀀스 x₀가 주어지면 마스킹 비율 t~U[0,1]를 샘플링하고, 각 토큰을 확률 t로 독립적으로 마스크 토큰 M으로 대체해 손상 시퀀스 xₜ를 얻는다. 모델은 마스크된 모든 토큰을 예측하도록 학습된다.

백본은 RMSNorm, SwiGLU, RoPE를 쓰고 어텐션·MLP 바이어스가 없는 밀집 트랜스포머다. 멀티헤드 어텐션을 쓰는 LLaDA와 달리 iLLaDA는 GQA를 사용한다. KV 캐시류 메커니즘이 디퓨전 언어 모델에도 적용될 수 있음이 최근 밝혀졌고, 이런 캐시형 구현하에서 GQA는 캐시된 키/값 상태의 메모리 사용량을 줄인다. 파라미터 수를 더 통제하기 위해 입력 임베딩과 LM 헤드 파라미터를 묶는다.

최대 시퀀스 길이 8192로 사전학습하며, 마스크 디퓨전의 무작위 길이 학습에서 영감을 받아 30% 확률로 8192 시퀀스를 두 짧은 세그먼트로 무작위 분할한다. 가변 길이 예시를 배치마다 패킹하고 FlashAttention 기반 가변 길이 어텐션 커널로 패딩 없이 예시를 분리한다. 학습률은 2×10⁻⁴까지 선형 워밍업 후 일정하게 유지하다가, 사전학습 손실이 더 줄지 않을 때 최소 5×10⁻⁶의 코사인 감쇠로 전환하니 손실이 다시 개선됐다. 가중치 감쇠 0.1의 AdamW를 쓴다.

2.2 지도 미세조정 SFT

선행 연구는 보통 프롬프트와 전체 참조 응답을 이어붙여 SFT 인스턴스를 구성하고, 학습 중 프롬프트 토큰은 가시 상태로 두고 응답 영역에만 마스크를 적용했다. 미니배치 내 짧은 응답은 가장 긴 응답 길이에 맞춰 |EOS|로 패딩됐다.

iLLaDA는 대신 사전학습과 동일한 데이터 처리·마스킹 방식을 쓴다. 각 지시 예시를 프롬프트-응답 시퀀스와 단일 종단 |EOS|로 포맷하고, 모든 예시를 연속 지시 코퍼스로 이어붙인 뒤 8192 토큰 학습 시퀀스를 샘플링한다. 그다음 전체 시퀀스에 무작위 마스크를 적용해 식 (1)을 최적화하므로, 프롬프트·응답·|EOS| 토큰이 모두 마스크될 수 있다. 사전학습과 같은 무작위 길이 학습도 쓴다. 이 SFT 포맷은 가변 길이 블록 생성을 자연스럽게 지원한다.

SFT 코퍼스는 약 250억 토큰이며 12 에폭 미세조정한다. 학습률은 5×10⁻⁶까지 선형 워밍업 후 일정하게 유지하다가 마지막 10% 구간에서 5×10⁻⁷까지 선형 감쇠한다.

2.3 추론 confidence scoring · variable-length

iLLaDA는 LLaDA와 동일한 확률적 정식화를 쓴다. 두 모델 모두 식 (1)의 마스크 디퓨전 목적으로 학습되며, 이는 모델 분포의 음의 로그우도 상계에 해당한다.

객관식 — 신뢰도 기반 채점

많은 벤치마크가 HellaSwag, PIQA, ARC-Challenge처럼 객관식으로 정식화된다. 접두 p와 유한한 후보 연속들이 주어지면, 평가는 각 후보에 점수를 매겨 최고점을 고른다. iLLaDA는 결정론적 신뢰도 기반 채점 규칙을 쓰며, 이는 객관식에서 로그우도 상계보다 경험적으로 더 낫다. 전부 마스크된 후보에서 시작해, 남은 마스크 위치 중 모델이 가장 높은 신뢰도를 부여하는 토큰을 매 단계 하나씩 정답 토큰으로 드러낸다.

개방형 — 가변 길이 생성

프롬프트에 마스크 토큰 블록을 덧붙이고 그 블록 안에서 디퓨전 샘플러를 돌린다. 각 샘플링 단계에서 모델은 모든 마스크 위치를 예측하고, MaskGIT·LLaDA의 저신뢰 리마스킹 전략을 따라 가장 확신하는 예측을 가시 토큰으로 전환하되 저신뢰 위치는 마스크로 유지한다. 블록이 디코딩되면 |EOS|나 정지 토큰이 나타날 때 종료하고, 아니면 새 마스크 블록을 덧붙여 최대 생성 예산까지 계속한다.

베이스·지시 평가 모두에서 iLLaDA는 기존 디퓨전 언어 모델을 실질적으로 개선하며, 여러 추론 벤치마크에서 강력한 자기회귀 베이스라인과 경쟁력을 유지한다.

핵심 향상 — Base, vs LLaDA

3.2 절제 연구 ablation

처음부터 학습된 8B 완전 양방향 디퓨전 언어 모델 iLLaDA를 제시했다. 사전학습을 12T 토큰으로 확장하고 모델 설계, 학습률 스케줄, SFT 포맷, 신뢰도 기반 객관식 채점, 가변 길이 생성 등 실용 레시피의 여러 부분을 갱신했다. iLLaDA가 다중 에폭 SFT에서 계속 이득을 얻음도 확인했다. 베이스·지시 평가 전반에서 이 변경들은 일반·수학·코드 벤치마크에서 LLaDA 대비 실질적 개선으로 이어져, 처음부터의 완전 양방향 디퓨전 학습이 강한 언어 모델링 성능을 달성할 수 있음을 시사한다.

한계

• 강화학습 정렬 부재. iLLaDA는 RL로 추가 정렬되지 않았으며, 이것이 iLLaDA-Instruct와 강한 AR 지시 모델 간 남은 격차를 부분적으로 설명할 수 있다. 마스크 디퓨전 LLM용으로 개발된 최신 RL 기법(VRPO, diffu-GRPO, MDPO, ESPO)을 iLLaDA에 직접 적용해 지시 따르기·추론 능력을 더 높일 수 있다.
• 규모 제약. 연산 한계로 연구가 8B 규모에 국한되며 자기회귀 모델과 완전히 맞춘 비교를 제공하지 못한다. 대신 12T 토큰 사전학습에 연산을 배분했다. RL 정렬과 더 큰 규모 연구를 향후 과제로 남긴다.

iLLaDA-8B-Base

BBH·GSM8K·MATH·HumanEval·MBPP에 개방형 생성을 쓴다. BBH·GSM8K·MATH·MBPP는 최대 생성 길이 1024, 블록 길이 32로 설정한다. HumanEval은 준자기회귀 블록 샘플링이 성능을 해치는 것을 관찰해 최대 생성 길이와 블록 길이를 모두 512로 둔다.

iLLaDA-8B-Instruct

벤치마크별 추론 설정을 쓴다. 답 글자 하나만 생성하면 되는 MMLU·MMLU-Redux는 최대 생성/블록 길이를 각각 4/4, 3/3으로 둔다. GSM8K·HumanEval은 2048/32, MMLU-Pro·MATH는 4096/32, MBPP는 2048/16으로 설정한다.