Tailored Pruning Strategies for Mamba Models

Abstract

Large language models (LLMs) have grown significantly in size alongside their performance, making model compression an essential research focus. While transformers, the backbone of most LLMs, have been extensively studied for pruning, Mamba, an alternative architecture, also requires effective pruning strategies. However, existing pruning methods for LLMs often fail to perform well on Mamba. In this paper, I propose an N:M pruning approach tailored to Mamba by leveraging the characteristics of its underlying SSM (Structured State Space Model). Additionally, I adapt existing LLM pruning techniques to better suit Mamba's architecture. Through an in-depth analysis of Mamba's hidden states, I achieved higher sparsity levels without compromising accuracy. Notably, even at over 50% sparsity, the accuracy drop was limited to just 8.47%.;대규모 언어 모델(LLMs)인 GPT-3과 같은 모델들은 대규모의 학습 능력과 뛰어난 성능으로 다양한 언어 작업 들에서 활용되고 있다. 그러나 이러한 효과는 수백억 개의 매개변수를 포함하는 엄청난 크기로 인해 발생하는 높은 비용을 요구 한다. 이는 상당한 메모리 사용과 계산 오버헤드를 초래하여 많은 에너지와 하드웨어 자원이 필요하며, 특히 자원이 제한된 환경에서 지속 가능성과 접근성에 대한 우려를 야기한다. LLMs에서 주로 사용되는 트랜스포머 모델은 입력 크기와 함께 계산이 제곱적으로 증가하여 매개변수 수가 증가함에 따라 비효율성에 특히 취약하다. 트랜스포머의 계산 문제를 해결하기 위해 연구자들은 딥러닝에 상태 공간 모델(SSMs)을 통합하는 것을 탐구했다. 상태 공간 모델들의 계산적 이점은 시퀀스 길이에 따라 선형적으로 확장되고, 메모리 사용이 줄어들며, 계산량을 줄여 트랜스포머 기반 모델보다 월등한 자원 효율성을 달성한다. 상태 공간 모델들은 처음에는 크게 주목받지 못했으나 HIPPO 행렬 초기화와 같은 발전을 거치면서 언어 작업에서 뛰어난 성능을 보이는 모델들이 등장했다. Mamba 모델은 성공적인 성능을 보이는 한 대표적인 모델로, 음성 인식 및 텍스트 처리를 포함한 다양한 언어 작업에서 뛰어난 성능을 보이면서 점점 더 중요하게 평가받고 있다. 상태 공간 모델의 장기 기억을 효율적으로 관리할 수 있는 능력 덕분에 복잡한 언어 애플리케이션에서 특히 효과적으로 사용될 수 있다. Mamba는 Llama와 같은 트랜스포머 기반 모델보다 작지만 여전히 수십억 개의 매개변수를 가지고 있다. Mamba는 종종 하이브리드 모델로 사용되며, 이는 전체 매개변수 수를 더욱 증가시킨다. 따라서 Mamba의 크기를 줄이기 위해 프루닝 시도들이 이루어졌다. 기존 Mamba 프루닝 연구는 주로 하이브리드 모델의 비 Mamba 구성 요소를 프루닝 하거나 비전 작업에서 Mamba 프루닝 기술을 적용하는 데 중점을 두었다. 그러나 Mamba의 대부분의 매개변수는 채널 확장과 압축을 처리하는 선형 레이어에 집중되어 있다. Mamba의 매개변수 수를 효과적으로 줄이려면 이러한 선형 레이어의 프루닝이 필수적이다. 선형 레이어가 매개변수 수를 지배하지만, LLM 작업에 필수적인 상태 공간 모델의 장기 기억 능력은 상태 공간 모델 블록에 의해 처리된다. 성능을 유지하기 위해서는 상태 공간 모델 블록의 특성을 이해하는 것이 필요하다. 상태 공간 모델 의 핵심 메커니즘은 현재의 입력 정보와 과거를 포함하고 있는 현재 시스템의 상태를 n차원 상태 벡터로 압축하는 것이다. Mamba에서는 각 차원마다 독립적인 상태 공간 모델이 구성되어 있으며, 이는 트랜스포머의 다중 헤드 구조와 유사하게 다양한 학습 방향을 가능하게 한다. 이러한 독립적인 SSM들은 각각 상태 정보를 n차원 벡터에 저장하며 임베딩 차원으로 쌓여 있다. 다양한 정보를 포착하기 위해 Mamba는 채널 확장, 축소 및 컨볼루션 작업을 포함한다. 그러나 상태 벡터의 크기는 압축된 정보의 양에 상태 벡터 크기에 제한되며, 상태 벡터 내의 높은 중복성을 초래한다. 다양한 그룹 크기에서 가중치의 유사성을 분석한 결과, 상태 벡터 크기 수준의 그룹 크기에서 약간 더 높은 유사성이 나타났으며, 델타 투영 레이어에서 더 높은 유사성이 관찰되었다. 이에 따라 N:M 프루닝을 적용하여 m을 상태 벡터 크기로 설정하고 유사성 데이터를 기반으로 델타 투영 레이어의 sparsity를 증가시켰다. 트랜스포머에서의 이전 연구에 따르면, 일반 활성화 값보다 100배 이상 큰 크기의 이상 활성화 값이 LLM 성능에 큰 영향을 미친다. Mamba가 유사한 행동을 보이는지 확인하기 위해, 상위 90%의 활성화 값의 분포를 그렸다. Llama와 비교 했을 때 유사한 패턴을 보였으며, 상태 공간 모델 블록을 통과한 후 활성화 값이 더 다이내믹하게 분산되는 것이 관찰되었다. 이 관찰은 Mamba에서도 활성화 값을 고려할 필요가 있음을 강조한다. 대규모 언어 모델(LLMs)의 효율적인 관리와 관련된 문제를 해결하기 위해, Wanda에서 사용한 활성화 값을 사용한 프루닝 방법을 Mamba 모델에 맞게 수정했다. 이 수정된 접근 방식은 계산 복잡도를 크게 줄이고 처리 속도를 향상시켰다. 이 접근 방식은 또한 레이어 활성화가 높은 유사성을 보일 때 극단적인 sparsity를 적용할 수 있게 하면서도 모델 성능이 저하되지 않도록 한다. 이 방법의 효과를 평가하기 위해 추가적인 훈련이나 후처리 조정 없이 모델의 성능을 평가하는 제로샷 테스트를 수행했다. 결과는 다음과 같았다.: 50%의 sparsity 수준에서 모델은 단 4%의 정확도 하락을 보였다. 놀랍게도, 78.85%의 높은 sparsity 수준에서도 비교 가능한 정확도를 유지할 수 있었다. 이는 이 Mamba에 맞는 프루닝 전략의 효율성을 입증한다.