Adaptive MLP Pruning for Large Vision Transformers

이 논문은 기존 MLP 가지치기 방법의 한계를 극복하기 위해 레이블 없는 정보 엔트로피 기준과 이진 탐색을 활용한 적응형 MLP 가지치기 (AMP) 를 제안하여, CLIP 및 DINOv2 와 같은 대규모 비전 트랜스포머 모델의 파라미터와 연산량을 약 40% 줄이면서도 성능 저하 없이 효율성을 극대화합니다.

핵심

🎒 거대한 배낭을 가볍게 만드는 방법

1. 문제: 너무 무거운 배낭 (대규모 비전 트랜스포머)
지금까지 개발된 최신 인공지능 모델들은 엄청난 성능을 보여주지만, 그 크기가 너무 커서 일반 컴퓨터나 스마트폰에서 실행하기 어렵습니다. 마치 등산객이 너무 무거운 배낭을 메고 있는 것과 같습니다. 이 배낭에는 수많은 물건 (데이터) 이 들어있지만, 실제로 필요한 것만 골라내면 훨씬 가볍고 빠르게 이동할 수 있습니다.

2. 분석: 배낭의 어디가 가장 무거운가?
연구자들은 이 거대한 AI 모델의 구조를 분석해 보니, 전체 무게의 80% 이상을 차지하는 것이 'MLP(다층 퍼셉트론)'라는 부품임을 발견했습니다.

• 비유: 배낭 속에 들어있는 '무거운 돌멩이'들이 바로 이 MLP 모듈들입니다. 이 돌멩이들을 잘라내면 배낭이 훨씬 가벼워집니다.

3. 해결책 1: "무엇을 버릴지"를 판단하는 새로운 나침반 (정보 엔트로피)
기존에는 어떤 돌멩이를 버릴지 판단할 때, 정답이 있는 문제 (예: "이 사진은 고양이인가?") 에만 집중했습니다. 하지만 이 방법은 AI 가 "고양이일 수도 있고, 강아지일 수도 있고, 토끼일 수도 있다"는 다양한 가능성을 고려하지 못해 중요한 정보를 잃을 수 있었습니다.

• 새로운 방법: 연구자들은 **'정보 엔트로피 (Information Entropy)'**라는 새로운 나침반을 만들었습니다.
• 비유: 기존 방법은 "정답이 '고양이'일 때만 점수를 매겼다"면, 새로운 방법은 **"고양이, 강아지, 토끼 등 모든 가능성이 어떻게 분포되어 있는지"**를 모두 살펴봅니다. 마치 모든 방향을 훑어보는 360 도 카메라처럼 AI 가 세상을 어떻게 바라보는지 전체적인 그림을 보고, 가장 덜 중요한 부분 (중복된 정보) 을 정확히 찾아냅니다.
• 장점: 이 방법은 정답 (라벨) 이 없어도 작동합니다. 그래서 공개된 데이터나 정답이 없는 모델 (예: DINOv2) 도 쉽게 다룰 수 있습니다.

4. 해결책 2: "얼마나 잘라낼지"를 자동으로 정하는 스마트 가위 (이진 탐색)
기존 방법들은 "무조건 50% 를 잘라내자"처럼 미리 정해진 비율로 자르곤 했습니다. 하지만 모델마다 불필요한 부분의 양이 다릅니다.

• 새로운 방법: 연구자들은 '이진 탐색 (Binary Search)' 알고리즘을 사용했습니다.
• 비유: 마치 금광에서 금을 캐는 것과 같습니다.
  1. 배낭을 반으로 잘라봅니다. (너무 많이 잘랐나요? 아니면 너무 적게 잘랐나요?)
  2. AI 의 성능 (정보 엔트로피) 이 너무 많이 떨어지지 않는 선에서, 가장 많이 잘라낼 수 있는 최적의 지점을 찾아냅니다.
  3. 이 과정을 반복하며 각 부품마다 "얼마나 redundantly(불필요하게) 많은지"에 맞춰 자동으로 잘라냅니다.
  • 결과: "무조건 50%"가 아니라, "이 부분은 30%, 저 부분은 60%"처럼 각 상황에 맞는 최적의 크기로 조정됩니다.

5. 마무리: 잃어버린 기억을 되찾는 과정 (지식 증류)
무언가를 잘라내면 원래의 성능이 떨어질 수 있습니다. 이때 **원래의 거대한 AI(선생님)**가 **잘라낸 작은 AI(학생)**를 가르쳐 줍니다.

• 비유: 큰 형님이 동생에게 "이걸 잘라내도 괜찮아, 내가 너를 도와줄게"라고 가르쳐 주는 것입니다. 작은 AI 는 원래 AI 의 생각 방식을 따라 배워서, 크기는 작아졌지만 똑똑함은 거의 잃지 않게 됩니다.

