기술 요약: Vision Transformer를 위한 깊이 프러닝(Depth Pruning)의 재고: 이질성 인지적 관점

1. 문제 정의

Vision Transformer(ViT)는 다양한 도메인에서 최첨단 성능을 달ach성했지만, 높은 계산 비용과 방대한 파라미터 수로 인해 긴 추론 지연 시간이라는 문제를 안고 있습니다. 구조적 프러닝(Structured pruning)은 효과적인 압축 전략이지만, 기존 방식들은 다음과 같은 중대한 한계에 직면해 있습니다:

• 너비 프러닝(Width Pruning): 레이스 내의 채널이나 어텐션 헤드를 줄이지만, 주류 하드웨어에서 실질적인 속도 향상을 얻는 데는 한계가 있는 경우가 많습니다.
• 깊이 프러닝(Depth Pruning): 전체 레이어를 제거하며, 훨씬 더 높은 속도 향상 잠재력을 제공합니다(그림 2a 참조). 그러나 공격적인 레이어 제거 시 정확도를 회복하는 것이 매우 어렵습니다.
• 핵심 병목 현상: 본 논문은 기존의 깊이 프러닝 방식이 실패하는 이유가 두 가지 간과된 문제 때문이라고 주장합니다:
  1. 이질성 무시(Neglect of Heterogeneity): ViT는 서로 다른 유형의 레이어(예: Attention 레이어 vs. Feed-Forward Network 내의 Activation 함수 레이어)를 포함하고 있습니다. 기존 방식은 이러한 레이어들을 균일하게 취급하거나, 레이어 유형 간의 상당한 그래디언트 규모 차이로 인해 실패하는 그래디언트 기반 메트릭에 의존합니다.
  2. 차원 불일치(Dimension Mismatch): Attention 레이어와 FFN 내의 Linear 레이어를 동시에 프러닝하면 구조적 불일치가 발생합니다. FFN은 확장-축소 메커니즘(예: $D \to 4D \to D$)을 사용하므로, Linear 레이어를 제거하면 후속 레이어에 필요한 텐서 형태의 연속성이 깨집니다(그림 3 참조).

2. 방법론: HetDPT

저자들은 이질성을 해결하고 차원 불일치를 방지하기 위해 설계된 2단계 프레임워크인 HetDPT(Heterogeneity-aware DePTh pruning)를 제안합니다.

핵심 통찰

이 방법은 레이어 이질성에 관한 두 가지 결정적인 관찰에 근거합니다:

1. 그래디언트 격차(Gradient Disparity): Attention 레이어와 Activation 레이어는 역전파 과정에서 매우 다른 그래디언트 크기를 보입니다. 표준 그래디언트 기반 중요도 메트릭을 사용하면 편향된 프러닝 결정을 내리게 되어, 특정 유형을 지속적으로 과하게 혹은 과소하게 프러닝하게 됩니다.
2. 회복 비대칭성(Recovery Asymmetry):
   • Attention 레이어를 프러닝하면 초기 정확도 저하는 완만하지만, 회복을 위해 광범위한 재학습이 필요합니다.
   • Activation 레이어를 프러닝하면 심각한 초기 성능 저하(종종 $\ge 90\%$ 정확도 손실)가 발생하지만, 파인튜닝 과정에서 매우 빠르게 회복됩니다.
   • 시사점: 프러닝 직후의 정확도를 평가하는 근시안적인 메트릭은 효과적이지 않습니다.

2단계 파이프라인

1단계: 중복 레이어 식별
이 단계에서는 서로 다른 유형의 레이어를 직접 비교하지 않고, 얼마나 많은 그리고 어떤 특정 레이어를 프러닝할지 결정합니다.

• 단계 1: 프러닝 예산 할당 (Model Accuracy Predictor - MAP): 레이어를 직접 비교하는 대신, 가벼운 데이터 수집 프로세스(반복적 프러닝 $\to$ 빠른 파인튜닝 $\to$ 평가)를 통해 데이터셋을 구축합니다. 다항 함수를 사용하여 **Model Accuracy Predictor (MAP)**를 근사하며, 이는 Attention 레이어($\tilde{m}_a$)와 Activation 레이어($\tilde{m}_g$)의 프러닝 비율이 주어졌을 때 프러닝된 모델의 최종 정확도를 예측합니다. 이를 통해 시스템은 두 레이어 유형 사이에서 프러닝 예산을 최적으로 할당할 수 있습니다.
• 단계 2: 특정 중복 레이어 제거: 예산이 할당되면, 각 동질적 그룹 내에서(예: 오직 Attention 레이어들 중에서만) 제거할 특정 레이어를 식별합니다. 학습 가능한 중요도 파라미터와 그래디언트 전파를 사용하여 레이어를 순위 매기며, 이를 통해 유형 간의 직접적인 비교가 발생하지 않도록 보장합니다.

2단계: 프러닝된 모델 최적화

