블록 디퓨전¶

자기회귀 모델은 왼쪽에서 오른쪽으로 한 번에 하나씩 되돌릴 수 없는 선택을 합니다. 디퓨전 언어 모델은 불확실한 토큰 영역을 유지하면서 여러 위치를 반복해서 수정하거나 확정합니다. 블록 디퓨전은 이 영역의 크기를 제한합니다. 블록 하나를 생성해 접두 시퀀스에 확정한 뒤 다음 블록을 생성하므로 블록 내부의 병렬 작업과 블록 사이의 인과 순서를 결합합니다.

mlxcel은 서로 연관되지만 다른 두 계열을 구현합니다.

• DiffusionGemma (diffusion_gemma, diffusion_gemma_text): 토큰 캔버스를 초기화하고 엔트로피 또는 신뢰도 기반 승인 규칙으로 노이즈를 제거합니다. 멀티모달 형태는 이미지를 지원합니다.
• LLaDA-2 MoE (llada2_moe): 생성 위치를 <|mask|>로 초기화하고 고정 크기 텍스트 블록 안에서 신뢰도가 높은 위치의 마스크를 반복해서 해제합니다.

둘 다 고유 대상 디코더이며 DFlash 초안 모델이 아닙니다. DFlash는 자기회귀 대상 모델이 검증할 토큰 제안을 만들 때만 디퓨전을 사용합니다. MTP와 DFlash를 참조하십시오.

자기회귀에서 노이즈 제거로¶

자기회귀 대상 모델은 p(x_t | x_<t)를 추정하고 t+1로 이동한 뒤에는 x_t를 수정할 수 없습니다. 이산 디퓨전 또는 마스크 모델은 손상된 상태 z_s를 보고 잡음이나 마스크 수준이 낮아질수록 더 깨끗한 토큰 값을 예측합니다.

z_S -> z_(S-1) -> ... -> z_0 = x

여러 위치가 캔버스에 동시에 존재하므로 한 번의 모델 평가에서 여러 위치를 바꿀 수 있습니다. 잠재적인 장점은 “순전파 한 번으로 답 전체를 생성한다”는 것이 아니라, 노이즈 제거 단계마다 유효한 병렬 토큰 작업을 수행할 수 있다는 점입니다. 실제 비용은 블록 길이, 모델 평가 횟수, 신뢰도가 높아지는 속도, 어텐션 연산 형태가 하드웨어를 효율적으로 사용하는지에 달려 있습니다.

블록 생성은 답 전체를 하나의 캔버스에 두었을 때 생기는 실전 문제도 해결합니다.

• 메모리와 어텐션 비용이 설정한 블록 또는 캔버스 크기로 제한됩니다.
• 완료된 블록이 안정된 접두 시퀀스가 되므로 긴 답 전체를 매번 다시 생성하지 않습니다.

대신 한 번 확정한 블록은 보통 뒤의 블록을 처리할 때 수정할 수 없습니다. 작은 블록은 인과성이 더 강하고 스트리밍을 일찍 시작할 수 있습니다. 큰 블록은 내부 병렬성과 문맥을 더 많이 제공하지만 확정이 늦어집니다.

mlxcel의 두 생성 반복문¶

DiffusionGemma 캔버스 노이즈 제거¶

1. 프롬프트를 밀집 FP16 KV 캐시의 읽기 전용 접두 시퀀스로 인코딩합니다.
2. min(max_canvas, max(remaining, min_canvas)) 길이의 캔버스를 할당하고 시드가 적용된 MLX 난수 생성기로 초기화합니다.
3. 학습된 자기 조건화와 함께 현재 캔버스를 임베딩합니다. 캔버스 위치끼리는 양방향 어텐션을 사용하고 확정된 접두 시퀀스에는 인과적으로 어텐션합니다.
4. 선택한 샘플러로 새 예측을 승인합니다.
5. 설정한 안정 구간 동안 캔버스가 안정되고 신뢰도가 충분히 높아지거나 노이즈 제거 단계 예산을 모두 쓰면 종료합니다.
6. 캔버스를 접두 시퀀스에 확정하고 출력 예산이 남으면 다음 캔버스를 시작합니다.

LLaDA-2 블록 마스크 해제¶

