동결된 멀티 토큰 예측으로 Pixel에서 Gemini Nano 모델 가속하기
기존 Gemini Nano v3 모델을 다시 학습하지 않고도 Multi-Token Prediction(MTP)을 추가해 Pixel 기기에서 온디바이스 추론을 가속하는 방법을 소개한다. 별도의 드래프터 모델 대신 본 모델에 가벼운 MTP 헤드를 붙이고, 기존 KV 캐시를 공유하는 zero-copy 구조를 사용해 메모리와 지연 시간을 줄였다. 그 결과 Pixel 9·10의 일부 기능에서 토큰 생성 속도가 50% 이상 향상되고, 인스턴스당 최대 130MB의 메모리를 절약했다. ## 모바일에서 자동회귀 생성이 느린 이유 - 기존 언어 모델은 한 번에 하나의 토큰만 생성하므로, N개 토큰을 만들려면 대형 모델의 추론을 N번 수행해야 한다. - 모바일 기기는 서버보다 RAM과 전력 예산이 제한적이다. - 토큰을 순차적으로 생성하는 과정은 연산 자원을 충분히 활용하지 못하고 메모리 대역폭과 배터리를 많이 소모한다. - 알림 요약, 메시지 교정처럼 짧은 응답도 빠르게 처리해야 하므로 추론 효율이 사용자 경험에 직접 영향을 준다. ## 별도 드래프터의 한계와 MTP - speculative decoding은 다음과 같은 두 단계로 동작한다. - **Draft:** 작은 드래프터 모델이 여러 후보 토큰을 빠르게 생성한다. - **Verify:** 대형 모델이 후보들을 병렬로 검증하고, 일치하는 토큰만 수용한다. - 별도 드래프터 모델은 추가 파라미터와 RAM을 필요로 한다. - 본 모델이 이미 계산한 풍부한 의미 정보를 활용하지 못하고, 텍스트 이력만으로 후보를 예측한다. - MTP는 별도 언어 모델 대신 본 모델의 마지막 층에 경량 Transformer 기반 MTP 헤드를 추가한다. - MTP 헤드는 본 모델의 hidden state를 이용해 여러 미래 토큰을 예측하고, 본 모델은 이를 병렬 검증한다. - 이러한 구조를 글에서는 본 모델의 깊은 지점에서 빠져나와 토큰을 예측하는 **“late exit” 전략**으로 설명한다. ## 동결된 백본에 MTP 헤드 추가 - 이미 배포된 Gemini Nano v3의 가중치는 그대로 동결한다. - 새로 학습하는 부분은 미래 토큰 예측을 담당하는 MTP 헤드뿐이다. - 기존 모델 전체를 다시 사전 학습하거나 별도 드래프터를 작업별로 미세 조정할 필요가 없다. - 백본을 동결하므로 기본 모델의 성능과 안전 정렬을 변경하지 않는다. - 잘못된 후보 토큰은 검증 단계에서 폐기되기 때문에 최종 출력은 기존 대형 모델과 비트 단위로 동일하다. - 따라서 기존 모델과의 하위 호환성을 유지하면서 추론 효율만 개선할 수 있다. ## KV 캐시를 공유하는 zero-copy 구조 - 일반적인 별도 드래프터는 자체 KV 캐시를 생성하고 유지해야 하므로 메모리를 중복 사용한다. - 제안된 MTP 헤드는 본 모델의 KV 캐시에 직접 cross-attention으로 접근한다. - 별도 프롬프트 처리(prefill)나 독립적인 과거 문맥 저장이 필요하지 않다. - 주요 효과는 다음과 같다. - 드래프터가 프롬프트를 다시 처리하지 않아 prefill 지연 감소 - 드래프터 전용 임베딩 테이블과 attention 구조 제거 - 애플리케이션별 튜닝 파라미터 및 중복 KV 캐시 절감 - standalone 드래프터와 비교해 인스턴스당 최대 130MB의 메모리를 절약했다. ## 풍부한 표현이 만드는 예측 정확도 향상 - MTP 헤드는 대형 백본이 이미 계산한 최종 hidden state를 활용하므로 별도 드래프터보다 정확한 후보를 생성한다. - Pixel 9에서 비슷한 규모의 standalone 드래프터보다 작업에 따라 50% 이상의 속도 향상을 보였다. - 복잡한 지시를 따르는 요약·재작성 작업에서 MTP가 크게 우수했다. - 스마트 답장처럼 문장 구조가 예측 가능한 작업에서는 본 모델의 구문 패턴을 잘 학습했다. - 이런 작업에서는 토큰 수용률이 최대 55% 향상됐다. ## Pixel 실사용 결과 - Gemini Nano MTP는 Pixel 9 및 Pixel 10 시리즈에 배포됐다. - 검증과 드래프팅 사이의 의존성을 처리하도록 온디바이스 추론 스택도 함께 재설계했다. - AI 알림 요약과 Proofread 같은 실제 작업에서 추론 한 번당 평균 약 2개의 추가 토큰을 정확히 예측했다. - 수용되는 토큰이 많아지면서 전체 검증 횟수가 감소했다. - 무거운 프로세서를 깨우는 횟수도 줄어들어 응답 시간이 짧아지고 배터리 사용량이 감소했다. ## 실용적인 결론 이미 배포된 온디바이스 모델을 유지해야 하는 경우, 별도 드래프터를 추가하는 것보다 동결된 백본에 MTP 헤드를 붙이고 KV 캐시를 공유하는 방식이 효율적이다. 특히 메모리와 전력이 제한된 모바일 환경에서 모델 출력의 동일성을 유지하면서 속도와 배터리 효율을 함께 개선할 수 있는 현실적인 최적화 전략이다.
원문 읽기(새 탭에서 열림)