향상된 대형 언어 확산 모델

대부분의 최신 LLM은 autoregressive 학습과 causal attention에 의존하는 반면, 본 논문은 완전 양방향 masked diffusion으로 처음부터 대규모 학습한 8B 모델(iLLaDA)이 경쟁력 있는 성능을 달성함을 보여준다. 특히 수학과 추론·코드 벤치마크에서 autoregressive 강력 모델과 근접하거나 우수한 결과를 얻었다.

대형 언어 모델의 전통적 훈련은 토큰을 왼쪽에서 오른쪽으로 순차적으로 예측하는 autoregressive factorization과 causal attention에 기반한다. 이 접근은 자연스럽게 생성 품질과 샘플링 단계를 단순화하지만, 토큰 간의 양방향 정보 활용과 마스킹 기반 역추론(reverse reasoning) 능력은 제한적이다. 따라서 양방향 문맥을 학습하는 대안적 학습 목표가 탐색될 필요가 있다.

Masked diffusion은 입력 문장 x₀(길이 L)와 마스킹 비율 t∈U[0,1]을 입력으로 받는다 → 각 토큰을 확률 t로 M(마스크)로 대체해 손상된 시퀀스 x_t를 만든다 → 모델은 마스크된 위치들에 대해 p_θ(x₀^i | x_t)를 예측하고, 손실은 마스크된 토큰에 대해 1/t·Σ 1[x_t^i = M] · (−log p_θ(·))로 계산된다 → 이 값은 마스크된 토큰당 평균 음의 로그우도를 반환하므로 다양한 마스킹 강도에서 일관된 복원 성능을 학습하게 된다. 이 방식은 좌우 문맥을 모두 보면서 마스크된 위치의 원토큰을 복원하도록 학습하므로 bidirectional attention과 자연히 결합된다.

본 논문은 위 원리를 대규모(12T 토큰)로 확장하고 SFT에서도 동일한 masked-diffusion 포맷을 사용한다. 이로 인해 모델은 in-context learning과 instruction-following 능력을 획득하면서도 비자동회귀 학습의 장점을 유지한다. 추가로 GQA와 임베딩-헤드 묶음(tied embeddings)을 통해 메모리와 파라미터를 효율화하고, variable-length generation과 confidence-based scoring으로 실제 추론·평가 효율을 개선했다. 결과적으로 iLLaDA는 여러 reasoning·math·code 벤치마크에서 기존 LLaDA 대비 큰 향상을 보이며, 일부 벤치마크에서는 autoregressive 계열과 경쟁 가능한 성능을 달성했다.

전체 접근은 masked diffusion 목표를 사전학습(pre-training)과 SFT 전단계에 동일하게 유지하면서, 스케일·아키텍처·데이터 처리·추론 규칙을 실무적 관점에서 수정한 것이다. 사전학습 데이터는 합계 12T 토큰으로 확장했고, 모델은 dense Transformer(32 layers, model dim 4096) 구조를 기본으로 한다. 아키텍처적 최적화를 위해 RMSNorm, SwiGLU, RoPE를 사용하고 attention/MLP bias를 제거했다.

핵심 메커니즘은 손상 시퀀스 x_t 생성과 마스크된 토큰 예측 손실이다. 입력으로 깨끗한 시퀀스 x₀(길이 L)와 t~U[0,1]가 주어지면 → 각 토큰을 확률 t로 M으로 대체해 x_t를 만들고 → 모델은 모든 마스크된 위치에 대해 p_θ(x₀^i | x_t)를 출력하며 → 손실 ℒ(θ)=−E_{t,x₀,x_t}[ (1/t) Σ_i 1[x_t^i=M] log p_θ(x₀^i|x_t) ]를 최소화한다. 1/t 인자는 마스크 비율에 따른 손실 스케일 정규화 역할을 한다.

아키텍처·구현 세부는 GQA 도입, 임베딩-헤드 파라미터 결합(tied embeddings), 최대 시퀀스 길이 8192, 30% 확률로 시퀀스를 두 개로 랜덤 분할(random-length), FlashAttention 기반 variable-length attention 커널(누적 오프셋으로 패딩 없이 예제 분리) 사용으로 구성된다. 옵티마이저는 AdamW(weight decay=0.1)이며 사전학습 시 학습률을 선형 warmup으로 2×10^{-4}까지 올린 뒤 고정하고, 손실이 멈추면 5×10^{-6}까지 cosine decay로 전환했다.

SFT는 전체 instruction 코퍼스를 8192-token 연속 스트림으로 합치고 무작위 마스킹을 적용해 pre-training 포맷을 그대로 유지한다. SFT 코퍼스는 약 25B 토큰이고 12 epochs로 fine-tune을 수행했다. SFT 학습률은 warmup→고정→마지막 10% 동안 5×10^{-7}까지 선형 감쇠로 설정했다.

