Technology Innovation Institute(TII)가 순수 State Space Model 기반 Falcon Mamba 7B를 개발해 Attention 메커니즘 없이도 Transformer 수준의 성능 달성

Context

Transformer 기반 언어모델의 Attention 메커니즘은 시퀀스 길이에 따라 연산량과 메모리 비용이 선형으로 증가하는 근본적 한계를 갖고 있다. State Space Language Models(SSLMs)은 이 한계를 극복하려 했으나 기존 최고 성능 Transformer 모델보다 성능이 떨어졌다.

Technical Solution

• Mamba 아키텍처 기반으로 설계: 기존 Mamba에 RMS normalization 레이어 추가하여 대규모 학습 시 안정성 확보
• 약 5500GT 규모 데이터로 학습: RefinedWeb 데이터에 공개 소스의 기술 데이터와 코드 데이터 추가
• 단계별 학습 전략 적용: 대부분의 학습에는 상수 학습률 사용, 후반 단계에 짧은 학습률 감소 기간 추가 및 고품질 큐레이션 데이터 혼합
• Instruct 버전 개발: 기본 모델에 지시 추종 작업을 위한 추가 학습 수행
• 배포 최적화: 4비트 양자화 버전과 torch.compile을 통한 빠른 추론 지원

Impact

• 새로운 LLM 리더보드에서 평균 점수 15.04: 같은 카테고리 순수 SSM 모델(mamba-7b-rw 6.25)보다 2.4배, Transformer 모델들(Mistral-7B 14.50~15.28)과 경쟁 수준
• 기존 LLM 리더보드에서 평균 점수 64.09: Falcon2-11B(64.28)와 동일 수준이며 같은 크기 Transformer 모델들(Llama 3 62.28~62.62)을 상회
• 메모리 사용량: 단일 A10 24GB GPU에서 임의 길이 시퀀스 처리 가능
• 토큰 생성 지연: 컨텍스트 크기와 무관하게 일정한 시간으로 새로운 토큰 생성

Key Takeaway

State Space Model이 Attention 없이도 Transformer 수준의 성능을 달성할 수 있음을 실증했으며, 아키텍처 개선(RMS normalization 추가)과 신중한 학습 전략(단계별 학습률 조절, 고품질 데이터 혼합)이 대규모 모델 학습의 핵심 요소임을 보여준다.