HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

이 논문은 비전 트랜스포머의 엣지 디바이스 배포 한계를 해결하기 위해, 단일 엔드투엔드 학습 단계에서 매크로 및 마이크로 수준의 확률적 게이트를 통해 최적의 서브네트워크를 자동으로 탐색하는 계층적 자동 가지치기 프레임워크 'HiAP'를 제안합니다.

핵심

🧠 비전 트랜스포머 (ViT) 의 '지능형 자동 정리' 시스템: HiAP

이 논문은 인공지능, 특히 **'비전 트랜스포머 (Vision Transformer, ViT)'**라는 최신 이미지 인식 모델을 더 가볍고 빠르게 만드는 새로운 방법인 HiAP를 소개합니다.

기존의 방법들이 너무 복잡하고 수동적이었던 반면, HiAP 는 "모델 스스로가 무엇이 불필요한지 찾아내어 정리하는" 방식을 사용합니다. 일상생활에 비유해서 쉽게 설명해 드릴게요.

---

1. 문제 상황: "너무 무거운 가방"

비전 트랜스포머는 사진을 보고 "이건 고양이야, 개야"라고 판단하는 아주 똑똑한 AI 입니다. 하지만 이 AI 는 너무 무겁습니다.

• 문제: 스마트폰이나 드론 같은 작은 기기 (에지 디바이스) 에 넣으려면 너무 무겁고, 전기도 많이 먹습니다.
• 기존 해결책 (수동 정리): 예전에는 전문가들이 "여기 있는 이 부분은 쓸모없으니 잘라내자"라고 직접 결정하거나, 여러 단계에 걸쳐 복잡한 계산으로 불필요한 부분을 찾아냈습니다. 마치 가방을 정리할 때, 전문가가 와서 "이건 버려, 이건 가져가"라고 일일이 지시하는 것과 같습니다.

2. HiAP 의 등장: "스스로 정리하는 지능형 가방"

HiAP 는 이 문제를 모델 스스로가 배우게 함으로써 해결합니다.

🏗️ 두 단계의 정리 시스템 (마이크로 & 매크로)

HiAP 는 정리 작업을 두 가지 크기로 나눕니다.

1. 매크로 (Macro) 정리: "방 전체를 없애기"

   • 비유: 아파트에서 아예 필요 없는 방 (거실, 주방 등) 을 통째로 없애는 것입니다.
   • 작동: AI 가 "이 주의력 (Attention) 헤드는 쓸모없네?"라고 판단하면, 그 헤드가 담당하는 방 전체를 아예 없애버립니다. 이렇게 하면 무거운 가구를 옮길 필요 자체가 사라져서 (메모리 부하 감소) 매우 효율적입니다.

2. 마이크로 (Micro) 정리: "방 안의 가구 줄이기"

   • 비유: 방은 남겨두되, 방 안에 있는 불필요한 의자나 책상 다리를 잘라내는 것입니다.
   • 작동: 방 (헤드) 은 유지하되, 그 안에서 계산하는 세부적인 숫자 (차원) 나 뉴런을 줄입니다. 이렇게 하면 계산량 (FLOPs) 을 줄일 수 있습니다.

핵심: HiAP 는 이 두 가지를 동시에 수행합니다. "방을 통째로 없애는 게 더 효율적인가, 아니면 방은 두고 가구만 줄이는 게 나을까?"를 스스로 판단합니다.

3. 어떻게 작동할까? "스스로 배우는 정리 과정"

• 자동화 (Auto-Pruning):
  전문가가 "이것은 50% 잘라내라"라고 지시하지 않습니다. 대신 AI 가 **"이게 없으면 성능이 얼마나 떨어질까?"**를 학습하면서 스스로 "아, 이 부분은 없어도 되네"라고 깨닫습니다.
• 한 번에 끝내기 (Single-Phase):
  기존 방식은 "1 단계: 정리할 부분 찾기" -> "2 단계: 다시 학습하기"로 나누어 복잡했습니다. 하지만 HiAP 는 정리와 학습을 한 번에 끝냅니다. 마치 집을 정리하면서 동시에 살아가는 것처럼, 정리하는 과정에서도 성능이 유지되도록 훈련됩니다.
• 안정성 (Feasibility Penalty):
  AI 가 너무 과하게 정리해서 "집이 무너질까 봐" (모델이 망가질까 봐) 걱정할 때, HiAP 는 **"최소한의 방과 가구는 반드시 남겨야 한다"**는 규칙을 자동으로 적용합니다. 그래서 AI 가 스스로 무너지지 않는 선에서 최적의 크기로 정리됩니다.

4. 결과: "가볍고 빠른 AI"