1. 프롬프트와 생성 영역을 블록으로 나눕니다. 미래 생성 영역은 <|mask|> 토큰으로 시작합니다.
2. 프롬프트 블록을 각 계층의 KV 캐시에 확정합니다.
3. 활성 생성 블록을 고정된 접두 시퀀스와 함께 반복해서 순전파합니다.
4. 선택된 토큰의 신뢰도가 현재 임계값을 넘는 모든 마스크 위치를 확정합니다.
5. 임계값을 통과한 위치가 없으면 활성 위치 중 신뢰도가 가장 높은 하나를 확정하여 매 반복마다 진행을 보장합니다.
6. 완전히 확정된 블록을 캐시에 반영하고 시퀀스 순서대로 디코딩한 토큰을 스트리밍한 뒤 다음 블록으로 이동합니다.

LLaDA-2 기본값은 생성 블록 길이 32, 32단계, 초기·최저 신뢰도 임계값 0.95입니다. 엔진 내부에는 선형으로 낮아지는 임계값이 있지만 현재 공개 명령줄은 단계 수 재정의만 제공합니다. 블록 길이와 임계값 일정은 모델과 런타임의 기본값을 사용합니다.

샘플러와 확정 정책¶

DiffusionGemma는 --diffusion-sampler로 단계별 샘플러 두 가지를 제공합니다.

• entropy-bound(기본값): 위치별 엔트로피를 계산하고 낮은 순서로 정렬한 뒤 누적 엔트로피 규칙이 체크포인트의 entropy_bound를 넘을 때까지 앞쪽 위치를 승인합니다. 최소 한 위치는 승인합니다. mlx-vlm 참조 규칙을 따릅니다.
• confidence-threshold: 1순위 확률이 --diffusion-threshold(기본값 0.9)를 넘는 미확정 위치를 승인합니다. 해당 위치가 없으면 미확정 위치 중 신뢰도가 가장 높은 하나를 승인하여 진행을 보장합니다.

신뢰도 정책은 쉬운 위치를 초기에 많이 확정하여 품질 일부와 단계 수를 교환할 수 있습니다. 엔트로피 정책은 체크포인트가 보정한 경계값으로 불확실성에 따라 작업량을 조절합니다. “단계당 승인 토큰 수”만 비교하지 말고 최종 과제 품질과 종단 간 지연을 평가하십시오.

LLaDA-2는 자체 신뢰도 기반 전이 마스크를 사용합니다. --diffusion-sampler와 --diffusion-threshold는 이를 바꾸지 않습니다. 온도가 0보다 높을 때 일반 --temp, --top-k, --top-p 설정이 선택 토큰 샘플링에 적용됩니다.

명령줄 참조¶

DiffusionGemma¶

mlxcel generate \
  --model diffusiongemma-26B-A4B-it-4bit \
  --prompt "Explain TCP and UDP in one paragraph." \
  --max-tokens 256 \
  --max-denoising-steps 48 \
  --diffusion-sampler entropy-bound \
  --seed 42 \
  --profile

LLaDA-2 MoE¶

--max-denoising-steps는 LLaDA-2에서 블록당 단계 수를 재정의할 때도 사용됩니다. 다른 --diffusion-* 샘플러, 임계값, 캔버스 제어는 DiffusionGemma용이며 LLaDA-2의 설정을 바꾸지 않습니다.

mlxcel generate \
  --model ./models/LLaDA2.0-mini \
  --prompt "Explain why retrieval and reranking are separate stages." \
  --max-tokens 256 \
  --max-denoising-steps 32 \
  --temp 0 \
  --seed 42 \
  --profile

LLaDA-2는 --draft-model을 거부합니다. 고유 디코더가 이미 블록의 마스크를 해제하므로 런타임은 그 위에 투기적 규약을 추가로 겹치지 않습니다.

멀티모달 입력¶

DiffusionGemma 멀티모달 체크포인트는 여러 번 지정할 수 있는 --image <path> 인자를 받습니다.

mlxcel generate \
  --model google/diffusiongemma-26B-A4B-it \
  --image photo.jpg \
  --prompt "Describe what you see." \
  --max-tokens 256 \
  --seed 42

Gemma 4 비전 파이프라인을 재사용합니다. 각 이미지를 768×768로 크기 조정하고 여백을 채운 뒤 256개의 소프트 토큰으로 인코딩하여 이미지 센티널 토큰 사이에 넣습니다. 이미지 블록은 프리필 중 양방향 어텐션을 사용합니다. 오디오와 비디오는 거부합니다.

LLaDA-2는 텍스트 전용이며 이미지, 오디오, 비디오 입력을 모두 거부합니다.

서버 방식¶

두 계열 모두 전용 단일 스트림 작업자를 사용합니다. 연속 배칭 스케줄러에 참여하지 않고 요청을 큐에 넣어 하나씩 처리합니다. 스트리밍은 지원하지만 내부 위치는 시퀀스 순서와 다르게 수렴할 수 있습니다. 따라서 mlxcel은 불안정한 중간 토큰을 노출하지 않고 블록이 확정되었을 때 일반 시퀀스 순서로 텍스트를 보냅니다.

