DivPrune: Diversity-based Visual Token Pruning for Large Multimodal Models

Large Multimodal Models (LMMs) have emerged as powerful models capable of understanding various data modalities, including text, images, and videos. LMMs encode both text and visual data into tokens that are then combined and processed by an integrated Large Language Model (LLM). Including visual tokens substantially increases the total token count, often by thousands. The increased input length for LLM significantly raises the complexity of inference, resulting in high latency in LMMs. To address this issue, token pruning methods, which remove part of the visual tokens, are proposed. The existing token pruning methods either require extensive calibration and fine-tuning or rely on suboptimal importance metrics which results in increased redundancy among the retained tokens. In this paper, we first formulate token pruning as Max-Min Diversity Problem (MMDP) where the goal is to select a subset such that the diversity among the selected tokens is maximized. Then, we solve the MMDP to obtain the selected subset and prune the rest. The proposed method, DivPrune, reduces redundancy and achieves the highest diversity of the selected tokens. By ensuring high diversity, the selected tokens better represent the original tokens, enabling effective performance even at high pruning ratios without requiring fine-tuning. Extensive experiments with various LMMs show that DivPrune achieves state-of-the-art accuracy over 16 image- and video-language datasets. Additionally, DivPrune reduces both the end-to-end latency and GPU memory usage for the tested models. The code is available here⋄.

한 줄 요약

DivPrune는 시각 토큰의 다양성 최대화를 통해 LMM의 추론 속도와 메모리 사용량을 줄이는 토큰 프루닝 방법이다.

핵심 기여도

DivPrune는 Max-Min Diversity Problem(MMDP) 기반으로 토큰 프루닝을 수행하여, 기존의 중요도 기반 프루닝 대비 토큰 간 중복을 최소화함.
16개의 이미지-언어 및 비디오-언어 데이터셋에서 기존 방법 대비 최대 +15.8% 성능 개선을 달성함.
모델 미세조정 없이도 80% 이상의 프루닝 비율에서도 성능 유지 가능함.
LLaVA 1.5-7B 모델 기준, 19.61% TFLOP 비율에서 Layer 0에서의 프루닝이 Layer 1, 2, 3 대비 정확도가 높음.

핵심 아이디어

기존 토큰 프루닝 방법은 주로 어텐션 점수나 중요도 기반으로 토큰을 선택하지만, 이는 토큰 간 중복을 유발하여 성능 저하를 초래한다. DivPrune은 이러한 문제를 해결하기 위해 토큰 간 거리 기반의 Max-Min Diversity Problem(MMDP)을 도입하여, 선택된 토큰 집합의 다양성을 최대화한다. 이는 토큰 간 최소 거리를 증가시키는 방식으로, 선택된 토큰이 원본 토큰을 더 잘 대표하도록 보장한다. 본 연구는 MMDP를 토큰 프루닝에 적용한 최초의 시도로, 기존의 Min-Max 전략이나 랜덤 프루닝과 비교해 15.8% 이상의 성능 향상을 보인다.

기술적 접근법

**MMDP 기반 프루닝**: 토큰 간 거리 기반으로 최대 다양성을 확보하는 Max-Min Diversity Problem을 토큰 선택에 적용.
**거리 측정**: Cosine, ℓ₁, ℓ₂ 거리 중 Cosine 유사도가 가장 효과적임을 실험적으로 검증.
**프루닝 타이밍**: LLM의 첫 번째 디코더 레이어(Layer 0)에서 프루닝을 수행하는 것이 Layer 1, 2, 3 대비 정확도가 높음.
**하이퍼파라미터**: TFLOP 비율 19.61% 기준, Layer 0에서의 프루닝이 가장 효과적임을 실험적으로 확인.
**모델 호환성**: DivPrune는 LLM 아키텍처와 비전 인코더와 상관없이 적용 가능하며, KV 캐싱과 같은 추론 최적화 기법과도 호환됨.

주요 결과

**성능**: 16개의 이미지-언어 및 비디오-언어 데이터셋에서 DivPrune이 기존 방법 대비 최대 +15.8%의 정확도 향상을 기록함.
**프루닝 비율**: 80% 이상의 프루닝 비율에서도 모델 미세조정 없이 성능 유지 가능.
**추론 속도**: LLaVA 1.5-7B 기준, DivPrune는 GPU 메모리 사용량과 end-to-end latency를 동시에 감소시킴.
**TFLOP 비율**: 19.61% TFLOP 비율에서 Layer 0에서의 프루닝이 Layer 1, 2, 3 대비 정확도가 높음.

의의 및 한계

DivPrune은 기존의 중요도 기반 프루닝 방법의 한계를 극복하고, 토큰 간 중복을 줄이며 성능을 유지하는 새로운 접근법을 제시한다. 특히, 모델 미세조정 없이도 높은 프루닝 비율에서도 효과적인 성능을 보이는 점에서 실용적 가치가 높다. 또한, 다양한 LLM 아키텍처와 비전 인코더에 적용 가능하며, 추론 최적화 기법과의 호환성도 뛰어나다. 그러나 본 연구는 특정 LMM 아키텍처에만 적용되었으며, 더 다양한 모델 구조에서의 일반화 가능성은 추가 연구가 필요하다.

실용적 활용

DivPrune는 저연산 비용으로 높은 성능을 유지하는 LMM의 추론 최적화에 적합하다. 특히, 실시간 이미지/비디오 분석이 필요한 IoT, 모바일, 클라우드 환경에서 유용하게 활용될 수 있다. 또한, 모델 미세조정 없이 즉시 적용 가능한 plug-and-play 속성 덕분에, 다양한 산업 분야에서 즉각적인 성능 향상과 리소스 절감을 기대할 수 있다.