실험 결과 (ImageNet 데이터셋 등) 에서 HiAP 는 다음과 같은 성과를 냈습니다.

• 성능 유지: 모델을 30~50% 정도 가볍게 만들었는데, 정확도는 거의 떨어지지 않았습니다.
• 빠른 속도: 불필요한 부분을 물리적으로 잘라내어, 실제 기기에서 실행할 때 약 1.4 배 더 빨라졌습니다.
• 간편함: 복잡한 설정이나 추가 학습 없이, 한 번의 훈련으로 바로 쓸 수 있는 가벼운 모델을 만들어냅니다.

📝 한 줄 요약

HiAP 는 AI 모델에게 "너 스스로가 가장 효율적인 몸무게를 찾아봐"라고 말해주는 시스템입니다.
전문가의 손길 없이, 모델이 스스로 "어떤 방을 없애고, 어떤 가구를 줄일지" 결정하여, 무거운 AI 를 가볍고 빠른 스마트폰용 AI 로 변신시킵니다.

이 기술은 앞으로 우리가 일상에서 더 빠르고 똑똑한 AI 를 더 작은 기기에서도 사용할 수 있게 해줄 핵심 열쇠가 될 것입니다.

기술 요약

1. 문제 정의 (Problem Statement)

비전 트랜스포머 (Vision Transformers, ViT) 는 컴퓨터 비전 분야에서 뛰어난 성능을 보이지만, 높은 계산 비용과 메모리 대역폭 요구 사항으로 인해 엣지 장치 (Edge Devices) 에 배포하기 어렵습니다. 기존 구조적 가지치기 (Structured Pruning) 방법론들은 다음과 같은 한계를 가지고 있습니다:

• 단일 세분성 (Single Granularity) 의 한계:
  • 마이크로 (Micro) 가지치기: 어텐션 헤드 내부 차원이나 FFN 뉴런만 제거하는 방식은 FLOPs 는 줄이지만, 메모리 대역폭 (DRAM/SRAM 접근) 병목 현상을 해결하지 못합니다. 모든 레이어와 어텐션 맵을 로드해야 하므로 실제 지연 시간 (Latency) 이 개선되지 않을 수 있습니다.
  • 매크로 (Macro) 가지치기: 전체 어텐션 헤드나 FFN 블록을 제거하는 방식은 메모리 오버헤드를 줄일 수 있지만, 네트워크의 표현 능력 (Representation Capacity) 을 과도하게 잃어 성능 저하가 큽니다.
• 복잡한 파이프라인 의존성: 기존 방법들은 대부분 다단계 (Multi-stage) 파이프라인을 사용하며, 가지치기 후 사후 처리 (Post-hoc) 를 위한 임계값 설정이나 수동 중요도 휴리스틱 (Heuristics) 에 의존합니다. 이는 전문가의 개입을 필요로 하고 배포 파이프라인을 복잡하게 만듭니다.

2. 제안 방법론: HiAP (Hierarchical Auto-Pruning)

저자들은 HiAP라는 새로운 프레임워크를 제안합니다. 이는 단일 엔드 - 투 - 엔드 (End-to-End) 학습 단계에서 최적의 서브네트워크를 자동으로 발견하는 연속 완화 (Continuous Relaxation) 기반의 확률적 가지치기 프레임워크입니다.

핵심 메커니즘

1. 계층적 확률적 게이트 (Hierarchical Stochastic Gating):

   • 매크로 게이트 (Macro-gates): 전체 어텐션 헤드와 FFN 블록의 유무를 결정합니다. (메모리 대역폭 오버헤드 감소)
   • 마이크로 게이트 (Micro-gates): 활성화된 매크로 구조 내부의 어텐션 차원 (Intra-head dimensions) 과 FFN 뉴런을 선택적으로 제거합니다. (연산량 FLOPs 감소)
   • Gumbel-Sigmoid: 이산적인 가지치기 결정 (0 또는 1) 을 학습 가능한 연속 변수로 완화하여, 역전파 (Backpropagation) 를 통해 엔드 - 투 - 엔드 최적화가 가능하도록 합니다.

2. 차분 가능한 비용 모델링 (Differentiable Cost Modeling):

   • 학습 목표 함수에 예상 MACs (Multiply-Accumulate Operations) 를 직접 포함시킵니다.
   • 구조적 유휴 비용 (Structural Overhead) 고려: 빈 어텐션 헤드를 유지하는 것보다 헤드를 완전히 제거하는 것이 메모리 오버헤드 측면에서 유리함을 수식화하여, 네트워크가 자연스럽게 불필요한 헤드를 제거하도록 유도합니다.
   • 실현 가능성 페널티 (Feasibility Penalties): 네트워크가 모든 레이어를 제거하는 '레이어 붕괴 (Layer Collapse)' 현상을 방지하기 위해, 각 레이어와 구조에 최소 유지 비율 (Minimum Retention Quotas) 을 강제하는 페널티를 도입합니다.

