프로즌 멀티-토큰 예측으로 픽셀 기기에서 Gemini Nano 모델 가속화하기

제목: 프로즌 멀티-토큰 예측으로 픽셀 기기에서 Gemini Nano 모델 가속화하기

Gemini Nano나 Gemma 같은 온디바이스 모델 덕분에 강력한 대규모 언어 모델(LLM)을 주머니 속에 넣어 다니는 게 이제 현실이 됐어요. 이 기술은 알림을 바로 요약해주거나 중요한 문자 메시지를 교정해주는 등 일상생활에서 휴대폰에 필요한 기능을 제공해주는데, 이 모든 게 개인 데이터를 외부 서버로 보내지 않고도 가능하죠. 하지만 이런 기능이 사용자들에게 유용하려면 정말 효율적으로 작동해야 해요.

모바일 기기에서 이런 속도를 구현하는 건 정말 쉽지 않은 과제예요. 방대한 서버 환경과 달리 휴대폰은 에너지 예산과 RAM(메모리)에 엄격한 제한이 있거든요. 게다가 일반적인 언어 모델은 텍스트를 '자기회귀적'으로 생성하는데, 이게 무슨 말이냐면 한 번에 한 단어(혹은 토큰)씩 처리하고 출력한다는 뜻이에요. 이렇게 한 단계씩 진행되는 과정이 병목 현상을 일으키고, 휴대폰의 처리 능력을 제대로 활용하지 못하면서 메모리 대역폭에 부담을 줘서 결국 사용자 경험을 저하시키고 배터리를 빨리 소모하게 만들 수 있어요.

이런 병목 현상을 해결하기 위해, 구글 연구팀은 기존의 '프로즌' Gemini Nano v3 모델에 Multi-Token Prediction (MTP)을 소급 적용하는 새로운 아키텍처를 발표하고 있어요. EAGLE 프레임워크와 Confident Adaptive Language Modeling (CALM) 같은 기존 접근 방식들을 기반으로, 구글 연구팀은 특히 모바일 환경에서 효율성을 극대화할 수 있도록 새로운 아키텍처 구성 요소를 설계했어요. 구글의 최근 발표에서는 MTP를 사용한 Gemma 4 가속화와 개발자들을 위한 가용성에 대해 강조했었죠.

오늘 글에서는 엣지 컴퓨팅의 독특하고 극심한 제약 조건들을 다루고 있어요. 최근 픽셀 9 및 10 시리즈에 적용된 이 접근 방식은 즉각적인 속도 향상을 가져다줘요. 사용자 입장에서는 AI 알림 요약이나 교정 같은 기능이 훨씬 빠르게 텍스트를 생성하고 에너지 소비도 줄어든다는 의미예요. 개발자 입장에서는 새로운 작업마다 별도의 메모리 소모가 큰 드래프팅 모델을 파인튜닝할 필요 없이 고속 온디바이스 AI를 제공할 수 있어서 큰 걸림돌 하나가 사라진 거죠.

'늦은 종료' 전략

MTP는 투기적 디코딩의 진화 위에 구축된 기술이에요. 전통적인 설정에서는 N개의 토큰을 생성하려면 대규모 모델의 N번의 순방향 패스가 필요하죠. 투기적 디코딩은 이 과정을 두 부분으로 분리해요:

1. 초안 작성 (Draft): 더 작고 빠른 근사 모델(이른바 '드래프터')이 후보 토큰의 짧은 시퀀스(예: 3개 토큰)를 생성해요.
2. 검증 (Verify): 대규모 모델(이른바 '검증기')이 이 후보들을 병렬로 처리해요. 만약 후보들이 대규모 모델이 예측했을 내용과 일치하면 수락되고, 그렇지 않으면 첫 번째 불일치 지점으로 되돌아가요.

