[논문 리뷰] RecurrentGemma: Moving Past Transformers for Efficient Open Language Models

Aleksandar Botev, Soham De외 60명|arXiv (Cornell University)|2024. 04. 11.

Topic Modeling인용 수 4

한 줄 요약

RecurrentGemma-2B는 글로벌 어텐션을 선형 재귀와 국소 어텐션으로 대체하여 고정 크기의 상태 표현을 달성하는 Griffin 아키텍처를 기반으로 한 오픈소스 언어 모델을 소개한다. 이는 장문의 시퀀스에서 더 빠른 추론과 낮은 메모리 사용을 가능하게 하며, 훈련에 사용된 토큰 수가 33% 적음에도 불구하고 Gemma-2B와 동등한 성능을 달성한다.

ABSTRACT

We introduce RecurrentGemma, a family of open language models which uses Google's novel Griffin architecture. Griffin combines linear recurrences with local attention to achieve excellent performance on language. It has a fixed-sized state, which reduces memory use and enables efficient inference on long sequences. We provide two sizes of models, containing 2B and 9B parameters, and provide pre-trained and instruction tuned variants for both. Our models achieve comparable performance to similarly-sized Gemma baselines despite being trained on fewer tokens.

연구 동기 및 목표

• 장문의 시퀀스에서 추론 속도와 메모리 효율성 면에서 트랜스포머를 능가하는 고도로 효율적인 오픈소스 언어 모델을 개발한다.
• 선형 재귀 기반 모델이 Gemma-2B와 같은 최신 트랜스포머 모델과 비교해 경쟁 가능한 성능을 달성할 수 있음을 입증한다.
• KV 캐시의 선형 증가를 방지하기 위해 고정 크기 상태로 대체함으로써 장기 맥락 생성을 가능하게 한다.
• 연구 및 자원 제약 환경에서의 활용을 위해 사전 훈련된 버전과 지시 테이닝 버전을 모두 공개한다.
• Gemma의 책임감 있는 AI 관행을 따르며 철저한 벤치마크와 인간 평가를 통해 모델의 안전성과 일치성(alignment)을 검증한다.

제안 방법

• 글로벌 어텐션 없이 선형 재귀(RG-LRU)와 국소 어텐션(윈도우 크기 2048)을 조합한 Griffin 아키텍처를 채택하여 시퀀스를 모델링한다.
• 입력 시퀀스를 압축하기 위해 고정 크기 상태 벡터를 사용하여 자동회귀 생성 중에 증가하는 KV 캐시가 필요 없도록 한다.
• 훈련 안정성을 높이기 위해 입력 임bedding에 대해 가속도 계수(sqrt(model width))를 적용하며, 이는 Gemma의 설계와 일관된다.
• TPU에서 효율적인 추론을 위해 특수 설계된 Pallas 커널을 사용하며, 참조용 PyTorch 구현도 함께 제공한다.
• Gemma-2B와 동일한 데이터셋에서 2T 토큰으로 훈련하며, 일반 혼합 후 고품질 데이터를 대상으로 하는 이중 단계 사전 훈련 프로세스를 사용한다.
• 지시어 테이닝과 새로운 RLHF 알고리즘을 적용하여 지시어 따르기 능력을 향상시키며, 제어 토큰을 포함한 정의된 대화 형식을 사용한다.

실험 결과

연구 질문

• RQ1 선형 재귀와 국소 어텐션을 기반으로 한 재귀 아키텍처가 표준 NLP 벤치마크에서 Gemma-2B와 같은 트랜스포머 모델과 경쟁 가능한 성능을 달성할 수 있는가?
• RQ2 KV 캐시를 고정 크기 상태로 대체함으로써 기존 트랜스포머에 비해 장문의 시퀀스에서 훨씬 더 빠르고 메모리 효율적인 추론이 가능한가?
• RQ3 더 적은 토큰 수(2T 대비 3T)로 훈련된 모델가 장기 훈련 제약 조건을 유지하면서도 성능을 동등하게 유지할 수 있는가?
• RQ4 Mistral 7B v0.2 Instruct와 비교했을 때 인간 평가에서 RecurrentGemma-2B는 어떤 성능을 보이는가?
• RQ5 Griffin 기반 아키텍처가 메모리 제약이 없는 임의의 길이의 생성을 지원할 수 있으며, 처리량 측면에서 어떻게 스케일링되는가?

주요 결과

• RecurrentGemma-2B는 학술 벤치마크 세트에서 평균 44.6%의 점수를 기록했으며, Gemma-2B의 45.0%와 근접한 성능를 보이며 훈련에 사용된 토큰 수가 33% 적음에도 불구하고 경쟁 가능한 성능를 입증했다.
• 인간 평가에서 RecurrentGemma-2B-IT는 1,000개의 지시어 따르기 프롬프트에서 Mistral 7B v0.2 Instruct와 비교해 43.7%의 승리 비율을 기록했으며, 강력한 일치성과 사용성(usable)을 보였다.
• TPUv5e 장치에서 모든 시퀀스 길이에 걸쳐 RecurrentGemma의 추론 처리량은 Gemma를 항상 초월했으며, 시퀀스 길이가 증가함에 따라 성능 저하가 없었다.
• RecurrentGemma는 자동회귀 샘플링 중에도 높은 처리량(6,000 토큰/초)을 유지했지만, Gemma는 긴 시퀀스에서 증가하는 KV 캐시로 인해 처리량이 크게 감소했다.
• 프롬프트 처리 속도는 두 모델 모두 유사하게 빠르며(~40,000 토큰/초) 이는 성능 우위가 자동회귀 생성에 국한됨을 확인한다.
• 모델의 고정 크기 상태 덕분에 메모리 제약이 없는 임의의 길이의 생성이 가능하며, 이는 트랜스포머가 메모리 기반 KV 캐시 증가로 인해 제한받는 것과 대비된다.