출시 블로그 포스트에서의 발표에 이어, 이 실험적인 모델을 이해하고, 서빙(serve)하며, 커스텀(customize)할 수 있도록 돕기 위한 개발자 가이드를 공유합니다.

Gemma 4 백본(backbone)을 기반으로 구축된 DiffusionGemma는 개발자 워크플로(workflow)를 위한 몇 가지 이정표를 제시합니다:

1. 연산 제한적 병렬 생성 (Compute-bound parallel generation): 병목 현상을 메모리 대역폭(memory-bandwidth)에서 연산(compute)으로 전환함으로써 메모리 대역폭의 제한을 우회하며, GPU에서 최대 4배 빠른 토큰 생성 속도를 제공합니다 (NVIDIA GeForce RTX 5090에서 초당 최대 700개 이상의 토큰, 단일 NVIDIA H100에서 초당 1000개 이상의 토큰 생성).
2. 양방향 컨텍스트 및 자기 수정 (Bidirectional context & self-correction): 생성 과정에서 전체 텍스트 블록을 동시에 평가하기 위해 양방향 어텐션(bidirectional attention)을 사용하여, 실시간 오류 수정 및 병렬 컨텍스트 전파(parallel context propagation)를 가능하게 합니다.
3. 개발자 친화적인 크기: 추론(inference) 시 3.8B 파라미터만 활성화되는 26B MoE (Mixture of Experts) 모델로 설계되어, 18 GB VRAM 제한 내에서 양자화(quantized) 배포가 가능합니다.

아키텍처 (The Architecture)

GPU에서 전통적인 LLM(대규모 언어 모델)을 사용하여 구축하는 개발자들에게 주요 병목 현상은 메모리 대역폭입니다. 자기회귀(Autoregressive) 언어 모델은 텍스트를 한 번에 하나의 토큰씩 생성하기 위해 메모리에서 모델 가중치(weights)를 반복적으로 로드해야 합니다. DiffusionGemma는 병목 현상을 메모리 대역폭에서 연산으로 전환하여 **256-토큰 캔버스(256-token canvas)**를 병렬로 생성하고 정제함으로써 이러한 제한을 우회합니다. GPU에 대규모 병렬 워크로드(workload)를 제공함으로써, 로컬 서빙(local serving) 중에 유휴 상태로 남을 수 있는 텐서 코어(tensor cores)를 활용합니다.

• 균일 상태 확산 (Uniform State Diffusion): 토큰을 순차적으로 예측하는 대신, DiffusionGemma는 무작위 플레이스홀더 (placeholder) 토큰으로 구성된 캔버스에서 시작하여 이를 병렬적으로 반복해서 정제합니다. 여러 번의 디노이징 (denoising) 패스를 거치면서, 높은 확신도를 가진 토큰들이 인접한 위치를 해결하는 데 도움을 주어 전체 시퀀스가 초점에 맞춰지듯 정렬됩니다.
• 가변 길이 생성을 위한 블록 자기회귀 확산 (Block Autoregressive Diffusion for Variable Length Generation): 256개 토큰보다 긴 시퀀스의 경우, 256개 토큰 블록이 완전히 디노이징되면 모델은 이를 처리하여 KV 캐시 (KV cache)에 저장합니다. 그런 다음 모델은 이전에 저장된 이력 (history)을 조건으로 하는 새로운 256개 토큰 캔버스를 초기화하며 다음 블록으로 전환합니다. 이는 병렬 블록의 속도와 자기회귀 (autoregressive) 모델의 순차적 안정성을 결합한 방식입니다.

쇼케이스: 병렬 디노이징을 이용한 스도쿠 풀이

전통적인 자기회귀 (autoregressive) 모델은 스도쿠와 같이 엄격하고 다변수 제약이 있는 문제에서 어려움을 겪습니다. 텍스트를 엄격하게 왼쪽에서 오른쪽으로 생성하기 때문에, 미래의 플레이스홀더를 평가하거나 되돌아가는 (backtrack) 것이 불가능하기 때문입니다.

DiffusionGemma의 커스터마이징을 보여주기 위해, 우리는 모듈형 JAX 연구 툴박스인 Hackable Diffusion을 사용하여 파인튜닝 (fine-tuning) 레시피 및 결과를 공개합니다. 이 학습 설정은 고전적인 다변수 그리드 작업인 **스도쿠 해결사 (Sudoku Solver)**에 초점을 맞춥니다.

Diffusion에 스도쿠가 흥미로운 이유

81글자의 스도쿠 문자열 표현(빈 셀은 마침표로 표시됨)에서, 모든 숫자는 엄격하게 교차하는 가로, 세로 및 9x9 그리드 제약 조건에 묶여 있습니다.

양방향 문맥 전파 (Bidirectional Context Propagation): 자기회귀 (Autoregressive) 모델과 달리, DiffusionGemma의 노이즈 제거 (Denoising) 단계는 모든 캔버스 쿼리가 모든 위치를 병렬로 참조(Attend)할 수 있게 합니다. 정보는 보드 전체에 걸쳐 대칭적으로 흐르며, 각 단계에서 전역적 의존성 (Global dependencies)을 해결합니다.

