STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

기술 요약: STARFISH

문제 정의
신경망 프루닝(Pruning)은 가중치, 뉴런 또는 채널을 제거함으로써 모델의 크기, 메모리 점유율 및 추론 지연 시간을 줄이는 표준적인 기술이다. 그러나 공격적인 프루닝은 종종 예측 정확도의 상당한 저하를 초래한다. 성능 손실을 회복하기 위해 복구 과정(Healing)이 일반적으로 사용되지만, 기존 방법들은 중대한 한계에 직면해 있다. 전통적인 접근 방식은 원래의 훈련 데이터셋에 대한 전체 재학습 또는 미세 조정(Fine-tuning)을 요구하는 경우가 많은데, 이는 계산 비용이 많이 들 뿐만 아니라, 훈련 데이터가 독점적이거나 접근 불가능한 현대적 환경에서는 실행이 불가능한 경우가 많다. 또한, 최첨단 복구 방법들은 높은 희소성 조건(예: 가중치의 75% 이상 제거) 하에서, 특히 현대적인 Vision Transformer(ViT) 아키텍처에 대해 정확도를 복구하는 데 어려움을 겪는다.

방법론: STARFISH
저자들은 원래의 훈련 데이터나 레이블에 대한 접근 없이 작동하는 사후 프루닝 복구 방법인 STARFISH(faST Accuracy Recovery in pruned networks From Internal State Healing)를 제안한다. 핵심 가설은 프루닝된 네트워크와 원래의 밀집(Dense) 네트워크 사이의 중간 표현(Intermediate representations, 내부 상태)의 정렬을 최적화하는 것이 최종 출력 예측만을 최적화하는 것보다 정확도 복구에 더 효과적인 대리 지표(Proxy)라는 것이다.

알고리즘은 훈련 세트와는 구별되지만 데이터 분포에서 샘플링된 작은 미표기 교정 세트($S_{cal}$)를 사용하여 두 단계로 진행된다:

1. 타겟 표현 캐싱(Caching Target Representations): 원래의 밀집 모델을 $S_{cal}$에 실행하여 모든 계산 블록(예: 정규화, 비선형성 및 잔차 합산 이후)의 벡터화된 출력을 캐싱한다.
2. 정렬 최적화(Alignment Optimization): 프루닝된 모델($f_{\theta \odot m}$)은 모든 블록 $\ell$에 대해 프루닝된 모델의 은닉 상태($f^\ell_{\theta \odot m}(x)$)와 캐싱된 타겟($r^\ell_x$) 사이의 정렬 손실을 최소화하도록 훈련된다. 손실 함수는 다음과 같이 정의된다:
   $$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \ndot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$$
   프러닝된 모델의 남은 가중치들은 교정 세트에 대한 경험적 손실을 최소화하도록 업데이트된다.

이론적 동기는 최종 은닉 표현(분류 헤드로의 입력) 사이의 거리가 출력 분포의 쿨백-라이블러(KL) 발산의 상한선 역할을 한다는 것을 보여줌으로써 제공된다. 이는 내부 상태를 정렬하는 것이 본질적으로 출력 오차를 제약한다는 것을 시사한다.

주요 기여

1. 표현 기반 복구(Representation-Based Recovery): 이 논문은 특정 프루닝 유형, 마스크 또는 아키텍처에 독립적인 복구 방법을 소개한다. 저자들은 중간 표현 정렬이 높은 희소성 수준에서 직접적인 출력 기반 목적 함수(예: Cross-Entropy 또는 KL 발산 매칭)보다 우수함을 경험적으로 입증한다.
2. 고희소성 효율성(High-Sparsity Efficiency): 이 방법은 극단적인 희소성 수준(예: 85% 희소성)에서 원래 정확도의 높은 비율을 효율적으로 복구하는 첫 번째 방법이다. 이는 매우 작은 교정 세트(단 1,000개의 예시, 즉 훈련 세트 크기의 0.4%)를 사용하여 달성되며, 원래의 훈련 데이터에 대한 접근을 요구하지 않는다.
3. 모듈성 및 속도: 이 접근 방식은 계산적으로 가벼워 작은 데이터셋에 대해 몇 번의 순전파(Forward) 및 역전파(Backward) 패스만을 필요로 한다. 이는 구조적(Structured), 반구조적(Semi-structured) 및 비구조적(Unstructured) 프루닝 패턴 모두에 적용 가능하다.

실험 결과
저자들은 다양한 Vision Transformer(DeiT, ViT) 및 CNN(MobileNet) 아키텍처에 대해 ImageNet-1K에서 STARFISH를 평가했다.

• 고희소성 (0.75–0.85): 75% 희소도로 프루닝된 DeiT-B 모델에서, STARFISH는 원래 밀집 정확도의 82%를 복구한 반면, 경쟁 방법(CORP 등)은 40%만을 복구했다. 85% 희소도에서의 DeiT3-H 모델에 대해, STARFISH는 밀집 정확도의 92.08%를 유지하며, 프루닝된 모델의 성능을 72.36 퍼센트 포인트 향상시켰다.
• 중등도 희소성 (0.5): 다양한 DeiT 아키텍처와 프루닝 범위(MLP, QK, 그리고 합동 MLP+QK)에 걸쳐, STARFISH는 일관되게 최첨단 방법들을 능가했다. 예를 들어, 50% 구조적 프루닝(MLP+QK)이 적용된 DeiT-T 모델에서, STARFISH는 밀집 정확도의 81.95%를 유지하며 CORP 방법보다 31.54 퍼센트 포인트 높은 성능을 보였다.
• 비교: STARFISH는 특히 프루닝이 모델을 심하게 저하시키는 영역에서 Hessian 근사(SNOWS)나 폐쇄형 아핀 복구(CORP)와 같은 방법들에 비해 우수한 성능을 입증했다.

의의 및 주장
이 논문은 STARFISH가 복구 패러다임을 출력 매칭에서 내부 상태 치유(Internal state healing)로 전환함으로써 모델 압축 분야의 중요한 진보를 나타낸다고 주장한다. 저자들은 이 접근 방식이 훈련 데이터를 사용할 수 없고 전체 재학습이 불가능한 현대의 대규모 모델에서 특히 필수적임을 강조한다. 고희소성 조건 하에서 내부 표현이 최종 출력보다 모델의 지식을 더 견고하게 포착한다는 원리를 활용함으로써, STARFISH는 예측 성능을 희생하지 않고도 매우 효율적인 프루닝된 네트워크를 배포할 수 있게 한다. 이 방법은 초기 프루닝 전략과 관계없이 모든 프루닝된 네트워크에 적용할 수 있는 일반적인 플러그 앤 플레이(Plug-and-play) 솔루션으로 제시된다.