구글이 최근 공개한 260억 개 매개변수(26B) 규모의 오픈 가중치 언어모델인 디퓨전젬마(DiffusionGemma)가 특정 워크로드에서 기존 젬마 4(Gemma 4) 대비 약 4배 빠른 속도를 달성하며 주목받고 있습니다. 이 실험적인 모델은 공식 백서 없이 간략한 문서만 제공되었는데, ML 앳 버클리(ML@B) 연구진이 콜랩(Colab) 노트북 기반의 심층 분석을 통해 그 내부 작동 방식을 상세히 파헤쳤습니다.
ML 앳 버클리 연구진은 디퓨전젬마의 오픈 가중치를 활용해 모델을 재구현하고, 그 과정에서 드러난 독특한 설계 선택들을 설명했습니다. 특히, 모델 아키텍처, 셀프 컨디셔닝(self-conditioning), 샘플링 절차, 인코더/디코더 가중치 공유 구현 방식 등 불투명했던 세부 사항들을 명확히 했습니다. 이들은 모델 구성과 가중치를 면밀히 검토하여 '레이어 스칼라(layer scalar)', '부분 로터리 임베딩(Partial RoPE)', 그리고 구글의 '스칼라 QK 정규화(Scalar QK Norm)'와 같은 흥미로운 디자인 선택들을 발견하고 이에 대한 가능한 설명을 제시했습니다.
이 중 가장 핵심적인 발견은 '스칼라 QK 정규화'입니다. 일반적인 QK 정규화 모델에서는 쿼리(query)와 키(key)의 크기가 가변적이어서 어텐션 스코어가 쿼리마다 다른 유효 온도(effective temperature)를 가지거나, 키의 크기가 쿼리와 무관한 '전역 중요도(global importance)'를 인코딩할 수 있습니다. 이는 긴 문맥(long-context) 환경에서 초기에 등장한 키가 관련성이 떨어진 후에도 과도하게 영향을 미 미칠 수 있는 단점을 가집니다. 그러나 디퓨전젬마의 스칼라 QK 정규화는 쿼리와 키의 크기를 고정된 스칼라 값으로 정규화하여, 어텐션 스코어가 오직 쿼리와 키 사이의 유사도(코사인 유사도)만을 순수하게 반영하도록 만듭니다. 이는 어텐션 계산의 효율성을 높이고, 모델이 불필요한 노이즈 없이 핵심적인 정보에 집중할 수 있게 하여 전반적인 속도 향상에 기여하는 것으로 해석됩니다. 이러한 최적화된 어텐션 메커니즘은 대규모 언어모델의 추론(inference) 속도를 획기적으로 개선하는 중요한 기술적 진보로 평가될 수 있습니다.