향상된 대형 언어 확산 모델

왜 중요한가

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

핵심 기여

완전 양방향 masked diffusion으로 대규모 학습한 8B 모델(iLLaDA)

iLLaDA는 8B 파라미터 규모의 dense Transformer로서 masked diffusion 목표를 사전학습과 SFT 전 과정에 유지하며, 사전학습을 12T 토큰까지 확장해 비자동회귀(diffusion) 방식으로 대형 모델을 학습했다.

사전학습·SFT 실무 레시피 개선

사전학습에서 최대 시퀀스 길이 8192, 30% 확률로 시퀀스 분할(random-length), 학습률 warmup→고정→cosine decay 전환 등 학습률 스케줄 조정, SFT는 25B 토큰으로 12 epochs 수행해 성능 향상을 확보했다.

구성적 아키텍처 변경으로 메모리·파라미터 절감

Grouped-Query Attention(GQA)를 도입해 KV 캐시 유사 구현의 메모리 부담을 줄였고, 입력 임베딩과 LM-head를 tied하여 전체 파라미터 수를 감소시켰다(총 파라미터 7.62B, LLaDA 8.02B 대비).

실전용 추론·평가 기법: variable-length generation 및 confidence-based scoring

생성은 블록 단위 마스크와 low-confidence remasking 전략으로 가변 길이생성을 수행해 효율을 높였고, 다중선택 평가에서는 confidence-based scoring을 사용해 우도 기반 대비 일관된 성능 향상을 확인했다.

다양한 벤치마크에서 LLaDA 대비 큰 성능 향상

iLLaDA-Base는 BBH에서 +21.6p, ARC-Challenge에서 +14.9p 등 LLaDA 대비 다수 벤치마크에서 대규모 향상을 보여주었고, instruct 버전에서는 MATH +14.5p, HumanEval +16.5p 등 유의미한 개선이 관찰되었다.

핵심 아이디어 이해하기

대형 언어 모델의 전통적 훈련은 토큰을 왼쪽에서 오른쪽으로 순차적으로 예측하는 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 모델과 완전히 매칭되는 대규모 비교는 수행하지 않았다.

실무 활용

iLLaDA는 양방향 masked diffusion 학습과 variable-length generation을 결합해 추론 효율과 추론 품질 사이에서 균형을 제공한다. 공개된 모델·코드(레포 링크)에 따라 연구·프로토타입 단계에서 수학·추론·코드 관련 태스크에 적용 가능성이 높다.

추론·추론형 QA 시스템에서 bidirectional reasoning이 필요한 서술형 문제 해결
수학 문제 풀이·단계적 추론을 요구하는 학습 보조 시스템
코드 생성·보충(코드 스니펫 완성)에서 비자동회귀 생성으로 추론 블록 단위 효율화
다중선택 평가가 많은 벤치마크 또는 시험 시스템에서 confidence-based scoring 적용

코드 공개 여부: 공개

코드 저장소 보기

키워드

masked diffusion language model(마스크드 디퓨전 언어 모델)bidirectional attention(양방향 어텐션)grouped-query attention(그룹드-쿼리 어텐션)supervised fine-tuning (SFT)(지도 파인튜닝)variable-length generation(가변 길이 생성)confidence-based scoring(신뢰도 기반 점수화)MMLU(다중과제 언어 이해 벤치마크)

용어 해설

마스크드 디퓨전(Masked Diffusion): — 문장 내 일부 토큰을 무작위로 M(마스크) 토큰으로 대체한 후 모델이 마스크된 위치의 원래 토큰을 복원하도록 학습하는 확률적 생성 목표이다. 본 논문에서는 마스킹 비율 t를 U[0,1]에서 샘플링하고 손실을 마스크된 토큰 수로 정규화(1/t)하여 모델이 다양한 마스킹 강도에서 정확한 복원을 학습하도록 만든다.
양방향 어텐션(Bidirectional Attention): — 토큰이 좌우 문맥을 모두 볼 수 있게 하는 Transformer 어텐션 방식이다. 본 논문에서는 완전 양방향 attention을 사용해 비자동회귀(masked diffusion) 학습을 수행하며, 이를 통해 역추론(reversal)과 양방향 추론 능력 확보가 가능하다.
그룹드-쿼리 어텐션(Grouped-Query Attention (GQA)): — 쿼리 헤드를 그룹화하고 key/value 헤드를 줄여 KV 캐시 형태의 메모리 사용량을 절감하는 어텐션 변형이다. 캐시 기반 구현에서 key/value 상태의 메모리 풋프린트를 줄여 긴 컨텍스트에서 효율적이다.
가변 길이 생성(Variable-Length Generation): — 생성 시 고정 길이 대신 블록 단위로 마스크를 붙이고 디퓨전 샘플러를 돌려 자신감 높은 토큰을 고정한 뒤 필요한 경우 다음 블록을 이어 붙여 생성하는 방식이다. 불필요한 패딩을 줄여 추론 효율을 높인다.
신뢰도 기반 점수화(Confidence-based Scoring): — 다중 선택 평가에서 후보 정답의 각 토큰을 반복적으로 가장 모델이 자신 있어 하는 위치부터 순차적으로 공개하고 해당 토큰들의 로그확률 합을 스코어로 사용하는 규칙이다. 우도 기반 점수보다 다중선택에서 성능이 높은 것으로 보고됐다.

향상된 대형 언어 확산 모델

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드