기술 요약: SnapViT – 사전 학습된 모델으로부터의 재학습 없는 탄력적 비전 트랜스포머

문제 정의

비전 파운데이션 모델들은 다양한 작업에서 최첨단 성능을 달ato 달성했지만, 일반적으로 제한된 세트의 정해진 크기(예: DINOv3 또는 DeIT 제품군의 특정 파라미터 수)로 출시됩니다. 이러한 경직성은 실무자들이 계산량, 지연 시간, 비용과 같은 실제 환경의 제약에 직면했을 때 최적이 아닌 배포 선택을 하도록 강요합니다. 기존의 모델 유연성 솔루션들은 다음과 같은 중대한 한계에 직면해 있습니다:

• 지식 증류(Knowledge Distillation): 미리 결정된 타겟 아키텍처를 필요로 하며, 종종 공개되지 않은 사전 학습 데이터셋에 의존하므로 견고성이 저해됩니다.
• 탄력적 추론(Elastic Inference): Matryoshka 표현법이나 Slimmable Networks와 같은 방법들은 사전 학습 단계에서 특정 중첩 구조를 갖도록 설계되어야 하므로, 기존의 기성(off-the-shelf) 사전 학습 모델에는 적용할 수 없습니다.
• 구조적 프루닝(Structured Pruning): 메모리와 계산량을 줄이는 데 효과적이지만, 현재 대부분의 프루닝 기술은 특정 작업이나 제약 조건에 맞춰져 있으며 일반적으로 재학습을 요구합니다. 이는 레이블이 있는 데이터나 분류 헤드가 필요 없는 보편적으로 적응 가능한 모델 사이의 간극을 남깁니다.

방법론: SnapViT

저자들은 SnapViT(Single-shot network approximation for pruned Vision Transformers)를 제안합니다. 이는 재학습이나 레이블 데이터 없이 단일 사전 학습 모델로부터 연속적인 탄력적 서브 네트워크를 생성하는 포스트 프리트레이닝(post-pretraining) 구조적 프루닝 방법입니다. 이 방법의 핵심은 제거할 파라미터를 순위 매기기 위한 프루너빌리티 점수(prunability score, $P$)이며, 이는 두 가지 구성 요소로 이루어집니다.

1. 로컬 헤시안 근사 (자기 지도 학습 그래디언트)

분류 헤드 없이 로컬 파라미터 민감도를 포착하기 위해, 이 방법은 자기 지도 학습 그래디언트를 사용하여 헤시안 행렬의 대각 성분을 근사합니다. 구체적으로, 이미지 데이터셋에 대해 DINO 목적 함수를 사용하여 그래디언트를 계산합니다. 로컬 곡률 $H^{(l)}$은 다음과 같이 그래디언트 제곱 노름의 평균으로 추정됩니다:
$$H^{(l)} \approx \frac{1}{N_D} \sum_{i=1}^{N_D} \|\nabla_\theta L_i\|^2$$
이는 개별 파라미터가 특정 블록(예: 피드 포워드 뉴런 또는 어텐션 헤드) 내에서 얼마나 중요한지에 대한 기초적인 척도를 제공합니다.

2. 글로벌 헤시안 추정 (진화 알고리즘)

블록 간 의존성(오프-다이아고날 헤시안 요소)을 무시하는 한계를 해결하기 위해, SnapViT은 **지수적 자연 진화 전략(Exponential Natural Evolution Strategy, xNES)**을 사용하는 새로운 근사법을 도입합니다.

• xNES는 직접적인 계산이 불가능한 전체 헤시안을 구하는 대신, 로컬 민감도를 재가중하는 구조별 스케일링 인자($c$)의 집합을 학습합니다.
• 알고리즘은 이러한 스케일링 인자에 대해 다변량 가우시안 분포를 최적화합니다. 적합도 함수(fitness function)는 레이블이 없는 방식으로, PCA 압축 후 원본 모델과 프루닝된 모델 간의 임베딩 코사인 유사도를 측정합니다.
• 자연 경사 업데이트를 통해, 진화 전략의 공분산 행렬($\Sigma$)은 실제 인터-블록 헤시안의 역행렬($H^{(g)} \approx \alpha \Sigma^{-1}$)을 근사하도록 진화합니다. 이를 통해 어떤 블록들을 공동으로 프루닝해도 표현 품질이 저하되지 않는지 식별할 수 있습니다.

3. 통합 프루너빌리티 점수 및 탄력적 생성

최종 프루너빌리티 점수는 로컬 민감도와 글로벌 스케일링 인자의 요소별 곱입니다:
$$P = \text{diag}(H^{(l)}) \odot (M c)$$
여기서 $M$은 블록 수준의 인자를 개별 파라미터로 확장하는 멤버십 행렬입니다.