하지만 이 방식에는 몇 가지 비효율성이 있어요. 별도의 '독립형' 드래프터 모델(예: 128M 파라미터)을 실행하면 제한된 RAM을 두고 경쟁하게 돼요. 게다가 독립형 드래프터는 메인 모델의 풍부한 내부 상태를 '모른 채' 작동해서, 메인 모델이 이미 계산한 의미론적 맥락 없이 오직 텍스트 기록에만 의존해서 다음 토큰을 예측하죠. MTP는 독립형 아키텍처에서 통합 아키텍처로 전환함으로써 이러한 비효율성을 해결했어요. 별도의 소규모 언어 모델을 훈련해서 토큰 초안을 작성하는 대신, 메인 모델의 마지막 레이어에 경량 트랜스포머 헤드인 MTP 헤드를 추가하는 방식이에요.

이 아키텍처는 초안 작성을 위해 딥 엑싯(deep exit) 레이어를 사용하는데, 메인 모델 백본이 이미 수행한 작업을 활용해요. MTP 헤드는 메인 모델의 최종 고차원 활성화(은닉 상태)를 받아 미래 토큰 시퀀스를 자기회귀적으로 예측하는 데 사용해요.

프로즌 백본의 장점

MTP 헤드는 일반적으로 백본과 함께 사전 학습되지만(예: 구글 연구팀의 최근 Gemma 4 모델 릴리즈에서처럼), 이미 배포된 온디바이스 파운데이션 모델을 활용할 때는 이런 방식이 어려워요. 대신, 구글 연구팀의 연구는 드래프터 헤드를 사전 학습 파이프라인과 독립적으로 작동하도록 소급 적용하는 데 중점을 뒀어요.

구글 연구팀은 완전히 훈련된 Gemini Nano v3 모델을 가져와 가중치를 고정하고, 밀집 트랜스포머 스택(MTP 헤드)을 마지막 레이어에 부착했어요. 그리고 미래 토큰 예측 오류를 최소화하기 위해 오직 이 파라미터들만 훈련했죠. 백본이 고정되어 있기 때문에 MTP는 순전히 효율성 최적화에 불과하며, 기본 모델의 기능이나 안전성 정렬에 전혀 손상이 없다는 걸 보장해줘요.

잘못된 초안은 검증 중에 폐기되므로 최종 출력은 메인 모델과 비트 단위로 동일하게 유지돼요. 덕분에 완벽한 하위 호환성을 유지하면서 효율성 업데이트를 배포할 수 있었어요.

제로-카피 아키텍처

표준 MTP 구현은 메인 모델과 드래프터 간에 정적 파라미터(예: 임베딩 가중치)를 공유해서 훈련 효율성을 최적화하지만, 온디바이스 추론은 훨씬 더 엄격한 병목 현상, 즉 동적 메모리에 직면해요. 가중치를 공유하더라도 드래프터가 컨텍스트를 독립적으로 처리한다면 자체적인 키-값 (KV) 캐시를 생성하고 유지함으로써 메모리에 '이중 세금'을 부과하게 돼요. 모바일의 제한된 메모리를 고려할 때 이런 중복을 피하는 것이 매우 중요해요.

이 문제를 해결하기 위해, 구글 연구팀은 MTP 헤드가 메인 모델의 상태를 효과적으로 활용하는 제로-카피 아키텍처를 설계했어요. MTP 헤드는 자체 기록을 유지하는 대신, 메인 모델의 고정된 KV 캐시에 직접 교차 어텐션(cross-attend)하도록 설계되었죠. 이 덕분에 드래프터는 백본이 이미 계산한 '기억'과 컨텍스트를 중복 없이 조회할 수 있게 돼요.

이 설계는 두 가지 효율성 향상을 가져다줘요. 첫째, 드래프터의 사전 채우기(prefill) 레이턴시를 없애줘요. 기존 캐시를 활용함으로써 헤드가 프롬프트를 처리하는 데 추가 시간이 필요 없게 되는 거죠. 둘째, 런타임 메모리 사용량을 줄여줘요. 구글 연구팀은 드래프터 임베딩 룩업 테이블, 사전 채우기 점 어텐션 변형, 그리고 애플리케이션별 튜닝 파라미터들을 절약함으로써 독립형 드래프터 대비 인스턴스당 130MB의 절감 효과를 확인했어요.