• 파인튜닝: 프러닝된 모델을 파인튜닝하며, 정확도 복구를 위해 선택적으로 원본 미프러닝 ViT에 의한 자기 증류(self-distillation)를 사용합니다.
• 병합(Merging): 차원 불일치를 해결하고 추론을 가속화하기 위해, 프러닝된 Activation 레이어에 의해 분리된 인접 Linear 레이어들을 병합합니다. 예를 들어, Activation 레이어가 제거되면 $L_{l,1}$과 $L_{l,2}$가 하나의 Linear 레이어 $L_l$로 병합됩니다. 이는 엄격한 차원 정렬을 유지하면서 모델의 깊이를 줄입니다.

HetDPT+: 저자들은 또한 극단적인 ViT 압축을 위해 HetDPT를 기존의 너비 프러닝 방식과 결래하여 **HetDPT+**라는 파이프라인을 구축하여, 깊이와 너비 차원을 모두 다룹니다.

3. 주요 기여

1. Activation Function 프러닝: 본 논문은 ViT에서 Activation 함수(GELU) 레이어를 제거하는 첫 번째 전략을 제안합니다. 이러한 비선형성을 제거함으로써 인접한 Linear 레이어들을 자연스럽게 병합할 수 있으며, 이는 이질적 깊이 프러닝에 내재된 차원 불일치 문제를 해결합니다.
2. 이질성 분석 및 솔루션: 저자들은 그래디언트 격차와 회복 비대칭성을 밝혀내고 이를 해결했습니다. 이들은 최적의 구성을 보장하는 이론적 분석과 노이즈가 있는 빠른 파인튜닝 환경에서도 일관성을 유지하는 예측 모델을 갖춘 2단계 프레임워크(HetDPT)를 도입했습니다.
3. 최첨단 기록(SOTA):
   • 깊이 프러닝: HetDPT는 거의 손실 없는 정확도로 DeiT-B에서 1.58$\times$, DeiT-S에서 1.39$\times$의 속도 향상을 달야냈으며, 기존의 깊이 프러닝 연구들을 능가했습니다.
   • 극단적 압축: HetDPT+는 극단적인 ViT 압축을 위한 새로운 벤치마크를 설정했습니다. Isomorphic-Pruning-2.6G 구성에서, 거의 손실 없는 정확도를 유지하면서 가속 비율을 4.24$\times$에서 5.19$\times$로 높였습니다.

4. 실험 결과

이 방법은 ImageNet-1K, CIFAR-100, COCO, ADE20K에서 평가되었습니다.

• 정확도-속도 트레이드오프: ImageNet-1K에서 HetDPT는 Attention 또는 Activation 레이어만을 개별적으로 프러닝하는 것보다 일관되게 우수한 성능을 보였습니다. DeiT-B의 경우, 베이스라인과 일치하는 81.8%의 Top-1 정확도를 달성하면서 1.58$\times$의 속도 향상을 이루었습니다.
• SOTA와의 비교: HetDPT+는 기존의 극단적 압축 방식(예: Isomorphic Pruning, NViT)을 능가합니다. 약 81%의 정확도에서 HetDMT+는 Isomorphic Pruning의 4.24$\times$ 대비 5.19$\times$의 속도 향상을 달성했습니다. 초극단적 설정(1.1G MACs)에서는 9.19$\times$의 속도 향상을 달성했습니다.
• 전이 가능성: 이 방법은 CIFAR-100에서 강력한 전이 학습 능력을 보여주었으며, 밀집 예측 작업(COCO 객체 탐지 및 ADE20K 시맨틱 세그멘테이션)에서도 최소한의 정확도 저하를 유지하며 성능을 유지했습니다.
• 확장성: 이 접근 방식은 DINO-V2-Giant에서 15개의 레이어를 프러닝하면서도 정확도 저하를 단 0.06%로 제한함으로써, 수십억 개의 파라미터를 가진 모델로 확장 가능함을 입증했습니다.
• 직교성(Orthogonality): HetDPT는 토큰 프러닝 방식(예: GTP-ViT, ToMe)과 직교하며, 결합 시 추가적인 속도 향상을 제공합니다.

5. 의의 및 주장

본 논문은 주요 기여가 구조적 및 동적 과제를 해결함으로써 Vision Transformer의 효율적인 배포를 가능하게 하는 것이라고 주장합니다.

• 민주화: 계산 요구 사항을 대폭 줄임으로써(최대 9.19$\times$ 속도 향상), 이 방법은 제한된 자원을 가진 연구자와 실무자들이 최첨단 비전 모델을 사용할 수 있도록 만듭니다.
• 지속 가능성: 압축된 모델의 에너지 소비 감소는 환경적으로 지속 가능한 AI 개발에 기여합니다.
• 엣지 배포: 효율성 이득은 의료 진단이나 자율 주행 시스템과 같이 낮은 지연 시간이 필요한 애플리케이션을 위한 엣지 장치에서의 실시간 배포를 용이하게 합니다.

저자들은 한계점으로 현재 Model Accuracy Predictor (MAP)가 각 아키텍처 패밀리에 맞춰 다시 조정되어야 한다는 점을 언급하며, 이를 언어 모델이나 다른 모달리티로 확장하는 것은 향후 과제로 남겨두었습니다. 그러나 본 연구는 기존 문헌의 "일률적인(one-size-fits-all)" 접근 방식을 넘어, 이질적 깊이 프러닝을 위한 새로운 이론적, 실무적 토대를 마련했습니다.