Sumi: Scratch에서 시작한 오픈 Uniform Diffusion 언어 모델

왜 중요한가

Uniform diffusion은 모든 토큰을 어느 시점에든 업데이트할 수 있어 생성 유연성을 높인다. Sumi는 1.5T 토큰으로 scratch에서 pretrained된 7B UDLM의 첫 공개 사례로, UDLM의 확장성, 생성 역학, 제어성 및 기존 AR/MDLM과의 트레이드오프를 연구하는 기준점을 제공한다. 모델 가중치와 학습 레시피의 공개는 재현성과 비교 연구를 촉진한다.

핵심 기여

7B UDLM을 Scratch에서 사전학습

Sumi-7B는 7B 매개변수의 Uniform Diffusion Language Model을 scratch에서 pretrained하며 약 1.5T 토큰으로 학습된다. 공개 가중치와 체크포인트, 전체 학습 레시피를 제공한다.

데이터 혼합 및 학습 레시피의 공개

사전학습 데이터는 llm-jp-corpus-v4의 영어 부분과 codex/megamin 코드 데이터를 포함하며, 1.3T 토큰의 사전학습과 250B 토큰의 중간 학습으로 구성된다. 중간 학습은 도메인 특화 데이터 혼합으로 구성되고, 학습은 WSD 스케줄과 함께 진행된다.

지식/추론/코딩 벤치마크에서 AR 대비 경쟁력

일반 지식 및 코딩 벤치마크에서 Sumi가 프로토콜 내 평가모델 중 최고 점수에 가까운 성능을 보였으며, Reasoning 및 수학 벤치에서도 AR 대비 경쟁력을 보인다. 단 commonsense 벤치에서는 격차가 크다.

생성 시퀀스의 canvass 근거 분석 및 인퍼런스 프로브

canvas 길이에 따른 생성 fluency, 확신 샘플링에 의한 commit order 형성, 1-토큰 단위의 병렬 디코딩 가능성, 수정 예산(revision budget)에 의한 자기 수정의 한계 등을 탐색한다.

무대화형 UDLM에 대한 재현 가능한 연구 설계

모델 가중치와 중간 체크포인트, 학습 레시피를 공개하여 UDLM의 스케일링 연구를 재현 가능하게 한다.

핵심 아이디어 이해하기

출발점: 전통적 AR/MDLM은 특정 시퀀스의 일부 토큰을 고정적으로 채우는 제약이 있다. UDLM은 토큰 업데이트를 시간 축 전반으로 허용해 self-correction의 가능성을 열지만, 학습과 평가에서의 거대한 데이터/모델 규모를 요구한다. 해결 원리: Sumi는 GIDD 프레임워크와 SNR 재매개변수화를 기반으로 시간-비가역적 Bidirectional Transformer를 구성하고, uniform noise 아래에서 토큰 업데이트를 수행한다. 결과적으로 UDLM은 일반 지식, 추론, 코딩 벤치에서 AR 기반 모델과 비교해 competitive한 성능을 보여주나, 교육 중심의 데이터 혼합의 효과로 commonsense 벤치에서 성능 저하가 나타난다. 달라지는 점: canvas 길이 Band 내에서의 생성 유창성, confidence sampling이 도메인별 구조를 형성하는 효과, revision budget에 의한 self-correction의 부재가 관찰되었다.

방법론

아키텍처: 36-layer, 4096(hidden), SwiGLU, 12288 FFN, 32 heads, 8 KV groups, 128 head dimension, RMSNorm, untied input/output embeddings, bias 없음, dropout 없음, attention에서 off-by-one softmax를 사용한다. RoPE(θ=500,000) 적용. 토크나이저는 OLMo 3를 사용하며 어휘 수는 100,278이다. Megatron-LM 기반 프레임워크로 학습한다. 하드웨어: 288개의 NVIDIA H100 GPU에서 43,308 GPU-hours 소요. 데이터: pre-training 1.3T token, sequence length 1184, 배치 4,608 시퀀스. 학습은 AdamW(β=(0.9,0.95), weight_decay=0.1, gradient_clip=1.0)와 z-loss(계수 1e-5) 사용. 중간 학습은 도메인 혼합 데이터로 130B 토큰(길이 1184) 및 120B 토큰(길이 4864)으로 진행되며, 학습 스케줄은 2,000스텝 warmup에서 peak 2e-4, 상수 구간, 2,000스텝 cooldown에서 2e-5로 조정한다. 데이터 구성: pre-training은 llm-jp-corpus-v4의 영어 부분 및 python 코드 포함; mid-training은 llm-jp-corpus-midtraining-v2 및 MegaMath-Web-Pro-Max 계열 데이터를 포함하도록 조정했다.

주요 결과

