희소 오토인코더를 활용한 비전 트랜스포머 동적 헤드 프루닝 제어

본 논문은 Vision Transformer(ViT)에서 동적 헤드 프루닝을 수행할 때, 기존의 잔차 임베딩 기반 정책이 불투명하고 제어가 어려운 문제를 해결하고자 한다. 이를 위해 저자들은 k‑sparse 오토인코더(SAE)를 도입해, 최종 레이어의 CLS 토큰을 고차원 희소 잠재 공간으로 변환하고, 이 잠재를 증폭하거나 억제함으로써 프루닝 마스크를 직접 조작한다. 먼저, AdaViT를 베이스라인으로 사용해 각 레이어마다 헤드 중요도 로짓을 출력하는 경량 결정 네트워크를 학습한다. 이때 CLS 토큰을 입력으로 사용해 전역적인 클래스 정보를 반영한다. 이후, 학습된 ViT에서 최종 레이어 residual stream의 CLS 토큰 x(384 차원)를 추출하고, 이를 인코더 W_enc와 디코더 W_dec을 갖는 k‑sparse 오토인코더에 입력한다. 인코더는 top‑k 활성화를 통해 가장 큰 k=64개의 차원만 남기고, 디코더는 이를 다시 384 차원으로 복원한다. 재구성 손실은 MSE(0.0228)로 매우 낮아, SAE가 원본 정보를 거의 손실 없이 보존함을 확인한다. 핵심 단계는 잠재 z를 선택된 차원 집합 S에 대해 증폭(α)하거나 감소시키는 것이다. 증폭된 잠재 z′는 다시 top‑k를 적용해 ˜z를 만들고, 디코더를 통해 ˜x를 재구성한다. ˜x는 원래 결정 네트워크에 입력되어 새로운 프루닝 마스크를 생성한다. 저자들은 세 가지 S 정의 방식을 실험했는데, (1) 클래스별 빈도 상위 k, (2) 전체 데이터셋에서 전역 빈도 상위 k, (3) 무작위 선택이다. 실험 결과는 클래스별 전략이 가장 효과적이었다. α를 0에서 1.5까지 증가시켰을 때, 클래스별 스티어링은 헤드 사용률을 크게 낮추면서 정확도는 거의 유지하거나 소폭 향상했다. 반면 전역 및 무작위 전략은 α가 커질수록 정확도가 급격히 떨어졌다. 특히 ‘bowl’ 클래스는 정확도가 76%에서 82%로 상승하고 헤드 사용률이 0.72에서 0.33으로 감소했으며, 이는 h₂와 h₅ 헤드가 주요하게 활용된 결과였다. ‘pine tree’도 유사하게 h₂와 h₃을 중심으로 헤드 사용이 최적화되었다. 의미적으로 유사한 클래스(예: bowl과 plate)는 동일한 헤드 서브셋을 공유하는 경향을 보여, SAE 잠재가 클래스 간 의미 구조를 반영한다는 점을 시사한다. 추가 분석에서는 부정적 스티어링(α<0)이 헤드 사용을 증가시키는 반면, 긍정적 스티어링(α>0)은 사용을 감소시키는 것을 확인했다. 이는 잠재 차원의 증폭이 프루닝 신호를 강화하거나 약화시키는 메커니즘을 명확히 보여준다. 또한, SAE 재구성만을 사용했을 때도 프루닝 성능에 큰 변화가 없었으며, 이는 잠재 조작이 실제 프루닝 변화의 원인임을 뒷받침한다. 마지막으로, 저자들은 현재 연구가 최종 레이어와 소규모 데이터셋(CIFAR‑100)에서만 수행됐으며, 향후에는 초기 레이어와 대규모 사전학습된 ViT(예: ViT‑Base, ViT‑Large)에도 적용해 전반적인 효율성과 해석 가능성을 확대할 계획이라고 밝힌다.