---

🚀 이 기술이 가져온 놀라운 결과

이 방법을 적용한 결과, 다음과 같은 기적이 일어났습니다:

1. 반으로 줄어든 크기: 모델의 파라미터 (두뇌 세포) 와 계산량 (FLOPs) 이 약 40% 감소했습니다. 배낭이 반으로 가벼워진 셈입니다.
2. 거의 손실 없는 성능: 크기가 줄었는데도, 성능은 원래 모델과 거의 똑같아졌습니다. (심지어 일부 경우엔 더 좋아지기도 했습니다!)
3. 다른 방법보다 훨씬 강력함: 기존에 있던 다른 가지치기 방법들보다 성능이 월등히 좋았습니다. 특히, 추가 학습 (파인튜닝) 을 하지 않아도 다른 방법들보다 훨씬 잘 작동했습니다.
4. 범용성: CLIP, DINOv2 등 다양한 최신 모델에 적용 가능했습니다.

💡 한 줄 요약

이 논문은 **"AI 의 가장 무거운 부분 (MLP) 을, 정답 없이도 정확하게 찾아내어 (정보 엔트로피), 각 부품에 맞춰 자동으로 최적의 크기로 잘라내고 (이진 탐색), 원래의 지식을 전수받아 (지식 증류) 가볍고 빠른 AI 를 만드는 방법"**을 제안했습니다.

이제 거대한 AI 도 스마트폰이나 개인용 컴퓨터에서 가볍게 돌아갈 수 있는 길이 열렸습니다!

기술 요약

논문 제목: Adaptive MLP Pruning for Large Vision Transformers (AMP)

저자: Chengchao Shen (중앙남부대학교)

1. 문제 제기 (Problem)

• 대규모 비전 트랜스포머 (Large Vision Transformers) 의 확장성 vs. 비용: 비전 트랜스포머 (ViT) 는 모델 용량이 커질수록 성능이 비약적으로 향상되지만, 방대한 파라미터 수로 인해 계산 비용과 메모리 요구량이 과도하여 실제 배포에 큰 장벽이 됩니다.
• MLP 모듈의 비효율성: 기존 트랜스포머 구조 분석 결과, 모델 파라미터의 대부분을 차지하는 것은 멀티헤드 셀프 어텐션 (Self-Attention) 이 아니라 멀티레이어 퍼셉트론 (MLP) 모듈임이 확인되었습니다. (예: EVA-CLIP-E 의 경우 MLP 가 전체 파라미터의 81.1% 차지).
• 기존 프루닝 (Pruning) 방법의 한계:
  • 기존 Taylor 기반 프루닝 방법들은 주로 One-hot Cross Entropy 손실 함수를 사용하여 중요도를 평가합니다. 이는 정답 레이블에 해당하는 예측 확률만 고려하고, 다른 클래스에 대한 잠재적 예측 정보는 무시하여 중요도 평가의 정확도 (Fidelity) 를 떨어뜨립니다.
  • 또한, 많은 최신 모델 (예: DINOv2, CLIP) 은 학습 시 사용된 손실 함수나 추가 모듈 (예: DINO head, 텍스트 인코더) 의 가중치가 공개되지 않아, 기존 방법론을 직접 적용하기 어렵습니다.
  • 대부분의 방법은 사전에 정의된 고정된 압축 비율 (Predefined compression ratio) 을 사용하므로, 각 MLP 모듈의 실제 중복성 (Redundancy) 에 맞춘 적응형 압축이 어렵습니다.

2. 제안 방법 (Methodology: AMP)

저자들은 대규모 비전 트랜스포머의 MLP 모듈을 적응형으로 프루닝하는 Adaptive MLP Pruning (AMP) 방법을 제안합니다. 주요 과정은 다음과 같습니다.

가. 레이블 없는 정보 엔트로피 기준 (Label-free Information Entropy Criterion)

• 문제 해결: 기존 One-hot Cross Entropy 의 한계를 극복하기 위해, 모델의 전체 예측 분포를 활용하는 **정보 엔트로피 (Information Entropy)**를 중요도 평가 기준으로 도입합니다.
• 구현:
  • 레이블이 없는 데이터셋 (Unlabeled dataset) 에서 미니배치 내 이미지 간의 유사성 행렬 (Inter-instance similarity matrix) 을 계산합니다.
  • 이 유사성 행렬에 Softmax 를 적용하여 예측 확률 분포를 유도하고, 이를 기반으로 정보 엔트로피를 계산합니다.
  • 장점: 원본 모델의 손실 함수나 추가 모듈 (예: DINO head, 텍스트 인코더) 없이도 중요도를 정밀하게 평가할 수 있어, 가중치가 공개되지 않은 모델에도 적용 가능합니다.

