[논문 리뷰] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

통

통통한아이멤버

Jun 20, 20242y ago

**해석과 맞춤법은 아이멤버 업무도우미를 이용하였습니다.

안녕하세요 이번 논문 리뷰를 맞게 된 통통한아이멤버입니다.

제가 이번에 리뷰할 논문은 "Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling"입니다.

Mamba, State Space Model (SSM), with Sliding Window Attention (SWA)를 합친 SAMBA로 무한한 context 길이를 가지는 sequence를 modeling하는 것을 해결했습니다.

위의 이미지에서 볼 수 있듯이 최대 1M 토큰까지 예측 성능 향상, 64K 디코딩에서 최신 아키텍보다 빠른 디코딩 성능을 확인했습니다.

이후 어떠한 점이 해당 아키텍처가 성능과 속도를 모두 잡았는지 살펴보겠습니다.

SSM의 약점을 보완한 하이브리드 접근법

Attention 기반 모델이 가진 장기 의존성 문제를 해결하기 위해 SSM을 이용한 새로운 모델이 제안되었고, 이는 서형 계산 복잡성과 더 긴 Sequence를 이해할 수 있게 되었습니다.

하지만 SSM 기반 모델들은 Markovian nature 때문에 memory recall에 문제가 있으며 retrieval-related tasks 굉장한 한계를 노출했습니다.

이를 해결하기 위해 어텐션 메커니즘과 혼합하는 다양한 방식이 탐구되었지만 선형 시간 복잡성 때문에 context extrapolation를 수행하지 못합니다.

해당 논문은 SSM과 어텐션 기반 모델의 강점을 조화시키면서 선형 시간 복잡성으로 무제한 context extrapolation을 달성하는 간단한 신경 아키텍처인 SAMBA를 소개합니다.

방법론

아키텍처

Mamba, Sliding Window Attention(SWA), 다층 퍼셉트론(MLP)으로 구성된 하이브리드 전략을 구성합니다.

Mamba는 반복되는 시퀀스 구조를 포착하고, SWA는 메모리를 정밀하게 기억하며, MLP는 사실적 지식을 기억하는 역할을 합니다.

위의 그림에서 볼 수 있듯 Mamba-MLP-SWA-MLP로 구성된 SAMBA의 모습을 확인할 수 있습니다.

Mamba Layer

Sliding Window Attention (SWA) Layer

Mamba가 가지고 있는 non-Markovian dependencies in sequences 특성을 해결하기 위해 SWA를 설계했습니다.

SWA는 입력 시퀀스 위에서 윈도우 사이즈 w=2048의 크기로 슬라이딩하여 시퀀스 길이에 비례하는 복잡성을 유지합니다.

Multi-Layer Perceptron (MLP) Layer

SAMBA의 MLP layer nonlinear transformation(비선형 변환)과 recall of factual knowledge(사실적 지식 회상)의 주요 메커니즘을 제공합니다.

SAMBA는 Mamba와 SWA layer에 의해 포착된 다른 유형의 정보를 위해 별도의 MLP를 적용합니다.

Experiments and Results

비슷한 조건으로 다양한 LLMs(Llama-3 1.6B, Mistral 1.6B, Mamba 1.8B)와 같은 모델들과 다양한 NLP 평가 테스트를 진행한 결과입니다.

같은 조건 (< 2.0B)에서 제안한 SAMBA의 성능이 우수한 것을 확인할 수 있습니다.

Long-Context Understanding

synthetic passkey retrieval task에서 오른쪽의 Mistral 1.6B와 비교하여 왼쪽의 SAMBA 1.7B의 모델이 압도적으로 긴 길이에서 retrieval 성능이 좋은 것을 확인할 수 있습니다.

Conclusion

논문에서 제안한 SAMBA 아키텍쳐는 무제한의 컨텍스트를 처리하기 위해 고안된 간단하지만 강력한 하이브리드 신경망입니다.

SAMBA는 SSM의 기반의 다른 모델들을 크게 능가하는 성능을 보여주며, 긴 컨텍스트를 처리하는데 있어 Transformer 아키텍쳐와 비교하여 매우 높은 효율성을 보여주었습니다.

더 적은 파라미터와 하이브리드 방식으로 더 효율적이고 효과적인 언어 모델링이 가능함을 증명했습니다.

Subscribe to 'lotte-innovate-ai'

Subscribe to my site to be the first to receive notifications and emails about the latest updates, including new posts.
Join Slashpage and subscribe to 'lotte-innovate-ai'!