기술 요약: Sumi – 처음부터 학습된 개방형 균등 확산 언어 모델 (Open Uniform Diffusion Language Model)

문제 정의

확산 모델(Diffusion Models)이 자기회귀(Autoregressive, AR) 언어 모델의 유망한 대안으로 부상했음에도 불구하고, **균등 확산 언어 모델(Uniform Diffusion Language Models, UDLMs)**을 대규모로 개발하는 데에는 상당한 격차가 존재한다. 마스크된 토큰을 비가역적으로 채우는 마스크 기반 확산 언어 모델(Masked Diffusion Language Models, MDLMs)과 달리, UDLM은 모든 토큰을 임의의 단계에서 업데이트할 수 있게 하여 이론적으로 더 유연한 생성 및 자기 수정(self-correction)을 가능하게 한다.

그러나 대규모 파라미터 규모(수십억 개의 파라미터)와 대규모 토큰 예산(수조 개의 토큰) 모두에서 처음부터 사전 학습된(pretrained from scratch) UDLM은 아직 존재하지 않는다. 기존의 대규모 UDLM은 계산 최적화된 체크포인트이거나 사전 학습된 AR 아키텍처(예: DiffusionGemma)를 변형한 모델들이다. 결과적으로, 데이터가 풍부한 환경에서 네이티브 대규모 균등 확산의 동작 방식은 아직 탐구되지 않은 영역이며, 스케일링 법칙(scaling laws), 생성 역학(generation dynamics), 그리고 확립된 AR 및 MDLM 베이스라인과의 트레이드오프를 연구하기 위한 깨끗한 참조점을 결여하고 있다.

방법론

모델 아키텍처 및 학습 목표

저자들은 완전 개방형 7B 파라미터 UDLM인 Sumi를 소개한다.

• 아키텍처: Sumi는 36개 레이어, 4,096의 은닉 크기(hidden size), SwiGLU MLP를 갖춘 시간 불가지론적(time-agnostic) 양방향 트랜스포머를 사용한다. 그룹 쿼리 어텐션(32개 헤드, 8개 KV 그룹), RMSNorm, 그리고 분리된(untied) 입출력 임베딩을 채택했다. 어텐션 싱크(attention sinks) 현상을 완화하기 위해 어텐션에 off-by-one 소프트맥스를 적용했다.
• 목표: 모델은 일반화된 보간 이산 확산(Generalized Interpolating Discrete Diffusion, GIDD) 목적 함수, 구체적으로는 신호 대 잡음비(SNR) 재매개변수화 형태를 사용하여 학습된다.
• 노이즈 스케줄: 학습에는 로그 SNR이 $\lambda \in [-9, 9]$ 범위로 제한된 순수 균등 노이즈(pure uniform noise)를 사용한다.
• 토큰화: OLMo 3 토크나이저를 사용하며 어휘 사전 크기는 100,278이다.

학습 체계

Sumi는 288대의 NVIDIA H100 GPU(총 약 43,308 GPU-시간)를 사용하여 1.5조 개의 토큰으로 처음부터 사전 학습되었다.

1. 사전 학습 (1.3T 토큰): 시퀀스 길이는 1,184, 글로벌 배치 크기는 4,608로 수행되었다. 데이터 혼합은 교육적 점수를 통해 필터링 및 재순위화된 교육 콘텐츠에 크게 가중치를 두었다(점수는 증류된 Qwen3-32B 분류기로부터 도출됨).
2. 중간 학습 (250B 토큰): 사전 학습 이후 두 단계가 이어진다:
   • 130B 토큰은 원래 길이(1,184)로 진행.
   • 120B 토큰은 확장된 길이(4,864)로 진행.
   • 데이터 혼합은 코드, 수학, 추론 중심으로 전환되었으며, 코딩 32.5%, 수학 29.7%, 일반 21.0%, 추론 16.8%로 구성된다.

추론 전략

• 캔버스 길이(Canvas Length): 모델은 모든 평가에서 2,048 토큰의 고정된 캔버스 길이를 사용한다. 이 길이는 상당히 짧거나 긴 캔버스에서는 성능이 급격히 저하되기 때문에 선택되었다.
• 샘플링: 저자들은 신뢰도 기반 적응형 샘플링(confidence-based adaptive sampling)과 조상 샘플링(ancestral sampling)을 비교 탐구한다.
• 평가 프로토콜: 모든 평가는 확산 점수 산출을 위해 수정된 lm-evaluation-harness를 사용한다. 가능도(likelihood) 태스크의 경우 정답 너머의 위치는 랜덤 토큰으로 채워지며, 생성 태스크의 경우 업데이트는 생성 길이에 제한된다.

주요 기여

