DiffusionGemma: 병렬 확산을 통한 4배 빠른 텍스트 생성

DiffusionGemma는 순차적 토큰 생성을 병렬 텍스트 확산으로 대체하여 4배 빠른 추론을 달성합니다

DiffusionGemma는 전통적인 자기회귀(autoregressive) 대규모 언어 모델(LLM)의 지연 시간 병목 현상을 제거하기 위해 설계된 실험적인 오픈 웨이트 모델입니다. 한 번에 하나의 토큰을 생성하는 대신 전체 텍스트 블록을 동시에 생성함으로써, 전용 GPU에서 최대 4배 빠른 텍스트 생성을 제공합니다. Apache 2.0 라이선스로 출시된 이 모델은 26B Mixture of Experts (MoE) 아키텍처를 사용하여 추론 시 3.8B 파라미터만 활성화하므로, 양자화 시 고사양 소비자용 GPU의 18GB VRAM 제한 내에 맞출 수 있습니다.

메모리 대역폭 병목 현상 해결

전통적인 LLM은 왼쪽에서 오른쪽으로 토큰을 순차적으로 생성하는 "타자기"처럼 작동합니다. 요청을 배치(batch)로 처리할 수 있는 고동시성 클라우드 서비스에서는 효율적이지만, 로컬의 단일 사용자 추론에는 비효율적입니다. 로컬 환경에서는 시스템이 실제로 계산을 수행하는 시간보다 RAM에서 프로세서로 가중치를 이동하는 데 더 많은 시간을 소비하기 때문에 GPU가 제대로 활용되지 못하는 경우가 많습니다.

DiffusionGemma는 병목 현상을 메모리 대역폭에서 연산으로 전환합니다. 단일 다음 토큰을 예측하는 대신, 256개 토큰의 단락을 동시에 초안 작성합니다. 이러한 "인쇄기" 방식은 하드웨어의 연산 능력을 포화시켜 로컬 가속기에서 상당한 속도 향상을 가져옵니다:

  • NVIDIA H100: 초당 1000개 이상의 토큰.
  • NVIDIA GeForce RTX 5090: 초당 700개 이상의 토큰.

기술적 아키텍처 및 기능

DiffusionGemma는 Gemma 4 제품군과 Gemini Diffusion 연구를 기반으로 구축된 새로운 확산 헤드(diffusion head)를 통합합니다. 작동 로직은 AI 이미지 생성과 유사합니다. 무작위 자리표시자 토큰으로 구성된 캔버스에서 시작하여 텍스트가 최종 출력으로 수렴될 때까지 여러 번의 패스를 통해 반복적으로 정제합니다.

주요 기술적 장점

  • 양방향 어텐션(Bi-directional Attention): 256개 토큰이 병렬로 생성되므로, 블록 내의 모든 토큰이 다른 모든 토큰을 참조할 수 있습니다. 이는 코드 인필링(code infilling), 인라인 편집, 수학적 그래프, 아미노산 서열과 같은 비선형 작업에 모델을 독특하게 적합하게 만듭니다.
  • 지능적 자기 수정: 모델은 전체 텍스트 블록을 한 번에 평가하여 정제 과정 중에 실시간으로 실수를 수정할 수 있습니다.
  • 하드웨어 최적화: 모델은 NVFP4 (4-bit floating-point) 커널을 지원하여 Hopper 및 Blackwell 아키텍처에서 손실에 가까운 정확도로 연산 처리량을 가속화합니다.

트레이드오프 및 실제 사용

DiffusionGemma는 원시 출력 품질보다 속도와 병렬 레이아웃을 우선시합니다. 전체적인 품질은 표준 자기회귀 Gemma 4 모델보다 낮습니다. 따라서 Google은 최대 품질을 요구하는 애플리케이션에는 표준 Gemma 4를 권장하며, DiffusionGemma는 속도가 중요한 대화형 로컬 워크플로우에 적합합니다.

개발자 생태계 및 통합

DiffusionGemma는 빠른 실험을 위해 설계되었으며 여러 주요 AI 프레임워크와 호환됩니다:

  • 서빙(Serving): MLX, vLLM (Red Hat 통합 포함), 그리고 Hugging Face Transformers를 통해 지원됩니다. llama.cpp에 대한 공식 지원은 대기 중입니다.
  • 미세 조정(Fine-tuning): 개발자는 Hackable Diffusion JAX 툴박스, Unsloth, 또는 NVIDIA NeMo를 사용할 수 있습니다. 주목할 만한 미세 조정 사례로는 Unsloth를 사용하여 모델이 Sudoku를 풀 수 있도록 하는 것이 있으며, 이는 양방향 어텐션이 자기회귀 모델보다 뚜렷한 장점을 제공하는 작업입니다.
  • 배포포함: Hugging Face, Gemini Enterprise Agent Platform Model Garden, 그리고 NVIDIA NIM을 통해 사용할 수 있습니다.

커뮤니티 인사이트 및 분석

개발자들 사이의 기술적 논의는 프롬프트와 응답 사이의 대시간을 줄임으로써 확산 모델이 "페어 프로그래밍" 경험을 혁사신할 수 있는 잠재력을 강조합니다.

"그것은 프롬프트하고 기다리는 최a SOTA agentic experience의 데어서가 아니라, 페어 프로그래밍 경험에 더 가까웠습니다... 프롬프트하고, 기다리고, 결과가 제대로 되었기를 바라기만 하는 슬롯머신 같은 느낌이 훨씬 줄어들었습니다." — @vineyardmike

다른 기여자들이 에지 디바이스에서의 메모리 대역폭 문제의 중요성을 강조합니다:

"에지 디바이스에서는 다른 문제가 있습니다. RAM에서 GB 단위의 가중치를 앞뒤로 옮기는 동안 추론 가속기가 굶주리고 있습니다... 확산 방식은 토큰을 병렬로 계산할 수 있어 메모리 대역폭 병목 현상을 완해해줍니다." — @samuelknight

특정 도메인에서의 모델 성능에 대한 질문은 남아 있습니다. 예를 들어 도구 호출(tool calling) 및 추론(reasoning)과 같은 분야입니다. 일부 사용자들은 모델의 재편집(re-edit) 기능이 단일일적인 모놀리식 도구 호출보다는 여러 파일에 걸친 일련의 편집 작업이나 "변경 스트림(change streams)"에 더 적합하다고 제안합니다.

Sources