벤치마크 결과: Sumi-7B는 일반지식 및 코딩 벤치에서 프로토콜 내에서 최고점에 근접하거나 최고치를 달성한다. MMLU: 51.1(5) vs AR 46.0, LLaDA-8B 65.9, Llama 3-8B 65.4. RACE: 41.4 vs 38.3~~39.5; TruthfulQA: 46.6 vs 34.3~~38.8. GSM8K: 32.8(4)로 AR 대비 낮지만 ARC-Easy 70.0, ARC-Challenge 43.0에 근접 또는 더 높다. BBH-Logical3: 31.8; MBPP: 26.6; HumanEval: 22.6; GPQA: 26.1. PIQA: 66.4; HellaSwag: 60.0; WinoGrande: 60.0. 결론적으로 Sumi는 지식·추론·코딩 벤치에서 강한 편이나 commonsense 벤치에서는 하락한다. 중간 학습 데이터 구성은 지식/추론 벤치에서 긍정적 영향을 주나, commonsense에서의 갭은 교육 중심 데이터의 영향으로 보일 수 있다.

기술 상세

아키텍처: 36-레이어, 은닉 차원 4096, SwiGLU, FFN 12288, 32-head, 8 KV 그룹, head_dim 128, RMSNorm, untied embedding, bias 없음, dropout 없음, attention의 off-by-one softmax. RoPE 500,000. 입력/출력 임베딩 정규화 없이 사용. 토크나이저: OLMo 3, vocab 100,278. 프레임워크: Megatron-LM. 학습: 7B 파라미터, 1.3T 토큰, canvas 2048, 배치 ≈ 5.5M 토큰/스텝. 최적화: AdamW, β=(0.9,0.95), weight_decay=0.1, gradient_clip=1.0, z-loss 1e-5. 데이터 구성: pre-training(llm-jp-corpus-v4 영어 파트 + code), mid-training(130B + 120B tokens)와 4_BUCKET 구성(코딩 81.4B, 수학 74.3B, 일반 52.4B, 추론 42.0B; 총 약 250B). 학습 일정: 2k warmup → 상수 구간 → 2k cooldown. 데이터 혼합: English llm-jp-corpus-v4.1 및 MegaMath-Web-Pro-Max 계열을 포함, nemotron_pretraining_code_v2의 비중 축소 및 수록 중복 제거 등 데이터 재정렬 포함.

한계점

Limitations: 추론 시간 분석은 방향성 있는 결과이며 30개의 샘플로 제한된다. Falcon 퍼플렉시티를 대략적 유창도 척도로 사용한다. mid-training 데이터 구성의 영향으로 commonsense 벤치에서의 격차가 크게 나타나며, 다른 UDLM 또는 AR 모델과의 직접 매칭 비교가 수행되지 않았다. Sumi는 instruction tuning이나 안전성 정렬 없이 pretrained되어, 안전 필터링이나 해석 가능성의 한계가 있다.

실무 활용

Sumi의 오픈형 UDLM 파라미터 및 학습 레시피를 통해 UDLM 연구 및 재현 가능한 비교 연구가 가능하다.

UDLM의 스케일링 및 데이터 구성 영향 연구
AR/MDLM와의 직접 비교 연구를 위한 재현 가능 벤치 구축
생성 도구에 대한 자체 검증 및 한계 분석
교육/코드 중심 데이터의 효과에 대한 추가 분석

코드 공개 여부: 공개

코드 저장소 보기

키워드

uniform diffusion language models (UDLMs)GIDD (Generalized Interpolating Discrete Diffusion)SNR 재매개변수화캔버스 길이confidence sampling(확신 샘플링)revision budget(리비전 예산)데이터 혼합(교육/코드 중심)

용어 해설

일반화된 보간 이산 확산(Generalized Interpolating Diffusion (GIDD)): — Generalized Interpolating Diffusion은 확산 기반 모델의 목적 함수를 보간적으로 재구성하는 프레임워크로, 토큰 수정의 순서와 흐름을 학습 시점에서 유연하게 제어하는 방법을 제공한다. GIDD는 ELBO의 재매개변수를 통해 확산 과정의 신호 대 잡음비(SNR) 관점에서 최적화를 수행하도록 한다.
SNR 재매개변수화: — SNR 재매개변수화는 GIDD의 ELBO를 신호 대 잡음비(SNR) 관점으로 재표현하는 형식으로, 모델이 노이즈 수준에 따른 업데이트를 보다 안정적으로 학습하도록 돕는다.
OLMo 토크나이저(OLMo tokenizer): — OLMo 토크나이저는 어휘 크기가 100,278로 설정된 토크나이저로, 학습 데이터 구성에 맞춰 토큰 효율성을 높이고 처리량을 최적화한다.
Megatron-LM: — Megatron-LM은 대규모 모델 학습을 위한 모델 파병 및 병렬화 프레임워크로, Sumi의 학습 프레임워크의 기반으로 사용된다.
RoPE: — RoPE는 rotary position embedding으로, Transformer의 위치 정보를 더 효과적으로 모델에 주입하기 위한 기법이다.