기술 요약: 비전 트랜스포머의 해석 가능성에 기반한 어텐션 헤드 소프트 프루닝 (SAPER)

문제 정의

DINOv2와 같은 자기지도 학습 아키텍처를 포함한 비전 파운데이션 모델들은 다양한 다운스트림 태스크에서 놀라운 성능을 달나 성취했습니다. 그러나 이들의 성공은 거대하고 과잉 매개변수화된 아키텍처에 의존하며, 이는 상당한 계산 병목 현상과 메모리 요구량을 초래합니다. 멀티 헤드 셀프 어텐션(multi-head self-attention) 메커니즘은 강력하지만, 모델의 불투명성과 중복성을 심화시키는 원인이 됩니다.

본 논문이 다루는 핵심 문제는 개별 어텐션 헤드의 구체적인 기능적 역할에 대한 이해 부족과 모델 내에 존재하는 기능적 중복성의 존재입니다. 단순히 성능 지표만을 근거로 모델 크기를 줄이는 방식은 네트워크의 내부 메커니즘을 고려하지 못해 실패하는 경우가 많습니다. 또한, 이러한 모델들이 에지 디바이스로 이식됨에 따라, 유지되는 모든 헤드는 레이턴시, 메모리 사용량, 에너지 소비와 직결되므로, "그린 AI(Green AI)" 의제에 부합하는 구조적인 헤드 수준의 압축 전략이 필요합니다. 저자들은 이러한 중복성이 결함이 아니라, 방대한 데이터셋에서 범용 표현을 학습하는 과정에서 발생하는 창발적 결과라고 상정합니다. 즉, 잠재적인 다운스트림 태스크 전반에 걸쳐 견고함을 보장하기 위해 많은 병렬 경로가 유지되는 것입니다.

방법론

본 논문은 해석 가능성과 효율성을 연결하는 엔드 투 엔드 미분 가능한 프루닝 프레임워크인 **SAPER (Soft Attention PrunER)**를 소개합니다. 방법론은 세 가지 주요 단계로 진행됩니다.

1. 스펙트럼 분석 및 시각화 (라플라스 맵, Laplace Maps)

어텐션 헤드의 기능적 구조를 이해하기 위해 저자들은 그래프 이론적 접근 방식을 제안합니다. 어텐션 맵은 토큰을 정점으로, 어텐션 가중치를 간선으로 하는 가중 그래프로 취급됩니다.

• 라플라시안 고유벡터 (Laplacian Eigenvectors): 이 방법은 각 헤드에 대해 어텐션 그래프의 정규화된 라플라시안(normalized Laplacian)을 계산합니다.
• 시각화: 두 번째, 세 번째, 네 번째로 작은 고유값에 대응하는 고유벡터들을 추출합니다. 두 번째로 작은 고유값에 대응하는 고유벡터(피들러 벡터, Fiedler vector)는 헤드의 지배적인 분할 동작을 포착합니다.
• RGB 매핑: 이 세 개의 고유벡터를 R, G, B 채널에 매핑하여 **라플라스 맵(Laplace maps)**을 생성합니다. 이 시각화는 초기 레이어의 헤드들이 종종 "컨볼루션(convolutional)" 동작(글로벌 패턴이나 에지에 반응하는 방식)을 보이는 반면, 깊은 레이어의 헤드들은 객체 수준의 구조에 집중한다는 것을 보여줍니다.

2. 시맨틱 클러스터링 (Semantic Clustering)

저자들은 네트워크 전반의 기능적 중복성을 식별하기 위해 시맨틱 클러스터링을 수행합니다.

• 특징 추출: 이미지 데이터셋에 대해 소프트맥스 전 단계의 어텐션 로짓 행렬(pre-softmax attention logit matrices)을 추출합니다. 공간적 어텐션 확률 벡터(CLS 토큰에서 공간 패치로의 흐름)를 분석합니다.
• 유사도 행렬: 공간 패치 전반에 걸친 최대 어텐션 값을 사용하여 특징 행렬을 구성합니다. 모든 헤드 쌍 사이의 피어슨 상관 계수를 계산하여 유사도 행렬을 생성합니다.
• 클러스터링: 스펙트럴 클러스터링(spectral clustering)을 적용하여 헤드들을 별개의 클러스터로 분할합니다. 연속적인 블록 구조를 가정하는 기존 연구(예: RAPTOR)와 달리, 이 방식은 생성된 표현에 기초하여 개별 헤드를 클러스터에 할당하며, 이를 통해 기능적으로 유사한 헤드가 서로 다른 네트워크 깊이에 존재할 수 있음을 밝혀냅니다.

3. 미분 가능한 헤드 선택 (LapSum Soft Top-K)

SAPER는 이산적이고 미분 불가능한 연산 없이 헤드를 프루닝할 수 있는 미분 가능한 선택 메커니즘을 사용합니다.