더 풍부한 표현의 잠금 해제

구글 연구팀의 실험 결과, MTP 드래프터는 비교 가능한 파라미터 수의 '독립형 드래프터' 대비 작업에 따라 픽셀 9 기기에서 50% 이상의 일관된 속도 향상을 보여줬어요.

이러한 성능 격차는 MTP가 더 풍부한 표현에 접근할 수 있다는 점에서 비롯돼요. 메인 모델을 블랙박스로 취급하는 독립형 드래프터와 달리, MTP 헤드는 더 큰 백본이 이미 처리한 최종 활성화 값을 직접 활용하거든요:

• 명령어 수행 (Instruction following): 요약이나 복잡한 제약 조건이 있는 재작성 같은 작업에서 MTP는 독립형으로 파인튜닝된 드래프터를 훨씬 능가하는 성능을 보였어요.
• 예측 가능한 텍스트 구조: 스마트 답장처럼 구조적 예측 가능성이 높은 작업의 경우, MTP 헤드는 메인 모델의 문법 패턴을 효과적으로 학습해서 토큰 수락률을 최대 55%까지 향상시켰어요.

실제 영향

픽셀 9 및 10 기기에 MTP를 배포하기 위해, 구글 연구팀은 검증 단계와 초안 작성 단계 사이의 복잡한 종속성을 처리할 수 있도록 온디바이스 추론 스택을 재설계했어요.

그 결과는 아키텍처 선택의 타당성을 입증해줬어요. AI 알림 요약 및 교정과 같은 실제 작업 부하에서 MTP는 추론 패스당 평균 거의 두 개의 추가 토큰을 올바르게 예측해요. 게다가 검증 단계가 줄어든다는 것은 무거운 프로세서를 깨우는 시간이 줄어들어 에너지 소비를 감소시키고 배터리 수명을 향상시킨다는 의미도 있죠.

미래 방향

구글 연구팀은 MTP를 미래 픽셀 기기에 통합하는 것은 물론, 엄격한 모바일 제약 조건 하에서 초안 작성 레이턴시를 더욱 줄이고 동시 토큰 검증을 늘리기 위해 병렬 디코딩 및 보조 헤드가 없는 패러다임을 포함한 대체 아키텍처를 탐색할 계획이에요.

또한 언어 생성의 내재된 모호성을 더 효율적으로 처리하는 방법도 연구하고 있어요. 표준 투기적 디코딩은 단 하나의 최적의 미래 경로를 가정하지만, 구글 연구팀은 모델이 병렬로 분기 가능한 가능성을 탐색할 수 있는 기술을 개발하고 있어요. 이는 불확실한 맥락에서도 긴 시퀀스를 수락할 확률을 최대화하는 것을 목표로 해요. 나아가 구글 연구팀은 검증 관용성(verification leniency)도 연구하고 있는데, 특정 사용 사례에 대해 초안과 검증 간의 엄격한 토큰 일치 조건을 완화하여 엣지 기기에서 추가적인 효율성을 가져오려는 시도예요.

감사 말씀

이 작업은 Filippo Galgani, Omri Homburger, Pooja Consul, Matthew Markwell, 그리고 Vivek Kumar와 함께 온디바이스 LLM 효율성을 최적화하기 위한 구글 연구팀의 노력의 일환이에요. 특정 요소들은 Google DeepMind의 Gemini 팀인 Tal Schuster, Ziwei ji, Ivan Korotkov, Ganesh Jawahar의 개발을 기반으로 만들어졌어요. 또한 Nadav Bar, Utku Evci, Nir Shabat, Joe Zou, 그리고 Google Research, Google Deepmind, Platforms & Devices 팀의 검토와 소중한 피드백 및 지원에 진심으로 감사드립니다.