[논문 리뷰] DenseMamba: State Space Models with Dense Hidden Connection for Efficient Large Language Models

Wei He, Kai Han외 5명|arXiv (Cornell University)|2024. 02. 26.

Topic Modeling인용 수 5

한 줄 요약

DenseMamba는 DenseSSM을 도입합니다. 이는 SSM(State Space Models)에서 층 간 정보 흐름을 개선하기 위한 밀집된 은닉 연결 메커니즘으로, 최소한의 매개변수 및 계산 오버헤드로 RetNet/Mamba 계열 아키텍처의 성능을 향상시킵니다.

ABSTRACT

Large language models (LLMs) face a daunting challenge due to the excessive computational and memory requirements of the commonly used Transformer architecture. While state space model (SSM) is a new type of foundational network architecture offering lower computational complexity, their performance has yet to fully rival that of Transformers. This paper introduces DenseSSM, a novel approach to enhance the flow of hidden information between layers in SSMs. By selectively integrating shallowlayer hidden states into deeper layers, DenseSSM retains fine-grained information crucial for the final output. Dense connections enhanced DenseSSM still maintains the training parallelizability and inference efficiency. The proposed method can be widely applicable to various SSM types like RetNet and Mamba. With similar model size, DenseSSM achieves significant improvements, exemplified by DenseRetNet outperforming the original RetNet with up to 5% accuracy improvement on public benchmarks. code is avalaible at https://github.com/WailordHe/DenseSSM

연구 동기 및 목표

• 대형 언어 모델용 Transformer의 대안으로서 효율적인 시퀀스 모델(SSM)의 사용을 촉진합니다.
• SSM의 계층 간 은닉 상태 저하와 그것이 정보 흐름에 미치는 영향을 식별합니다.
• DenseSSM을 제안하여 얕은 은닉 상태를 더 깊은 층에 촘촘하게 연결하되 학습/추론 효율성을 유지합니다.
• DenseSSM이 공개 벤치마크에서 RetNet 및 Mamba 변형의 성능을 향상시킵니다.
• 선택적 전이 모듈과 은닉 융합 모듈의 실용적이고 하드웨어 친화적인 구현 전략을 제공합니다.

제안 방법

• 선택적 전이 모듈 φ와 은닉 융합 모듈 Fuse를 통해 이전 m개의 층의 은닉 상태를 현재 층으로 촘촘히 연결하여 DenseSSM을 제안합니다.
• φ에서의 프로젝션+게이트 메커니즘을 사용하여 얕은 은닉 상태를 대상 서브스페이스로 투영하고 유용한 구성 요소를 선택합니다.
• 선택된 얕은 상태를 현재 층의 은닉 상태와 더하기 융합(Additive fusion)을 통해 융합합니다(h_t^L = h_t^L + sum φ(h_t^{l-i})).
• DenseSSM을 재발생(DenseRetNet) 및 컨볼루션/GAU와 유사한 블록 변형에도 확장하여 병렬 가능한 학습 및 효율적인 자기 회귀 추론을 유지합니다.
• DenseSSM을 RetNet 및 Mamba 아키텍처에 적용하여 유사한 모델 크기이지만 향상된 성능을 갖는 DenseRetNet과 DenseMamba를 만듭니다.
• 56GB The Pile-derived 데이터, 15B 토큰, LLaMA 토크나이저를 사용하고 제로샷 및 4샷 능력을 다양한 QA/일상상식 작업에서 평가합니다.

실험 결과

연구 질문

• RQ1 SSM 계층 간의 밀집된 은닉 연결을 도입하면 전통적인 계층 한정 은닉 상태 전파에 비해 정보 흐름이 유지되거나 향상됩니까?
• RQ2 DenseSSM이 RetNet- 및 Mamba 기반 모델의 학습/추론 효율성을 저하시키지 않고 정확도와 perplexity를 개선할 수 있습니까?
• RQ3 선택적 전이 모듈과 은닉 융합 모듈이 성능 및 매개변수 효율성에 어떤 영향을 미칩니까?
• RQ4 제로샷 및 소수 샷(4-shot) 평가에서 상식 QA 및 추론 벤치마크에서 개선이 일관되게 나타나나요?
• RQ5 DenseSSM이 최소한의 아키텍처 변경으로 다양한 SSM 변형(RetNet, Mamba)에 적용 가능한가요?

주요 결과

• DenseSSM은 얕은 레이어의 은닉 정보를 더 깊은 층으로 밀집하게 연결하는 선택적 프로젝션 및 게이팅 메커니즘을 통해 은닉 정보 흐름을 향상시킵니다.
• DenseRetNet 및 DenseMamba가 공개 벤치마크에서 제로샷 및 소수샷 설정하에서 Baseline RetNet, Mamba, 및 Transformer 기반 모델보다 성능이 개선됩니다.
• DenseRetNet은 원래 RetNet 대비 공개 벤치마크에서 최대 5% 정확도 향상을 달성합니다(초록에 기재된 바와 같이).
• DenseSSM은 매개변수 증가가 거의 없이 학습 병렬성 및 자기 회귀 추론 효율성을 유지합니다.
• Ablation 연구에 따르면 선택적 전이를 위한 Identity 프로젝션과 MLP를 사용하고 additive fusion을 사용하는 것이 성능과 매개변수 수 간의 유리한 균형을 제공합니다.