Steering Sparse Autoencoder Latents to Control Dynamic Head Pruning in Vision Transformers (Student Abstract)
이 논문은 희소 오토인코더 (SAE) 를 활용하여 비전 트랜스포머 (ViT) 의 밀집 임베딩을 해석 가능한 희소 잠재 공간으로 변환하고, 이를 클래스별로 증폭하여 동적 헤드 가지치기 정책을 제어함으로써 효율성과 기계적 해석 가능성을 동시에 달성하는 새로운 프레임워크를 제안합니다.
1. 문제 상황: 너무 많은 화가들 (비효율성) 이 인공지능 (ViT) 은 이미지를 분석할 때 '주목 (Attention)'이라는 기능을 가진 **여러 명의 화가 (Head)**로 구성된 팀을 운영합니다.
현실: 12 층의 건물 (레이어) 마다 6 명의 화가가 함께 그림을 그립니다.
문제: 그런데 사실 모든 화가가 매번 필요한 건 아닙니다. 어떤 그림은 2 명만 그려도 충분하고, 어떤 그림은 5 명이 필요할 뿐입니다. 하지만 기존 시스템은 "모든 화가를 다 쓰자"라고 하거나, "어떤 화가를 빼야 할지"를 어떻게 결정하는지 알 수 없는 블랙박스 상태로 운영했습니다. (효율은 떨어지고, 왜 그렇게 했는지 설명할 수 없음)
2. 새로운 도구: "의미 분류기" (Sparse Autoencoder, SAE) 저자들은 이 팀의 뇌 (데이터) 를 분석할 수 있는 새로운 도구인 **SAE(희소 오토인코더)**를 도입했습니다.
비유: 이 도구는 팀원들이 서로 섞어서 말하던 복잡한 속삭임 (밀집된 데이터) 을 받아, **"이 화가는 '그릇' 그릴 때만 특출나고", "저 화가는 '나무' 그릴 때만 잘한다"**는 식으로 **명확한 역할 (단일 의미)**을 가진 개별적인 신호로 분리해 줍니다.
마치 팀원들의 성격을 "그릇 전문가", "나무 전문가"로 딱딱 분류해 주는 명찰을 달아주는 것과 같습니다.
3. 해결책: "지시자의 손짓" (Latent Steering) 이제 이 분리된 신호들을 이용해 팀을 지휘합니다.
기존 방식: "그림을 그려!"라고만 하면, 팀이 알아서 (알 수 없는 이유로) 화가들을 골라냅니다.
새로운 방식: "이번엔 그릇을 그리는 거야!"라고 특정 신호 (Latent) 를 **강조 (증폭)**해 줍니다.
결과: 지시자가 "그릇" 신호를 키우자, 팀은 자연스럽게 **그릇을 그리는 데만 특화된 화가들 (예: 2 번, 5 번 화가)**만 남기고 나머지는 쫓아냅니다.
📊 실제 성과: "그릇"을 그릴 때의 변화
논문의 실험 결과, 이 방법이 얼마나 효과적인지 보여줍니다.
상황: '그릇 (Bowl)'이라는 이미지를 인식할 때.
기존: 화가 6 명 중 0.72 명 (약 4 명) 을 썼는데, 정확도는 76% 였습니다.
새로운 방법 (지시자 활용): '그릇' 신호를 키우자, **2 명 (2 번, 5 번 화가)**만 남겼습니다.
결과: 화가 수는 절반으로 줄었는데, 정확도는 오히려 76% 에서 82% 로 올라갔습니다!
즉, 불필요한 화가들을 쫓아내니 오히려 더 집중해서 잘 그리는 효과가 나왔습니다.
💡 핵심 요약
효율성: 인공지능이 이미지를 볼 때, 상황에 따라 필요한 '화가 (Head)'만 골라내어 연산량을 줄였습니다.
이해 가능성: "왜 이 화가들을 골랐을까?"라고 물으면, **"그림의 종류 (그릇, 나무 등) 에 따라 특화된 역할이 있기 때문"**이라고 명확하게 설명할 수 있게 되었습니다.
통제 가능성: 우리가 "이런 그림은 이렇게 그려!"라고 신호를 보내면, 인공지능이 그 신호에 맞춰 화가들을 자동으로 조정할 수 있게 되었습니다.
🚀 결론
이 연구는 인공지능이 "더 적은 일로 더 잘" 하도록 만들면서, 동시에 "왜 그렇게 했는지" 인간이 이해하고 통제할 수 있게 하는 다리를 놓았습니다. 마치 복잡한 오케스트라에서 악기 소리를 분리해, 필요한 악기만 연주하게 지휘하는 것과 같습니다.
앞으로 이 기술은 더 큰 인공지능 모델에서도 적용되어, 더 빠르고 투명하게 작동하는 AI 를 만드는 데 기여할 것으로 기대됩니다.
1. 연구 배경 및 문제 정의 (Problem)
비전 트랜스포머 (ViT) 의 비효율성: ViT 는 다양한 토큰 상호작용을 포착하기 위해 멀티헤드 셀프 어텐션을 사용하지만, 많은 어텐션 헤드가 중복되어 계산 비용만 증가시키고 성능 향상에는 기여하지 않는 경우가 많습니다.
기존 동적 가지치기 (Dynamic Pruning) 의 한계: AdaViT 와 같은 기존 적응형 프레임워크는 입력에 따라 불필요한 헤드를 제거하여 효율성을 높입니다. 그러나 이러한 가지치기 정책은 주로 잔여 임베딩 (residual embeddings) 에 의존하여 결정되므로, 어떤 기준으로 헤드가 선택되는지 그 과정이 불투명 (opaque) 하고 잠재 수준 (latent level) 에서 제어하기 어렵습니다.
해결 필요성: 가지치기 효율성과 기계적 해석 가능성 (mechanistic interpretability) 을 동시에 달성하기 위해, 가지치기 결정을 해석 가능하고 제어 가능한 방식으로 전환할 수 있는 새로운 접근법이 필요합니다.
2. 제안 방법론 (Methodology)
저자들은 희소 오토인코더 (Sparse Autoencoder, SAE) 를 ViT 의 동적 가지치기 프레임워크에 통합하여 가지치기 결정을 잠재 수준에서 제어하는 새로운 프레임워크를 제안합니다.
기본 아키텍처:
베이스라인: AdaViT 를 기반으로 하며, 각 층 (layer) 의 CLS 토큰을 입력으로 받아 헤드의 중요도 (logits) 를 예측하는 경량 결정 네트워크 (decision network) 를 사용합니다.
SAE 학습: ViT 의 최종 층 잔여 임베딩 (CLS 토큰) 을 입력으로 받아 k-sparse 오토인코더를 학습시킵니다.
인코더는 고밀도 임베딩을 희소한 잠재 표현 (z) 으로 변환합니다.
$TopK활성화함수를사용하여상위k$개의 값만 보존하고 나머지는 0 으로 만듭니다.
재구성 손실 (MSE) 을 최소화하도록 학습합니다.
잠재 공간 조정 (Latent Steering):
학습된 SAE 의 특정 잠재 차원 (latent dimensions) 을 증폭 (amplify) 시킵니다.
특정 클래스에 해당하는 빈번하게 활성화된 잠재 변수 집합 (S) 을 선택하고, 그 값에 α만큼 더합니다 (zi′=zi+α).
증폭된 벡터를 다시 $TopK처리하여조정된희소벡터(\tilde{z})를얻고,이를디코더를통해재구성된임베딩(\tilde{x}$) 으로 변환합니다.
이 조정된 임베딩을 결정 네트워크에 입력하여 가지치기 마스크 (pruning mask) 를 생성합니다.
3. 주요 기여 (Key Contributions)
해석 가능하고 제어 가능한 가지치기 프레임워크: SAE 를 활용하여 ViT 의 밀집 임베딩을 해석 가능한 희소 특성으로 분리하고, 이를 통해 동적 가지치기 결정을 잠재 수준에서 직접 제어할 수 있음을 입증했습니다.
클래스별 가지치기 전략 (Per-class Steering): 특정 클래스에 빈번하게 활성화되는 잠재 특성을 증폭시킴으로써, 해당 클래스에 최적화된 헤드의 하위 집합을 자동으로 발견하고 가지치기 효율을 극대화하는 전략을 제시했습니다.
메커니즘적 해석 가능성 확보: 어떤 헤드가 어떤 클래스의 의미론적 특성을 담당하는지 (예: 'bowl' 클래스는 h2,h5 헤드를 주로 사용) 를 시각화하고 설명할 수 있게 되었습니다.
4. 실험 결과 (Results)
실험 설정: CIFAR-100 데이터셋에서 ImageNet 사전 학습된 ViT-Small 을 미세 조정 (Fine-tuning) 하였습니다.
SAE 성능: SAE 로 재구성된 임베딩을 사용해도 정확도 저하가 거의 없었으며 (-0.12%), 가지치기 비율에도 큰 변화가 없어 SAE 가 핵심 정보를 잘 보존함을 확인했습니다.
가지치기 전략 비교:
클래스별 빈도 기반 (Per-class frequent): 정확도를 유지하면서 헤드 사용량을 크게 줄였습니다.
전체 빈도 기반 (Global frequent) 및 무작위: 정확도가 급격히 떨어지는 경향을 보였습니다. 이는 SAE 가 클래스별 구별 가능한 개념을 포착했음을 의미합니다.
구체적 사례 (Case Study):
'bowl' 클래스: 정확도가 76% 에서 82% 로 상승하고, 헤드 사용량이 0.72 에서 0.33 으로 감소했습니다. 주로 h2와 h5 헤드를 활용했습니다.
의미론적 연관성: 'bowl'과 'plate'와 같이 의미적으로 유사한 클래스들은 유사한 헤드의 하위 집합 (h2,h5) 을 공유하는 것을 확인하여, 조정된 잠재 공간이 클래스 간 의미론적 관계를 반영함을 보였습니다.
5. 의의 및 결론 (Significance & Conclusion)
효율성과 해석 가능성의 교량: 이 연구는 ViT 의 동적 가지치기가 단순히 계산 효율을 높이는 것을 넘어, 어떤 헤드가 왜 선택되는지를 이해하고 제어할 수 있는 길을 열었습니다.
향후 과제: 현재는 최종 층과 소규모 데이터셋에 국한되어 있으나, 향후 더 초기 층과 대규모 기초 모델 (Foundation Models) 로 확장하여 적용 범위를 넓힐 계획입니다.
요약하자면, 이 논문은 SAE 를 통해 ViT 의 내부 표현을 해석 가능한 희소 특성으로 변환하고, 이를 조정하여 클래스별 최적의 헤드 가지치기를 수행함으로써, 효율적인 모델 압축과 기계적 해석 가능성 (Interpretability) 을 동시에 달성하는 새로운 패러다임을 제시했습니다.