3. 단일 단계 엔드 - 투 - 엔드 발견 (Single-Phase End-to-End Discovery):

   • 기존 방법과 달리 '검색 (Search)'과 '파인튜닝 (Fine-tuning)' 단계를 분리하지 않습니다.
   • 온도 어닐링 (Temperature Annealing): 학습 초기에는 Gumbel-Sigmoid 의 온도 ($\tau$) 를 높게 유지하여 게이트가 확률적으로 동작하게 하고, 학습이 진행됨에 따라 온도를 낮추어 게이트가 결정론적인 0/1 로 수렴하게 합니다. 이를 통해 가중치와 구조가 점진적으로 적응 (Co-adaptation) 하도록 하여, 최종적으로 추가 파인튜닝 없이 바로 배포 가능한 물리적 서브네트워크를 추출합니다.

3. 주요 기여 (Key Contributions)

1. 통합된 다중 세분성 프레임워크: 매크로 (블록/헤드) 와 마이크로 (뉴런/차원) 가지치기를 단일 차분 가능 프레임워크로 통합하여, 메모리 대역폭과 연산량 (FLOPs) 을 동시에 최적화합니다.
2. 수동 개입 없는 예산 인식 학습: 사전 정의된 희소성 목표나 수동 중요도 순위 (Ranking) 없이, 모델이 스스로 최적의 서브 아키텍처를 발견하고 경화 (Harden) 시킵니다.
3. 실제 배포 가능한 결과물: 학습 종료 후 확률적 게이트를 임계값 (0.5) 으로 하드닝하여, 희소 연산 엔진이 필요 없는 밀집 (Dense) 서브네트워크를 물리적으로 추출합니다.

4. 실험 결과 (Results)

• 데이터셋: ImageNet-1K (DeiT-Small) 및 CIFAR-10 (ViT-Tiny).
• 성능 (ImageNet-1K):
  • DeiT-Small (4.6G MACs) 을 3.1G MACs (약 33% 감소) 로 압축했을 때, Top-1 정확도는 **79.1%**를 기록했습니다.
  • 이는 복잡한 다단계 파이프라인을 사용하는 최신 방법들 (GOHSP, ViT-Slim 등) 과 경쟁력 있는 정확도 - 효율성 파레토 프론티어를 달성하면서도, 훨씬 단순화된 배포 프로세스를 제공합니다.
• 성능 (CIFAR-10):
  • 33% 압축 시, 균일 비율 (Uniform-Ratio) 가지치기 및 $\ell_1$-Norm 휴리스틱 기반 방법보다 0.93% 높은 정확도를 달성했습니다.
  • 단일 GPU 에서 추론 지연 시간 (Latency) 이 5.57ms 에서 3.86ms 로 감소하여 약 1.44 배의 속도 향상을 확인했습니다. 이는 희소 연산 엔진 없이도 실제 하드웨어 가속이 가능함을 의미합니다.
• 구조 분석:
  • 학습 초기에는 계산 비용이 큰 매크로 구조 (전체 헤드, 마지막 레이어의 FFN 블록 등) 를 먼저 제거하고, 이후 마이크로 차원을 세밀하게 조정하는 계층적 가지치기 전략을 자동으로 학습하는 것을 확인했습니다.

5. 의의 및 결론 (Significance)

HiAP 는 ViT 가지치기를 "수동으로 튜닝된 휴리스틱"에서 "예산 인식형 단일 학습 문제"로 전환했습니다.

• 실용성: 복잡한 다단계 검색 과정과 추가 파인튜닝 없이, 학습과 동시에 최적의 서브네트워크를 생성하여 엣지 디바이스 배포 장벽을 낮춥니다.
• 효율성: 메모리 대역폭 병목과 연산 병목을 동시에 해결하여, 실제 하드웨어에서의 추론 속도와 에너지 효율을 극대화합니다.
• 확장성: 이 프레임워크는 토큰 가지치기 (Token Pruning) 나 양자화 (Quantization) 등 다른 최적화 기법과 결합하여 더 넓은 적용 가능성을 가집니다.

요약하자면, HiAP 는 ViT 의 구조적 복잡성을 줄이면서도 성능을 유지하는 자동화되고 효율적인 가지치기 솔루션을 제시하여, 트랜스포머 모델의 엣지 배포를 현실적으로 가능하게 하는 중요한 진전을 이룬 연구입니다.