손상 인지 다중 팔 밴딧을 활용한 트랜스포머 가지치기 기법
Damage-Aware Bandit Pruning for Vision and Language Transformers
arXiv트랜스포머 모델의 구조적 가지치기(pruning)를 위해 '손상 인지 다중 팔 밴딧' 문제를 활용한 새로운 최적화 방법을 제시했습니다.
- 마스킹 손실을 이용해 안정성을 높였으며, 기존 방식 대비 성능 저하를 줄이는 효과가 다양한 언어 및 비전 모델에서 확인되었습니다.
- 대규모 AI 모델(LLM)을 실제 환경에 배포할 때, 크기는 줄이면서도 핵심 성능 저하를 최소화하는 효율적인 경량화 기술로 활용될 수 있습니다.
- 성능 향상이 단순히 후보 단위 평가 예산 증가 때문만은 아니라는 등, 모델별 특성과 세부적인 조건들을 함께 고려해야 합니다.
본 연구는 모델의 구조적 후처리 가지치기(pruning)를 위해 손상 인지 다중 팔 밴딧 문제를 제안합니다. 이 방법은 고정된 후보 단위 평가 예산 하에서 작동하며, 언어 및 비전 트랜스포머에 대한 구조적 유닛 선택을 목표로 합니다. 특히 배치 간 변동성을 줄이기 위해 마스크 처리된 손실(masked loss)과 기본 손실(base loss)의 차이인 '페어드 손상(paired damage)'을 계산하여 안정적인 보상을 도출합니다.
제안된 방법론은 UCB 스타일 정책이나 분수 베타 톰슨 샘플링 같은 방식으로 진행되며, 선택된 유닛들은 원래의 밀집 체크포인트에서 기능적으로 제거됩니다. 실험은 WikiText-2, LAMBADA, Imagenette 등 다양한 데이터셋을 사용했으며, GPT-2, OPT, Pythia, Qwen2.5와 같은 언어 모델과 ViT-B/16, Swin-Tiny 등의 비전 모델에 적용되었습니다.
실험 결과, 밴딧 기반 방법은 일반적으로 예산이 제한된 탐욕적 선택(budgeted greedy) 방식 대비 성능 저하를 줄이는 경향을 보였습니다. 특히 언어 모델 비교에서 통계적으로 유의미한 개선이 관찰되었으며, ViT-B/16 및 Swin-Tiny의 경우 얻은 성능 향상이 단순히 더 큰 후보 단위 평가 예산 때문만은 아니라는 점도 확인되었습니다.
용어 풀이
- 트랜스포머
- 오늘날 대부분의 언어 모델이 쓰는 신경망 구조. 문장 속 단어들의 관계를 한꺼번에 계산해요.