나. 적응형 MLP 프루닝 (Adaptive MLP Pruning via Binary Search)

• 중요도 순위 매기기: 위 기준을 통해 MLP 의 숨겨진 뉴런 (Hidden neurons) 중요도를 계산하고 순위를 매깁니다.
• 이진 탐색 (Binary Search) 알고리즘:
  • 고정된 압축 비율 대신, 각 MLP 모듈의 중복성에 따라 최적의 숨겨진 뉴런 수를 동적으로 결정합니다.
  • 프루닝 후 모델의 정보 엔트로피 변화량 ($\Delta E$) 이 임계값을 초과하지 않는 범위 내에서 이진 탐색을 수행하여 최적의 숨겨진 크기를 찾습니다.
  • 이를 통해 각 레이어마다 다른 정도의 압축이 가능해지며, 성능 저하를 최소화하면서 최대의 압축을 달성합니다.

다. 지식 증류 (Knowledge Distillation)

• 성능 회복: 프루닝된 모델 (Student) 이 원본 모델 (Teacher) 의 성능을 회복하도록 지식 증류를 수행합니다.
• 구조적 이점: MLP 의 숨겨진 층만 제거되고 출력 차원은 원본과 동일하므로, 추가 정렬 모듈 없이도 클래스 토큰과 패치 토큰의 출력에 대해 MSE 손실 함수를 사용하여 효율적으로 증류가 가능합니다.

3. 주요 기여 (Key Contributions)

1. 정확한 중요도 평가 기준: 레이블이 없는 정보 엔트로피 기준을 도입하여, 기존 Cross Entropy 보다 정밀한 뉴런 중요도 점수를 제공하며, 공개되지 않은 가중치나 손실 함수를 가진 모델의 압축을 가능하게 함.
2. 적응형 프루닝 전략: 사전 정의된 압축 비율을 제거하고, 이진 탐색을 통해 각 모듈의 중복성에 맞춰 적응적으로 뉴런을 제거하는 방법을 제안.
3. 근사 손실 없는 압축: ImageNet-1K 에서 증류만 수행한 경우에도 약 40% 의 파라미터 및 FLOPs 감소와 동시에 거의 성능 저하 없는 가속화를 달성. 미세 조정 (Finetuning) 없이도 기존 프루닝 방법들보다 월등히 우수한 성능을 보임.

4. 실험 결과 (Results)

• 대상 모델: OpenCLIP-g, OpenCLIP-G, EVA-CLIP-E, EVA-CLIP-8B, DINOv2-g 등 최신 대규모 비전 트랜스포머 모델.
• 성능 지표:
  • 파라미터 및 FLOPs 감소: 모든 모델에서 약 40% 감소.
  • 추론 속도: 약 1.5 배 가속화.
  • Zero-shot 분류 (ImageNet variants, ObjectNet):
    • 미세 조정 (Distillation) 후 원본 모델과 동등하거나 일부 경우 (OpenCLIP-g, EVA-CLIP-E) 오히려 약간 더 높은 정확도를 기록.
    • 미세 조정 없이 프루닝만 수행한 경우에도 기존 방법 (Random, $\ell_2$, Taylor, SAViT 등) 보다 **매우 큰 마진 (예: OpenCLIP-g 에서 42.7% 성능 향상)**으로 우세함.
  • 검색 작업 (Retrieval): Flickr30K 및 COCO 데이터셋에서 원본 모델과 비교 가능한 성능 유지.
  • kNN 평가: DINOv2-g 와 같은 순수 비전 트랜스포머에서도 원본 성능을 회복하며 약 54.4% 의 파라미터만 사용.

5. 의의 및 결론 (Significance)

• 실용성: 대규모 비전 트랜스포머의 배포 장벽을 낮추는 효과적인 솔루션을 제시하며, 특히 가중치가 공개되지 않거나 레이블이 없는 환경에서도 적용 가능한 범용성을 가짐.
• 효율성: 고정된 압축 비율이 아닌 데이터와 모델 구조에 기반한 적응형 접근법을 통해, 불필요한 파라미터 제거와 성능 보존 사이의 최적 균형을 달성함.
• 미래 전망: MLP 프루닝의 성공을 바탕으로, 향후 멀티헤드 셀프 어텐션 모듈의 적응형 축소 및 대규모 언어 모델 (LLM) 가속화로 연구 범위를 확장할 계획.

이 논문은 대규모 비전 모델의 효율성을 극대화하면서도 성능을 유지하는 새로운 프루닝 패러다임을 제시했다는 점에서 중요한 의의를 가집니다.