1. 최초의 네이티브 대규모 UDLM: Sumi는 대규모 파라미터 규모(7B)와 대규모 토큰 예산(1.5T)에서 처음부터 사전 학습된 최초의 UDLM으로서, 네이티브 균등 확산을 연구하기 위한 베이스라인을 제공한다.
2. 개방형 공개: 저자들은 모델 가중치, 체크포인트, 그리고 공개된 코퍼스 상의 완전한 데이터 혼합 명세를 포함한 전체 학습 레시피를 공개한다.
3. 탐색적 추론 프로브: 논문은 캔버스 길이 민감도, 토큰 확정 순서(commitment order), 병렬 디코딩 능력, 자기 수정 잠재력을 포함한 생성 역학에 대한 방향성 있는 분석을 제시한다.

결과

벤치마크 성능

Sumi는 유사한 토큰 예산으로 학습된 오픈 AR 모델들(Falcon-7B, Llama 2-7B, OLMo-7B)과 함께 4개 카테고리의 13개 벤치마크에서 평가되었다.

• 일반 지식 및 코딩: Sumi는 경쟁력 있는 성능을 보이며, 동일한 프로토콜 하에 평가된 모델 중 MMLU, TruthfulQA, HumanEval, MBPP에서 최고의 점수를 기록했다. 이는 교육 및 코드 중심의 데이터 혼합 덕분이다.
• 추론 및 수학: Sumi는 GSM8K 및 ARC와 같은 태스크에서 Llama 2-7B와 경쟁할 만한 수준이며 Falcon-7B를 능가하지만, 더 크거나 다르게 학습된 모델(예: LLaDA-8B, Llama 3-8B)에는 뒤처진다.
• 상식(Commonsense): Sumi는 상식 벤치마크(PIQA, HellaSwag, WinoGrande)에서 크게 낮은 성능을 보인다. 저자들은 이를 교육 중심의 데이터 혼합 때문이라고 설명하며, 이러한 혼합이 상식 성능을 저하시킨다는 점이 관찰된 바 있다고 언급했다. 다만, 이 격차가 데이터 구성만으로 설명될 수 있는 것보다 더 클 수 있음을 시사했다.

추론 시 관찰 사항

• 캔버스 민감도: 생성 유창성은 캔버스 길이에 매우 민감하다. Sumi는 태스크 의존적 범위(2,048을 중심으로 함) 내에서는 유창하게 생성되지만, 이 범위를 벗어나면 특히 짧은 캔버스에서 성능이 급격히 저하된다.
• 확정 순서(Commitment Order): 학습 목적 함수는 순서에 무관함에도 불구하고, 신뢰도 기반 샘플링은 자기 조직화된 확정 순서를 유도한다. 모델은 결정 가능한 토큰을 일찍 확정하고 불확실한 토큰을 미루는 경향을 보이는데, 이는 조상 샘플링에서는 나타나지 않는 구조이다.
• 병렬 디코딩: 코딩 태스크(HumanEval, MBPP)에서 모델은 정확도 손실을 최소화하면서 적절한 병렬 디코딩(단계당 최대 4개 토큰 확정)을 지원한다. 그러나 다단계 산술(GSM8K)은 순서에 민감하여 병렬 단계가 적용되면 즉시 성능이 저하된다.
• 자기 수정(Self-Correction): 명시적인 수정 예산(추가 디노이징 패스 허용)을 두더라도 자기 수정은 발생하지 않는다. 모델이 확정된 토큰을 덮어쓰기는 하지만, 이러한 편집은 개선을 향한 직접적인 수정이라기보다 주로 "왕복(round trips, $A \to B \to A$)"에 불과하여 최종 정확도에는 변화를 주지 못한다.

의의 및 주장

본 논문은 Sumi를 네이티브 균등 확산을 대규모로 연구하기 위한 기초 자원으로 포지셔닝한다. 저자들은 자신들의 추론 프로브가 "결론적이라기보다는 방향성을 제시하는 것"임을 강조하며, 다음과 같이 잘 이해되지 않은 UDLM의 측면들을 연구하도록 동기를 부여하고자 한다:

• 마스크 기반 확산 대비 균등 확산의 본질적인 생성 행동.
• 자기 수정이 달성 가능해지는 조건.
• 데이터 구성이 특정 강점과 약점에 미치는 영향.

저자들은 Sumi가 지시어 튜닝(instruction tuning)이나 안전 필터링이 적용되지 않은 사전 학습된 베이스 모델로 공개되었음을 명시하며, 이는 직접적인 배포보다는 연구 지원을 목적으로 한다고 밝혔다. 또한, 이 모델이 유해한 출력 및 사실적 오류를 포함한 표준적인 베이스 LM의 위험을 그대로 가지고 있음을 경고한다.