mlxcel serve \
  --model diffusiongemma-26B-A4B-it-4bit \
  --diffusion-sampler entropy-bound \
  --max-denoising-steps 48 \
  --port 8080

서버는 채팅 및 완성 경로를 제공합니다. DiffusionGemma 이미지 요청은 일반 OpenAI image_url 콘텐츠 부분을 사용합니다. 지원하지 않는 미디어가 들어오면 해당 요청에 오류를 반환하지만 작업자를 종료하지는 않습니다.

DiffusionGemma 응답 동작¶

DiffusionGemma는 사고 과정 채널을 사용하는 모델입니다. 충분한 max_tokens를 지정하십시오. 64토큰 예산은 추론 과정에 모두 사용되어 사용자에게 보이는 content가 비어 있을 수 있으므로, 짧은 프롬프트에는 256이 더 실용적인 시작점입니다. 스트리밍 채팅 델타에는 reasoning_content가 들어갈 수 있고 비스트리밍 채팅 응답은 추론 과정을 생략하는 mlxcel 규칙을 따릅니다. 원시 /v1/completions도 기계적으로 동작하지만 채팅 템플릿 사용을 강하게 권장합니다.

취소 여부는 노이즈 제거 단계마다 한 번 확인하므로 즉시 중단되는 것이 아니라 현재 단계의 실행 시간 안에 처리됩니다.

재현성¶

--seed N으로 MLX 난수 생성기를 초기화합니다. 체크포인트, 프롬프트, 생성 설정이 같고 온도가 0이면 DiffusionGemma는 같은 캔버스 초기화와 확정 흐름을 사용하고 LLaDA 샘플링도 같은 시드의 난수원을 사용합니다. 커널이나 버전이 바뀌면 부동소수점 값이 거의 같은 토큰의 순위는 여전히 달라질 수 있습니다.

MLXCEL_DIFFUSION_DEBUG_CANVAS=1은 Python 참조 구현과의 동등성을 확인하기 위해 DiffusionGemma 난수 초기화를 고정된 디버그 무늬로 바꿉니다. 일반적인 품질 생성 방식이 아닙니다.

성능을 해석하는 법¶

유효 토큰 B개인 블록과 노이즈 제거 순전파 S회를 다음처럼 근사할 수 있습니다.

useful rate ~= B / (S * t_forward(B) + t_commit)

이는 의도적으로 자기회귀 모델의 초당 토큰 수와 다른 식입니다. 큰 블록은 유효 작업을 늘리지만 어텐션과 각 순전파 비용도 늘립니다. 낮은 단계 예산은 시간을 줄이지만 덜 수렴한 토큰을 확정할 수 있습니다. 조기 종료 또는 반복당 여러 위치 확정은 유효한 S를 바꿉니다.

M1 Ultra와 4비트 DiffusionGemma 26B-A4B 체크포인트로 측정한 값은 다음과 같습니다.

• 48단계 중 약 12단계에 수렴한 64토큰 캔버스의 짧은 답에서 약 20토큰/초
• 256토큰 캔버스에서 48단계 대부분을 쓴 긴 출력에서 약 8토큰/초
• 모델 상주 메모리 약 13.5GiB

특정 하드웨어와 프롬프트에 한정된 관찰이며 보편적인 모델 성능 주장이 아닙니다. 실제로 서빙할 작업과 같은 모델 리비전으로 다시 측정하십시오.

품질 및 조정 체크리스트¶

• 체크포인트의 샘플러 및 단계 기본값에서 시작합니다.
• 속도뿐 아니라 답 품질, 형식, EOS 동작, 반복을 평가합니다.
• 짧고 긴 출력을 비교합니다. 마지막의 일부만 채운 블록은 동작이 다를 수 있습니다.
• 최대 설정만 기록하지 말고 실제 단계 수와 조기 종료를 기록합니다.
• 확률적 샘플링에서는 여러 시드를 비교합니다.
• 서버는 요청을 직렬로 처리하므로 지연 시간에 큐 대기를 포함합니다.
• 디퓨전 모델과 자기회귀 모델을 원시 초당 토큰 수만으로 비교하지 말고 모델 크기, 과제 품질, 첫 토큰 시간, 출력 길이의 의미를 함께 봅니다.

현재 제한¶

더 읽을거리¶

• Google AI for Developers, DiffusionGemma 모델 카드
• Google AI for Developers, DiffusionGemma 동작 설명
• Nie 외, Large Language Diffusion Models