• 싱글 샷 생성: $P$가 계산되면 모든 파라미터의 순위가 전역적으로 매겨집니다. 이 단일 순위 덕분에 추가적인 최적화 없이 어떠한 원하는 희소도 수준($S \in [0, 1]$)에 대해서도 서브 네트워크를 추출할 수 있습니다.
• 범위: 이 방법은 피드 포워드 블록 내의 행-열 조합과 전체 어텐션 헤드를 모두 프루닝합니다.

주요 기여

1. 탄력적 프루닝 전략: 재학습 없이도 다양한 계산 제약에 적응할 수 있는 탄력적 모델을 생성하는, 사전 학습된 비전 트랜스포머를 위한 효율적이고 빠른 프루닝 전략을 제공합니다.
2. 진화적 헤시안 근사: xGES(xNES)를 사용하여 네트워크 구조에 대한 오프-다이아고날 성분을 근사함으로써, 전통적인 대각 근사법이 놓치는 교차 블록 상관관계를 포착하는 새로운 전략을 제시합니다.
3. 레이블 프리, 헤드 불가지론적 성능: 레이블 데이터나 분류 헤드 없이도 최첨단 성능을 달성하며, 감독 학습 및 자기 지도 학습 기반의 파운데이션 모델 모두에서 일반화됩니다.

실험 결과

저자들은 k-최근접 이웃(k-NN) 분류, 선형 프로빙(linear probing), 선형 시맨틱 세그멘테이션을 사용하여 8개의 데이터셋(ImageNet-1k, Pascal VOC 및 다양한 미세 조정 벤치마크 포함)에서 SnapViT을 평가했습니다.

• SOTA 대비 성능: SnapViT은 다양한 희소도 수준에서 SOTA 방법들(LAMP, LLM Surgeon, FPTP, SparseGPT, SNIP Magnitude, NViT)과 일치하거나 이를 능가하는 성능을 일관되게 보여줍니다.
  • DINOv1 ViT-B/16의 경우, 정확도 저하를 5% 미만으로 유지하면서 40% 희소도까지 프루닝할 수 있으며, 1.58배의 추론 속도 향상을 달ach 합니다.
  • AugReg 모델에서의 50% 희소도에서 SNIP보다 7%, FPTP보다 12.3% 높은 선형 분류 성능을 보였습니다.
  • SigLIPv2 및 DINOv3(대규모 파운데이션 모델)의 경우, 이 방법은 베이스라인을 크게 앞지르며, 두 번째로 우수한 방법 대비 선형 분류에서 최대 34.3%의 개선을 보였습니다.
• 효율성: DeIT-III와 같은 모델 제품군에 대한 전체 프루닝 과정은 단일 A100 GPU에서 5분 미만이 소요됩니다.
• 후처리: 이 방법은 포스트 프루닝 가중치 교정(SparseGPT 스타일) 및 전체 파인 튜닝과 호환됩니다. 단일 가중치 교정 단계를 거치면 극단적인 희소도(예: SigLIPv2에서 50%)에서도 성능이 거의 유지됩니다.
• 확장성: 대규모 사전 학습(예: DINOv3, SigLIPv2)은 분산된 표현 지식으로 인해 프루닝을 더 어렵게 만들지만, SnapViT은 가중치 교정 또는 더 긴 최적화 기간(500회 xNES 반복)과 결합될 때 여전히 효과적입니다.

의의 및 주장

본 논문은 SnapViT이 현재 파운데이션 모델 배포의 경직성과 유연하고 자원 적응적인 추론에 대한 필요성 사이의 간극을 메운다고 주장합니다. 단일 샷으로 하나의 사전 학습된 모델로부터 연속적인 서브 네트워크를 추출할 수 있게 함으로써, 사용자가 자신의 특정 계산 예산에 맞춰 최첨단 모델을 정밀하게 맞춤화할 수 있도록 합니다.

저자들은 다음과 같은 점을 강조합니다:

• 재학습의 필요성을 제거하여 상당한 계산 자원과 시간을 절약합니다.
• 레이블 데이터에 대한 의존성을 제거하여 분류 헤드가 없는 자기 지도 학습 파운데이션 모델에도 적용 가능하게 합니다.
• 진화적 근사를 통해 글로벌 구조적 상관관계를 포착하여, 로컬 전용 측정법의 한계를 극복합니다.
• 비전 트랜스포머의 계산 요구량과 전력 사용량을 낮춤으로써 환경적 영향(CO2 배출 감소 등)을 줄이는 데 기여합니다.

이 연구는 고성능을 유지하면서 제약된 환경에 대규모 비전 모델을 배포하기 위한 실질적인 솔루션으로 제시됩니다.