이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.
TL;DR
이 글은 Google의 DiffusionGemma 26B 공개 가중치를 기반으로 모델 구성과 샘플러·어텐션·MoE 구현을 재구성해 내부 설계 선택을 추적한 실전 노트입니다. 가중치 검사에서 Q/K 정규화가 사실상 레이어별 스칼라로 수렴해 효과적 온도를 고정한다는 점과 full-attention의 partial RoPE, W_k와 W_v의 공유 같은 특이 설계가 드러났습니다. 캔버스 단위의 엔트로피 기반 재노이즈 샘플링과 self-conditioning 메커니즘을 코드로 재현해, 산문형 입력에서 왼쪽→오른쪽으로 토큰이 점진 확정되는 경향과 재노이즈 기반 오류정정이 실제로는 드물게 일어남을 관찰했습니다.
섹션별 상세
이 글은 Google이 공개한 DiffusionGemma 26B 가중치를 직접 로드해 모델 구성과 구현 세부를 재구성하면서 설계 의도를 규명하려는 실전형 작업 노트입니다. 구현은 safetensors와 PyTorch 기반으로 진행되며 어텐션 블록, RoPE, MOE, 샘플러 등 핵심 모듈을 코랩 노트북 수준에서 재현합니다. 이렇게 내부 구현과 가중치를 대조하면 문서화가 추상적이었던 부분을 수치와 코드로 보완할 수 있습니다.
근거
- Google은 DiffusionGemma가 특정 워크로드에서 Gemma 4보다 약 4배 빠르다고 보고했으며, 이 블로그는 공개 가중치와 노트북을 통해 그 성능 특성의 원인을 재현 가능한 수준으로 추적하려고 했습니다. — 초반 모델 소개와 노트북 실행 예시, 참고 문헌의 Google 블로그·모델 카드; 노트북에서의 캔버스 병렬 디코딩(256토큰)과 KV 캐시 암당화 관련 코드 및 설명을 확인하면 됩니다.
가중치 관찰에서 핵심적으로 확인된 것은 QK 정규화 벡터가 각 요소마다 동일한 값을 갖는 패턴이며, 이는 레이어별 학습 스칼라 γ_q·γ_k로 쿼리·키 크기를 고정하는 결과를 낳습니다. 수학적으로는 어텐션 로짓이 d_h·γ_q·γ_k·cosθ로 표현되어 레이어 단위의 유효 온도 T = 1/(d_hγ_qγ_k)가 고정됩니다. 이 구조는 쿼리 의존적 온도와 키 크기에 기반한 전역 중요도 인코딩을 제거해 긴 문맥에서 오래된 높은 크기 키가 불필요하게 영향력을 유지하는 문제를 완화할 가능성이 있습니다.
python
def rms(x, w = 1):
return (w * x * (x.shape[-1] ** 0.5) / (torch.norm(x, dim=-1, keepdim=True, dtype=torch.float32) + model_config['rms_norm_eps'])).to(x.dtype)이 함수는 RMSNorm의 핵심 연산을 구현한 것으로 입력 벡터의 L2 노름으로 나눈 뒤 학습된 가중치 w를 곱해 스케일을 적용합니다. 계산은 reduction이 민감하므로 내부적으로 float32로 노름을 구한 다음 원래 dtype으로 복원해 정밀도 누적 오류를 완화합니다. 이 연산은 Q/K/V와 MLP 입력의 전처리에서 반복적으로 호출되어 수치 안정성을 확보하는 역할을 수행합니다.
근거
- 모든 QK 정규화 가중치 벡터의 요소가 동일한 값으로 학습되어 각 레이어에 고정된 효과적 온도 T가 존재함이 가중치 프린트에서 확인됩니다. — 노트북에서 'model.decoder.layers.{i}.self_attn.{q,k}_norm.weight' 값을 출력한 결과와 수학적 유도(내적 = d_h·γ_q·γ_k·cosθ)를 근거로 검증할 수 있습니다.
DiffusionGemma의 위치 인코딩은 두 형태의 RoPE를 병용하며, full-attention에서는 partial RoPE를 적용해 주파수 스펙트럼의 후반을 0으로 만들어 해당 차원에서 회전을 무효화합니다. 또한 RoPE 페어링은 인접(pair t,t+1)이 아니라 split-half(pair t,t+d/2) 순서로 저장되어 캐시 친화성과 구현 단순화를 꾀합니다. 결과적으로 한 헤드는 명시적 위치 민감 차원과 위치 불변 차원을 동시에 가질 수 있어 로컬·글로벌 특징을 동시에 표현할 수 있습니다.
python
def rotate(x, layer_type, start_idx=0):
head_dim = x.shape[-1]
rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)')
pos = (torch.arange(x.shape[-2])[:, None] + start_idx) * torch.cat([freq[layer_type], freq[layer_type]])
return torch.cos(pos).to(x.dtype) * x + torch.sin(pos).to(x.dtype) * rot_x이 코드는 RoPE(회전 위치 임베딩)를 구현하는 부분으로, DiffusionGemma에서 채택한 split-half(pair t, t+d/2) 페어링 규칙을 반영합니다. 주파수 벡터 freq는 일부 레이어에서 partial RoPE로 후반 주파수를 0으로 세팅하므로 특정 차원은 회전이 무효화되어 위치 불변 표현을 보존합니다. 반환식은 각 위치별로 cos·sin 가중 조합으로 원래 벡터와 90도 회전벡터를 혼합해 최종 회전된 쿼리/키를 생성합니다.
근거
- Full-attention 레이어는 RoPE의 전체 주파수 스펙트럼을 절반 이상 무효화하는 partial RoPE를 사용해 앞쪽 25% 주파수만 회전시키도록 구성되어 있습니다. — 모델 설정의 rope_parameters.full_attention.partial_rotary_factor 값과 노트북의 freq 초기화 코드에서 후반 인덱스를 0으로 세팅하는 부분을 확인하면 됩니다.
디코드와 인코드 모드의 어텐션 동작이 구분되어 캔버스 기반 생성에서 디코드는 '커밋된 과거|현재 캔버스'를 모두 참조하고 인코드는 인과적 전처리(prefill)처럼 동작합니다. 하이브리드 구조에서는 대부분의 레이어가 슬라이딩 윈도우 어텐션을 사용하고 일부 전역 어텐션 레이어가 존재하며, 전역 레이어는 K와 V 프로젝션을 공유(W_v = W_k)하는 특이점이 있습니다. 구현상에는 KV 캐시 절약이나 재구성 전략과 인코드 최적화(마지막 레이어 이후 연산 생략) 등 실용적 트레이드오프가 적용되어 있습니다.
python
q, k, v = x @ self.W_q.T, x @ self.W_k.T, x @ self.W_v.T
q, k, v = [rearrange(z, 'l (n h) -> n l h', h = self.head_dim) for z in [q, k, v]]
q, k = rms(q, self.q_norm), rms(k, self.k_norm)
q, k = rotate(q, self.layer_type, pos_idx), rotate(k, self.layer_type, pos_idx)
# k/v cache concatenation and attention follows이 코드 조각은 어텐션 블록에서 Q/K/V를 계산하고 헤드 차원으로 재배열한 뒤 Q/K에 대해 QK-norm과 RoPE를 적용하는 전형적 흐름을 보여줍니다. DiffusionGemma에서는 full-attention 레이어에서 W_v를 W_k로 공유하는 예외가 있어 V는 별도 정규화 규칙(단위 RMS)으로 처리됩니다. 이후 k/v 캐시와의 연결, 스코어 계산, softmax 절차가 이어지며 encode/decode 모드에 따라 마스킹이나 캐시 업데이트 시점을 달리합니다.
근거
- Full-attention에서 W_v가 W_k와 동일하게 설정되어 있어 K와 V 프로젝션이 내부적으로 공유되고, 실제로는 K에만 RoPE와 스칼라 RMSNorm이 적용되는 차이가 존재합니다. — AttentionBlock 초기화 코드에서 full_attention일 때 self.W_v = self.W_k로 설정한 부분과 이후 k/v에 적용되는 서로 다른 정규화·RoPE 처리 코드를 근거로 확인할 수 있습니다.
샘플링 루프는 캔버스별 반복 인코드와 최대 max_denoising_steps 디코드를 수행하며, 각 디코드 스텝에서 온도 t_i를 선형 스케줄로 내림으로써 분포을 점차 날카롭게 만듭니다. 각 위치의 엔트로피를 계산해 낮은 엔트로피의 접두사를 고정하고 나머지를 균등 무작위 토큰으로 다시 노이즈화하는 방식으로 재노이즈/수용을 반복합니다. 실험적 관찰에서는 산문형 프롬프트에서 캔버스가 좌→우로 '사실상' 확정되는 경향이 강했고, 퍼즐형 입력에서는 다른 패턴이 나타나며 Google이 주장한 재노이즈 기반 오류정정이 실제로 자주 일어나지는 않았습니다.
근거
- 샘플링 절차는 각 디코드 스텝에서 온도 스케줄을 선형으로 내리며 위치별 엔트로피를 계산해 엔트로피 바운드 내의 접두사를 고정하고 나머지를 재노이즈하는 알고리즘을 따릅니다. — generation_config의 t_max·t_min·max_denoising_steps·sampler_config.entropy_bound 항목과 노트북의 sampling 루프 코드를 확인하면 온도·엔트로피 기반 수용 로직을 재현할 수 있습니다.
- 실험적 관찰에서는 산문 프롬프트에서 캔버스 토큰이 좌→우로 사실상 확정되는 패턴이 자주 관찰되었고, 재노이즈로 인한 오류정정(다시 마스킹되는 빈도)은 전체적으로 매우 드물었습니다. — 저자가 저장한 plotting_data를 이용한 엔트로피·확정 시점 시각화(논문 내 캔버스 그림)와 'Preliminary Analysis' 절의 캔버스별 재마스킹 빈도 그래프를 근거로 합니다.
용어 해설
- 스칼라 Q/K 정규화(Scalar QK Norm)
- — 쿼리·키 정규화 벡터가 전체 차원에서 같은 값으로 학습된 경우를 가리키며, 각 레이어마다 학습된 두 스칼라 γ_q·γ_k로 q, k의 크기를 고정해 유효한 '온도'를 레이어 단위로 설정합니다. 이로 인해 q·k 내적은 크기 정보가 아닌 각도(유사도)에 의해 주로 결정되며 장문 문맥에서 고정 크기 키가 과도하게 영향력을 갖는 문제를 완화할 수 있습니다. 논리적으로는 q·k 곱에 상수인 d_hγ_qγ_k가 곱해져 소프트맥스의 샤프니스(효과적 온도)를 레이어 수준에서 고정합니다.
- 부분 RoPE(Partial RoPE)
- — RoPE 주파수 스펙트럼을 구성한 뒤 상위 또는 후반 주파수의 일부를 0으로 둬 회전(rotation)을 적용하지 않는 차원들을 만드는 기법입니다. DiffusionGemma에서는 full-attention 레이어의 RoPE에서 앞부분 25% 주파수만 유효하게 남기고 나머지는 회전을 무효화해 위치 민감 표현과 비위치 표현을 같은 헤드 안에서 공존시킵니다. 이 방식은 위치 의존성의 일부만을 유지하며 로컬·글로벌 특징을 동시에 다루는 목적에 맞추어 설계되어 있습니다.
- 레이어 스칼라(Layer Scalar)
- — 각 레이어의 잔차스트림 출력을 학습된 스칼라로 재조정하는 매커니즘이며, 특히 저정밀도(fixed dynamic range) 환경에서 활성값의 과도한 증가를 억제하기 위해 도입됩니다. 구현상 각 레이어 계산 후 x_{l+1} = s_l · F(...) 형태로 적용되어 누적된 잔차가 FP8·FP16 등으로 주조될 때 오버플로우·언더플로우 위험을 줄입니다. 학습된 s_l 값은 레이어별로 다른 스케일을 부여하므로 안정성·정밀도 관점에서 중요한 하이퍼파라미터 역할을 합니다.
- 자기 조건화(Self-conditioning)
- — 이전 디노이징 단계의 소프트맥스 분포 P를 임베딩 테이블 E에 곱해 기대 임베딩 P·E를 만들고 이를 작은 MLP로 가공해 현재 캔버스 입력에 더하는 방식입니다. 이렇게 하면 이전 단계의 확신도와 분포 모양이 현재 예측에 신호로 들어가서 불확실한 위치를 더 잘 다루거나 마스킹된 토큰을 추론하는 데 도움이 됩니다. 첫 디코드 스텝에서는 입력이 0이므로 영향이 없고, MLP가 영벡터를 영벡터로 매핑하도록 설계되어 초기 단계에서 일관성이 유지됩니다.
- 혼합 전문가(MoE)(Mixture-of-Experts)
- — 입력 토큰마다 상위 k개의 전문가에 라우팅해 각 전문가를 병렬로 실행한 뒤 가중합하는 MLP 구조입니다. DiffusionGemma는 토큰당 8개 전문가를 선택하고 별도의 공유 전문가(shared expert)를 포함해 합산하며, 라우팅 스코어는 라우터 키와 입력 내적의 지수값으로 산출한 뒤 재정규화·스케일링합니다. 구현 관점에서는 토큰 축 독립성 덕분에 전문가별 토큰 배치를 모아서 처리하는 방식이 성능상 유리합니다.
- 슬라이딩 윈도우 어텐션(Sliding-window Attention)
- — 캔버스 내부의 토큰 모두와 캔버스 앞의 고정 길이 접두사(prefix)의 KV만을 참조하는 국소 어텐션 스케줄입니다. DiffusionGemma의 하이브리드 구조에서는 슬라이딩 레이어 비중이 높고 일부 레이어에서만 전역(full) 어텐션을 사용해 연산·캐시 비용을 균형시킵니다. 이 방식은 로컬 컨텍스트를 병렬로 처리하면서 인접 블록과의 연결성을 보존하는 데 유리합니다.
기술
- PyTorch
- safetensors
- einops
- HuggingFace tokenizer
- vLLM
활용 사례
- 로컬 상호작용성이 중요한 저지연 대화형 환경에서 캔버스 병렬 디코딩으로 응답성을 높이는 용도
- 긴 문맥을 다루는 시스템에서 고정된 키 크기로 불필요한 장기 영향력을 억제해 안정성을 개선하는 환경
- 대규모 MoE 모델을 라우팅·캐시 최적화와 결합해 토큰별 전문가 분배를 효율적으로 처리하는 실무 배포
AI 분석 전체 내용 보기
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
원문 발행 2026. 08. 11.수집 2026. 08. 11.출처 타입 RSS
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.
