[논문 리뷰] Pruning as a Cooperative Game: Surrogate-Assisted Layer Contribution Estimation for Large Language Models

Xuan Ding, Pengyu Tong외 4명|arXiv (Cornell University)|2026. 02. 08.

Topic Modeling인용 수 0

한 줄 요약

이 논문은 Transformer 층들 사이의 협력 게임으로 LLM 가지치기를 다루고, 모듈화된 대리모델을 사용해 Stratified Monte Carlo 마스크를 통한 Shapley 유사 층 기여를 추정하며, 서로 다른 베이스라인에 비해 perplexity와 제로샷 정확도에서 일관되게 개선되면서 효율적인 가지치기를 가능하게 한다는 것을 보여준다.

ABSTRACT

While large language models (LLMs) demonstrate impressive performance across various tasks, their deployment in real-world scenarios is still constrained by high computational demands. Layer-wise pruning, a commonly employed strategy to mitigate inference costs, can partially address this challenge. However, existing approaches generally depend on static heuristic rules and fail to account for the interdependencies among layers, thereby limiting the effectiveness of the pruning process. To this end, this paper proposes a game-theoretic framework that formulates layer pruning as a cooperative game in which each layer acts as a player and model performance serves as the utility. As computing exact Shapley values is computationally infeasible for large language models (LLMs), we propose using a lightweight surrogate network to estimate layer-wise marginal contributions. This network can predict LLM performance for arbitrary layer combinations at a low computational cost. Additionally, we employ stratified Monte Carlo mask sampling to further reduce the cost of Sharpley value estimation. This approach captures inter-layer dependencies and dynamically identifies critical layers for pruning. Extensive experiments demonstrate the consistent superiority of our method in terms of perplexity and zero-shot accuracy, achieving more efficient and effective layer-wise pruning for large language models.

연구 동기 및 목표

• 인터-레이어 의존성을 반영하여 정적 휴리스틱이 아닌 가지치기를 동기화하고자 한다.
• 대리 네트워크를 이용한 층 기여를 추정하기 위한 확장 가능한 2단계 근사를 제안한다.
• 레이어 간 상호 작용을 보존하여 효율적인 가지치를 위한 핵심 층을 식별한다.
• 여러 모델과 태스크에서 언어 모델링 성능 및 제로샷 능력의 향상을 입증한다.
• 양자화와 Transformer 아키텍처를 넘어선 적용 가능성과 양자화와의 호환성을 보인다.

제안 방법

• 각 Transformer 층을 한 명의 플레이어로 보고 모델 성능을 효용으로 간주하는 협력 게임으로 가지치기를 정식화한다.
• 제한된 해밍 가중치를 가진 다양한 가지치기 마스크를 생성하기 위해 층화된 Monte Carlo 마스크 샘플링을 도입한다.
• 미지의 마스크에 대한 성능 하락을 예측하고 Shapley 값을 근사하기 위해 경량 대리 신경망을 학습한다.
• 대리 모델과 축적된 마스킹 평가를 이용해 층의 한계 기여도를 추정하고 Shapley 값을 기준으로 층의 순위를 매긴다.
• 목표 압축 비율을 달성하기 위해 추정 기여도가 가장 낮은 층을 제거하는 방식으로 가지치기를 수행한다.
• 필요시 사후 학습 양자화와 결합하여 효율성을 더욱 향상시킨다.

실험 결과

연구 질문

• RQ1 층 간 의존성을 효과적으로 포착하여 LLM의 층 가지치기에 이를 가이드로 활용할 수 있는가?
• RQ2 대리 모델이 샤플리 값 기반 가지치기를 활용하기 위해 층 기여도의 정확하고 확장 가능한 추정치를 제공할 수 있는가?
• RQ3 게임 이론적이고 맥락 의존적인 가지치기 전략이 다양한 아키텍처와 태스크에서 정적 휴리스틱 및 기존 베이스라인을 능가하는가?

주요 결과

• 제안된 방법은 여러 모델 및 가지치기 수준에 걸쳐 깊이별(depth-wise) 및 넓이별(width-wise) 가지치기 베이스라인에 비해 perplexity를 낮추고 제로샷 정확도를 더 높게 달성한다.
• 특히 더 큰 Meta-LLaMA-3-8B 및 LLaMA-2-13B-hf 모델에서 공격적인 가지치기 하에서도 생성 품질을 유지한다.
• 인터-레이어 의존성을 고려한 가지치기 결정은 적대적 추론 태스크(ANLI)에서 강건성을 향상시킨다.
• 이 방법은 비-트랜스포머 아키텍처(RWKV, Mamba)로 일반화되며 사후 학습 양자화와의 호환성을 유지해 추가적인 메모리 및 처리량 이득을 가능하게 한다.
• 실험 결과 이 방법이 추가 메모리 부담 없이도 속도/처리량/지연 시간의 유리한 트레이드오프를 달성한다.