• LapSum: LapSum 알고리즘을 활용하여 이산적인 Top-K 선택 문제를 소프트 완화(soft relaxation) 형태로 공식화합니다.
• 메커니즘: 헤드에 대한 점수 벡터가 주어지면, LapSum은 라플라스 분포의 누적 분포 함수(CDF)를 사용하여 소프트 선택 확률을 계산합니다. 음수의 온도 파라미터($\alpha$)는 선택의 날카로움을 조절하며, $\alpha \to 0$일 때 소프트 마스크는 하드 Top-K 선택에 근접합니다.
• 학습 파이프라인: 모델은 미분 가능한 마스크를 통해 Top-K 헤드가 선택되는 "소프트" 학습 단계를 거친 후, 선택된 헤드만을 사용하는 "하드" 파인튜닝 또는 지식 증류(knowledge distillation) 단계를 거칩니다.

주요 기여

1. 라플라스 맵 (Laplace Maps): 어텐션-그래프 라플라시안을 기반으로 개별 어텐션 헤드의 동작(예: 에지 검출 vs 글로벌 시맨틱 구조)을 해석할 수 있게 해주는 새로운 스펙트럼 시각화 기술입니다.
2. 헤드의 시맨틱 클러스터링: 헤드 전체가 아닌 개별 헤드를 클러스터링함으로써 네트워크 전반의 공유된 표현을 식별하는 방법입니다. 이는 비전 트랜스포머(ViT)의 표현 단계(representational phases)를 발견하는 것을 일반화하여, 기능적 중복성이 인접한 블록 내에서뿐만 아니라 다양한 깊이 전반에 걸쳐 존재함을 보여줍니다.
3. SAPER 프레임워크: 경쟁력 있는 분류 성능을 유지하면서 상당한 FLOPs 감소를 달하는 미분 가능한 Top-K 헤드 선택 프레임워크입니다. 이는 연속적인 블록 배수로 제한되는 RAPTOR의 방식보다 더 높은 유연성을 제공하며 임의의 개별 헤드를 선택할 수 있게 합니다.

실험 결과

저자들은 DINOv2 백본(ViT-Small 및 ViT-Base)을 사용하여 ImageNet-1K 및 CIFAR-100에서 SAPER를 평가했습니다.

• 정확도-효율성 트레이드오프 (Accuracy-Efficiency Trade-off): SAPER는 매우 유리한 트레이드오프를 보여줍니다. 예를 들어, 72개 중 6개의 헤드만 유지한 ViT-S 모델에서 SAPER는 지식 증류를 통해 약 **73%**의 분류 정확도를 달성하였으며, 이는 24개의 헤드를 가진 모델보다 약 8% 낮은 수준입니다.
• RAPTOR와의 비교: SAPER는 헤드 수 측면에서 경쟁적인 RAPTOR 베이스라인보다 우수한 FLOPs 감소를 보여줍니다. 동일한 헤드 수(예: 24개)에서 SAPER는 RAPTOR보다 GFLOPs를 약 2배 더 줄였습니다(ViT-B 기준 61.86 대비 32.17 GFLOPs). 이는 SAPER가 RAPTOR의 제약 조건인 연속적인 블록 배수에 얽매이지 않고 개별 헤드를 선택할 수 있기 때문입니다.
• 데이터셋 강건성: CIFAR-100에서 SAPER는 모든 유사한 헤드 예산 범위에서 RAPTOR보다 일관되게 우수한 성능을 보였으며, 훨씬 적은 GFLOPs를 사용했습니다.
• 학습 모드: 연구는 순수 파인튜닝과 지식 증류를 비교했습니다. 지식 증류가 특히 ImageNet-1K에서 더 나은 결과를 냈으며, 저자들은 파인튜닝 중에 사전 학습된 어텐션 헤드를 동결(frozen) 상태로 유지하는 것이 종종 유익했다고 언급했습니다.

의의 및 주장

본 논문은 해석 가능하고 투명한 비전 파운데이션 모델을 향한 진일보를 표방합니다. 본 연구의 주요 의의는 다음과 같습니다:

• 기계적 해석 가능성 (Mechanistic Interpretability): 스펙트럼 시각화를 통해 개별 어텐션 헤드의 동작을 미리 볼 수 있는 도구를 제공함으로써, 블랙박스 형태의 프루닝을 넘어섭니다.
• 유연한 프루닝: 프루닝을 결함의 수리가 아닌 일반 모델의 특화(specialization) 과정으로 바라봅니다. 기능적으로 중독된 헤드를 식별하고 제거함으로써, 치명적인 정확도 손실 없이 특정 효율성 제약에 맞춰 모델을 최적화할 수 있습니다.
• 그린 AI (Green AI): 계산 비용(FLOPs)을 대폭 줄이면서도 강력한 성능을 유지함으로써, 자원이 제한된 에지 디바이스에 대규모 파운데이션 모델을 배포하기 위한 실질적인 솔루션을 제공합니다.

저자들은 자신들의 접근 방식이 해석 가능성과 효율성 사이의 간극을 메우며, 기능적으로 유사한 헤드가 서로 다른 네트워크 깊이에 존재하고 이를 선택적으로 유지함으로써 모델 성능을 최적화할 수 있다는 가설을 입증한다고 결론짓습니다. 향후 연구로는 이 프레임워크를 토큰 레벨로 확장하고 MLP 가중치까지 포함하는 방안이 제시되었습니다.