추론과 평가에서는 두 가지 핵심 규칙을 사용한다. 다중선택 평가용 confidence-based scoring은 후보 연속문 y(길이 L)를 모두 마스크한 상태에서 반복적으로 가장 모델이 자신 있어 하는 위치 i_k를 선택해 그 토큰을 공개하고 로그확률을 누적한다: i_k = argmax_{i∈M_{k-1}} p_θ(y^i | p,  ilde{y}{k-1}) → S_conf(y|p)=Σ_k log p_θ(y^{i_k}|p, ilde{y}{k-1}). 생성은 block 단위의 마스크를 디퓨전 샘플러로 디코딩하고, 신뢰도가 높은 토큰을 고정한 뒤 필요 시 다음 블록을 붙이는 variable-length generation과 low-confidence remasking 전략을 따른다.

메인 벤치마크(베이스 모델)에서 iLLaDA는 LLaDA 대비 전반적 성능 향상을 보였다. Table 2 기준 iLLaDA 8B는 MMLU 74.8(대조: LLaDA 65.9), BBH 71.3(49.7), ARC-C 60.8(45.9), GSM8K 81.9(70.3), HumanEval 50.0(35.4)로 보고되며 전체 평균 63.9(대조: 51.1)을 기록했다. 이 표에서 iLLaDA는 일부 수학·추론 벤치에서 Qwen2.5 7B와 경쟁하거나 상회하는 항목들을 포함한다.

Instruction-tuned 결과에서도 iLLaDA가 LLaDA 대비 큰 개선을 보였다. Table 3 기준 iLLaDA-Instruct는 MMLU 71.6(대조 65.5), MMLU-Pro 52.3(37.0), GSM8K 89.0(77.5), Math 56.7(42.2), HumanEval 65.9(49.4), 평균 67.1(대조 54.5)을 기록했다. 다만 Qwen2.5 7B Instruct와 비교하면 일부 수학·코드 벤치에서는 차이가 남아있고, 논문은 그 원인으로 RL 기반 정렬(추가 적합)의 부재를 지목했다.

Ablation 결과에서는 confidence-based scoring이 우도 기반 다중선택 규칙 대비 일관된 개선을 보였다(PIQA +1.3, ARC-C +0.6, HellaSwag +2.3). SFT epoch ablation에서는 SFT epoch 수가 증가할수록 GSM8K, MATH, MMLU-Pro에서 성능이 전반적으로 개선되어 긴 SFT가 이 모델에서 유의미한 이득을 제공했다. 계산 자원 제약으로 12 epoch를 초과해 실험하지 않았다.

아키텍처는 32 layers, 모델 차원 4096, attention heads 32, key/value heads 8, FFN 차원 14,336, vocabulary 155,136, 최대 시퀀스 길이 8192로 구성되며 총 파라미터는 7.62B이다(비임베딩 파라미터 6.98B). 입력 임베딩과 LM-head를 tied하여 전체 파라미터를 줄였고, attention/MLP bias를 사용하지 않았다.

핵심 메커니즘은 masked diffusion objective ℒ(θ)=−E_{t,x₀,x_t}[ (1/t) Σ_i 1[x_t^i=M] log p_θ(x₀^i|x_t) ]로, t는 U[0,1]에서 샘플링되고 각 토큰은 확률 t로 마스크로 대체된다. 손실의 1/t 스케일링은 마스크된 토큰당 평균 음의 로그우도를 반환하여 다양한 마스킹 강도에서 안정적 학습을 유도한다.

구현 세부로는 grouped-query attention(GQA) 도입으로 KV 캐시 유사 구현에서 key/value 상태의 메모리 풋프린트를 절감했고, FlashAttention 기반 variable-length attention 커널을 통해 배치 내 서로 다른 길이 예제를 패딩 없이 처리했다(누적 시퀀스 오프셋 사용). 시퀀스들은 30% 확률로 랜덤하게 두 개의 짧은 세그먼트로 분할되어 학습 효율을 높였다.

학습 하이퍼파라미터는 다음과 같다. 사전학습: 총 12T 토큰, 학습률은 선형 warmup→2×10^{-4}까지 상승 후 고정, 손실 정체 시 cosine decay로 전환해 최소 5×10^{-6}까지 감쇠. 옵티마이저는 AdamW(weight decay=0.1). SFT: 약 25B 토큰, 12 epochs, 학습률 warmup→고정→마지막 10% 동안 선형 감쇠로 5×10^{-7}까지 감소, 동일한 AdamW 사용.

추론 알고리즘은 다중선택에 대해 confidence-based scoring을 사용하며, 생성은 block 단위 variable-length generation과 low-confidence remasking을 결합해 연속 블록을 디코딩한다. confidence-based scoring은 후보 y의 모든 토큰을 마스크한 뒤 반복적으로 모델이 가장 자신 있어 하는 위치의 토큰을 공개하고 해당 로그확률을 합산해 점수를 산출한다.

논문에서 명시한 한계는 다음 두 가지이다. 첫째, iLLaDA는 reinforcement-learning 기반 정렬(RL-based alignment)을 추가로 적용하지 않아 instruct 설정에서 일부 autoregressive instruct 모델과 격차가 존재한다. 둘째, 계산 자원 제약으로 연구를 8B 스케일에 한정했으며, autoregressive 모델과 완전히 매칭되는 대규모 비교는 수행하지 않았다.