• 재노이징 (Re-Noising)을 통한 오류 수정: 균일 상태 확산 (Uniform State Diffusion) 하에서 모델은 보드 전체를 동시에 평가합니다. 신뢰도가 떨어지면 샘플러(Sampler)가 숫자를 무작위 숫자로 교체하여 지속적인 자기 수정 (Self-correction)이 가능하게 합니다.
• 효율적인 조기 종료 (Early Stopping): 스도쿠 (Sudoku) 데이터셋을 통한 미세 조정 (Fine-tuning) 결과, 어댑터 (Adapters)가 조기 종료 성능을 향상시키는 것으로 나타났습니다. SFT (Supervised Fine-Tuning)가 완료된 모델은 베이스 모델보다 더 빠르게 안정화되어, 엔진이 더 일찍 멈출 수 있게 함으로써 지연 시간 (Latency)과 연산 비용을 줄여줍니다.

controls
loop
muted
autoplay
playsinline
width="100%"
style="border-radius: 8px; display: block; margin-bottom: 0; padding-bottom: 0px;"
aria-label="">

Your browser does not support the video tag.

왼쪽: DiffusionGemma가 스도쿠 출력을 생성하는 모습. 베이스 모델은 48단계 이후에도 스도쿠를 풀지 못합니다. 오른쪽: 미세 조정(SFT)된 DiffusionGemma는 12단계 만에 퍼즐을 해결합니다. 적응형 종료 (Adaptive stopping) 덕분에 조기에 완료할 수 있습니다.

성능 영향: 베이스 DiffusionGemma 모델은 스도쿠 퍼즐을 풀도록 특별히 훈련되지 않았지만 (~0% 성공률), 스도쿠 데이터셋에 간단한 JAX SFT 레시피를 적용하면 정확도가 80% 성공률로 상승하는 동시에 전체 추론 단계 수(Inference step count)는 감소합니다.

블록 자기회귀 노이즈 제거 (Block Autoregressive Denoising)

블록 자기회귀 노이즈 제거를 활성화하기 위해, DiffusionGemma는 추론 과정에서 점진적 프리필 (Incremental prefill)과 노이즈 제거 (Denoising)를 교대로 수행합니다:

• 프리필 / 증분 프리필 (Incremental Prefill) (인과적): _인과적 어텐션(causal attention)_을 사용하여 프롬프트 컨텍스트를 수집하고 KV 캐시에 기록합니다. 이는 초기 컨텍스트를 미리 채우기 위해 한 번 실행되고, 이후 다음 캔버스를 디노이징하기 전에 각 확정된 256 토큰 캔버스를 KV 캐시에 추가하는 작업을 블록마다 한 번 수행합니다.
• 디노이징 (Denoising) (양방향): _양방향 어텐션(bidirectional attention)_을 사용하여 캔버스를 반복적으로 디노이즈합니다. 캔버스상의 임의 위치에 있는 Query 토큰은 다른 모든 캔버스 토큰(및 KV 캐시)에 주의를 기울일 수 있어, 모델이 컨텍스트를 양방향으로 처리할 수 있게 합니다.

이러한 아키텍처 선택은 다음을 가능하게 합니다:

• 전역 컨텍스트 인식 (Global Context Awareness): 단순히

DiffusionGemma 서빙하기 (Serving DiffusionGemma)

이 실험적인 아키텍처를 효율적으로 서빙하기 위해, 우리는 vLLM 팀과 협력하여 DiffusionGemma를 vLLM에 구현했습니다. 이러한 통합을 통해 엔진은 배치된 요청 스트림(batched request streams) 전반에 걸쳐 반복적인 병렬 디노이징 루프(iterative parallel denoising loops)를 효율적으로 실행할 수 있습니다.

개발자들은 vLLM의 표준 OpenAI 호환 로컬 서버를 사용하여 DiffusionGemma를 즉시 배포할 수 있습니다.

vllm serve google/diffusiongemma-26B-A4B-it \
  --max-model-len 262144 \
  --max-num-seqs 4 \
...

지금 바로 시작하기 (Getting Started Today)

비자기회적(non-autoregressive) 텍스트 생성의 최전선을 탐험할 준비가 되셨나요? 더 자세한 내용을 확인하려면 다음 리소스를 참조하세요:

• 가중치 다운로드 (Download the Weights): Hugging Face에서 실험적 모델 (Apache 2.0 라이선스로 출시됨)의 가중치에 직접 접근할 수 있습니다.

• 선호하는 추론 프레임워크 사용 (Use Your Favorite Inference Frameworks): vLLM, Hugging Face Transformers, SGLang, 그리고 MLX를 사용하여 모델을 효율적으로 실행하세요.

• 적응 및 미세 조정 (Adapt & Fine-Tune): 신속한 실험을 위해, Hackable Diffusion을 사용하는 공식 학습 레시피 (official training recipes)를 공개합니다. 또한 Unsloth 또는 NVIDIA NeMo를 사용하여 효율적인 미세 조정 (fine-tuning)을 탐색할 수 있습니다.

• 원하는 방식으로 배포 (Deploy Your Way): Model Garden을 사용하여 Google Cloud에 즉시 배포하거나 NVIDIA NIM을 통해 배포하세요.

이 모델은 소비자용 RTX 4090 및 5090 카드부터 엔터프라이즈급 Hopper 및 Blackwell 서버에 이르기까지 하드웨어 스택 전반에 걸쳐 네이티브하게 최적화되어 있습니다.