최근 arXiv에 발표된 연구에 따르면, 대규모 언어 및 비전 트랜스포머(Transformer) 모델의 효율성을 높이기 위한 새로운 경량화 기법인 '손상 인지 밴딧 프루닝(Damage-Aware Bandit Pruning)'이 제안되었습니다. 이 기술은 모델의 핵심 구성 요소인 어텐션 헤드(attention head)나 MLP 채널 그룹(MLP channel group)과 같은 구조적 단위를 제거(pruning)할 때 발생하는 성능 저하를 최소화하면서, 제한된 평가 예산 내에서 최적의 가지치기 전략을 찾아내는 것을 목표로 합니다.
연구팀은 구조적 단위 선택 문제를 '손상 인지 다중 팔 밴딧(damage-aware multi-armed bandit)' 문제로 재정의했습니다. 이는 각 단위를 제거했을 때 발생하는 손상(masked loss - base loss)을 측정하고, 이 손상을 최소화하는 방향으로 제거할 단위를 순차적으로 선택하는 방식입니다. UCB(Upper Confidence Bound) 스타일 정책이나 톰슨 샘플링(Thompson Sampling)과 같은 밴딧 알고리즘을 활용하여, 어떤 단위를 제거할지 결정하며, 이 과정에서 배치(batch) 간 변동성을 줄여 정확도를 높였습니다. GPT-2, OPT, Pythia, Qwen2.5, SmolLM2 등 언어 모델과 ViT-B/16, DeiT-Tiny, Swin-Tiny 등 비전 모델을 대상으로 한 실험에서, 이 밴딧 프루닝 방식은 무작위, 크기 기반, 정적 살리언시(static-saliency), 예산 기반 그리디(budgeted-greedy) 등 기존 가지치기 방법론 대비 성능 저하를 효과적으로 줄이는 것으로 나타났습니다.
이 연구는 대규모 트랜스포머 모델을 실제 환경에 배포할 때 중요한 과제인 모델 경량화에 새로운 방향을 제시합니다. 모델의 크기를 줄이면 추론(inference) 속도가 빨라지고, 메모리 사용량이 감소하며, 운영 비용 또한 절감됩니다. 특히, 이 방법은 물리적인 압축이나 속도 향상보다는 '유효한 구조적 억제(effective structural suppression)'에 초점을 맞춰, 모델의 핵심 기능을 유지하면서 불필요한 부분을 제거하는 데 강점을 보입니다. 이는 자원이 제한적인 환경에서도 고성능 AI 모델을 효율적으로 활용할 수 있게 하여, AI 서비스의 접근성을 높이고 다양한 산업 분야에서의 적용을 가속화할 잠재력을 가지고 있습니다.