[논문 리뷰] Falcon Mamba: The First Competitive Attention-free 7B Language Model

Jingwei Zuo, Maksim Velikanov외 5명|arXiv (Cornell University)|2024. 10. 07.

DNA and Biological ComputingBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

Falcon Mamba 7B는 5.8조 토큰으로 훈련된 Mamba 아키텍처 기반의 처음으로 경쟁력 있는, 완전히 어텐션을 제거한 7B 파라미터 언어 모델이다. 이 모델은 Mistral 7B, Llama3.1 8B, Falcon2 11B와 같은 선도적인 오픈소스 모델들을 능가하며, 장기적 맥락 처리 시 선형 추론 비용과 낮은 메모리 사용을 유지한다. 이는 순수 Mamba 모델이 스케일링 시 기존 트랜스포머 기반 최고 수준의 성능을 달성하거나 초월할 수 있음을 입증한다.

ABSTRACT

In this technical report, we present Falcon Mamba 7B, a new base large language model based on the novel Mamba architecture. Falcon Mamba 7B is trained on 5.8 trillion tokens with carefully selected data mixtures. As a pure Mamba-based model, Falcon Mamba 7B surpasses leading open-weight models based on Transformers, such as Mistral 7B, Llama3.1 8B, and Falcon2 11B. It is on par with Gemma 7B and outperforms models with different architecture designs, such as RecurrentGemma 9B and RWKV-v6 Finch 7B/14B. Currently, Falcon Mamba 7B is the best-performing Mamba model in the literature at this scale, surpassing both existing Mamba and hybrid Mamba-Transformer models, according to the Open LLM Leaderboard. Due to its architecture, Falcon Mamba 7B is significantly faster at inference and requires substantially less memory for long sequence generation. Despite recent studies suggesting that hybrid Mamba-Transformer models outperform pure architecture designs, we demonstrate that even the pure Mamba design can achieve similar, or even superior results compared to the Transformer and hybrid designs. We make the weights of our implementation of Falcon Mamba 7B publicly available on https://huggingface.co/tiiuae/falcon-mamba-7b, under a permissive license.

연구 동기 및 목표

• 순수 Mamba 기반 언어 모델이 스케일링 시 최고 수준의 성능을 달성하거나 최적화된 트랜스포머 모델과 비교해도 열등하지 않음을 입증하는 것.
• 어 attention을 완전히 제거한 아키텍처가 일반적인 언어 이해 및 생성 능력에서 하이브리드 Mamba-트랜스포머 모델과 동등하거나 슈퍼어리어할 수 있는지 조사하는 것.
• 일관된 메모리 사용과 높은 추론 효율성으로 장기 맥락 생성에서 순수 Mamba 모델의 실현 가능성과 장점을 검증하는 것.
• 효율적이고 확장 가능한 시퀀스 모델링 기술의 발전을 위해 연구를 촉진하기 위해, 완전히 오픈소스이며 허용성 있는 라이선스를 적용한 7B 모델을 제공하는 것.

제안 방법

• 모델는 64층, 4096차원의 은닉 차원, 16차원 상태 표현을 가진 상태공간모델(SSM)을 사용하는 순수 Mamba 아키텍처를 채택한다.
• 7B 규모에서의 유연성과 성능 향상을 위해 입력 및 출력 임bedding을 분리(untied)한다.
• 어텐션 메커니즘을 포함하지 않고, 정교하게 구성된 데이터 믹스처를 기반으로 5.8조 토큰으로 훈련한다.
• 패딩 토큰 앞뒤에서 은닉 상태를 0으로 초기화하고, 왼쪽 패딩을 적용하여 인과적 컨볼루션 중 간섭을 방지한다.
• 손실 급등 현상이 흔한 Mamba 훈련에서의 안정성을 높이기 위해 철저한 가중치 초기화와 학습률 민감도 감소 전략을 적용한다.
• Hugging Face와 llama.cpp를 통해 배포되며, GGUF 양자화 및 CPU 추론을 지원하여 광범위한 접근성을 확보한다.

Figure 1: Data mixtures across training stages

실험 결과

연구 질문

• RQ1 순수 Mamba 기반 언어 모델이 7B 파라미터에서 최고 수준의 오픈소스 트랜스포머 모델들과 경쟁 가능한 성능을 달성할 수 있는가?
• RQ2 하이브리드 어텐션 성분 없이도, 장기 맥락 길이에서 순수 어텐션 없는 아키텍처가 선형 추론 비용과 메모리 효율성을 유지할 수 있는가?
• RQ3 순수 Mamba 모델의 성능은 하이브리드 Mamba-트랜스포머 모델 및 RWKV나 Griffin과 같은 다른 SSM 기반 아키텍처와 비교해 어떻게 되는가?
• RQ4 Mamba 모델의 인-컨텍스트 학습 한계를 고려할 때, 데이터 품질과 스케일링이 트랜스포머 모델 대비 얼마나 보완할 수 있는가?
• RQ5 순수 Mamba 모델에서 안정적인 훈련과 높은 성능을 가능하게 하는 주요 훈련 및 아키텍처 기법은 무엇인가?

주요 결과

• Falcon Mamba 7B는 Open LLM 리더보드에서 Mistral 7B, Llama3.1 8B, Falcon2 11B를 모두 능가하며, 어텐션 메커니즘 없이도 경쟁 가능한 성능을 보였다.
• RecurrentGemma 9B와 RWKV-v6 Finch 7B/14B와 비교해도 성능이 유사하거나 뛰어나며, 다양한 아키텍처 설계에서 강력한 결과를 확인했다.
• 모델은 시퀀스 길이에 관계없이 일정한 메모리 사용과 처리량을 유지하여, 효율적인 장기 맥락 생성이 가능하다.
• 중간 길이 맥락(8,000토큰)으로 훈련되었음에도 불구하고, 강력한 일반화 및 추론 능력을 보이며, 향후 장기 맥락 최적화 가능성을 시사한다.
• 특히 장기간의 시퀀스 처리에서 트랜스포머 모델 대비 뚜렷한 추론 속도 향상과 메모리 사용 감소를 보였다.
• 모델는 책임감 있는 사용 정책을 포함한 허용성 있는 Apache 2.0 기반 라이선스로 공개되었으며, Hugging Face와 llama.cpp 경유로 넓은 연구 및 구현 접근성을 제공한다.

Figure 2: We vary the context length of the prompt to determine the maximum sequence length that could be processed without encountering an out-of-memory (OOM) error. To ensure a fair comparison, all models were configured with a rescaled vocabulary size.