Gemma 4에서 다중 토큰 예측 (MTP)은 매우 효율적인 추측 디코딩을 지원하는 데 사용되는 특정 아키텍처입니다. 추측 디코딩은 대규모 언어 모델에서 추론 속도를 높이는 기법입니다. 대규모 타겟 모델에만 의존하여 자동 회귀적으로 토큰을 생성하는 대신 (한 번에 하나의 토큰을 생성하며 각 새 토큰은 이전 토큰에 종속됨) 더 작고 빠른 '초안 모델'이 여러 토큰을 미리 예측합니다. 그러면 타겟 모델이 이러한 초안 토큰을 병렬로 검증합니다. 타겟 모델이 초안 토큰을 거부하더라도 해당 위치에 올바른 토큰을 생성하고 (이 단계가 낭비되지 않도록 함) 초안 모델은 새로운 올바른 토큰부터 예측을 재개합니다.
Gemma 4는 이 더 작고 빠른 초안 모델로 기본 모델을 확장하여 MTP를 구현합니다. 이 초안 모델은 타겟 모델과 입력 임베딩 테이블을 공유하고 마지막 레이어 활성화를 기반으로 직접 빌드되므로 독립적이지 않습니다. 이로 인해 표준 자동 회귀 생성과 동일한 품질을 보장하면서 디코딩 속도가 크게 향상되어 이러한 체크포인트는 지연 시간이 짧은 온디바이스 애플리케이션에 적합합니다.
추측 디코딩은 여러 토큰을 초안 작성하고 단일 전달 패스에서 검증하는 방식으로 작동합니다. 밀집 모델의 경우 모든 토큰에 동일한 가중치가 사용되므로 여러 초안 토큰을 검증하면 최소한의 오버헤드가 추가됩니다. Gemma 4 26B A4B와 같은 전문가 망 (MoE) 모델은 다르게 작동합니다. 각 토큰은 서로 다른 전문가를 활성화할 수 있으므로 초안 토큰을 검증하려면 메모리에서 추가 전문가 가중치를 로드해야 할 수 있으며, 이는 초안 작성으로 인한 이점을 상쇄합니다. 배치 크기가 클수록 일반적으로 시퀀스 전반에서 활성화된 전문가가 더 많이 겹치므로 로드된 가중치의 재사용이 개선됩니다. 배치 크기가 1인 경우 이 겹침이 제한되므로 26B A4B 초안 작성기가 병렬 처리가 잘 되지 않는 하드웨어 플랫폼에서 속도 향상을 제공하지 않을 수 있습니다.
MTP 개선사항
Gemma 4는 초안 토큰의 품질과 효율성을 개선하기 위해 표준 추측 디코딩 파이프라인에 몇 가지 개선사항을 도입합니다.
- 공유 입력 임베딩: 초안 모델은 입력 임베딩 테이블을 타겟 모델과 공유합니다.
- 타겟 활성화: 초안 모델은 타겟 모델의 마지막 레이어에서 활성화를 사용하고, 이를 토큰 임베딩과 연결한 후 초안 작성기 모델의 측정기준으로 다운 프로젝션합니다.
- 효율적인 임베더: 전체 어휘에서 예측하는 비용이 많이 드는 작업을 방지하기 위해 모델은 유사한 토큰을 클러스터로 그룹화합니다. 먼저 가장 가능성이 높은 클러스터를 식별한 다음 최종 계산을 선택한 클러스터 내의 토큰으로만 제한합니다 (E2B 및 E4B만 해당).