[EViT] Not All Patches Are What You Need: Expediting ViTs via Token Reorganizations

Pruning ICLR 2022

Youwei Liang, Chongjian Ge, Zhan Tong, Yibing Song, Jue Wang, Pengtao Xie · UC San Diego / HKU / Tencent AI Lab

한 줄 요약. CLS attention으로 중요한(attentive) 토큰만 top-k 유지하고, 나머지 덜 중요한(inattentive) 토큰은 버리지 않고 1개로 fusion한다. 추가 파라미터·teacher 없이 학습에 녹여, DeiT-S를 +50% 빠르게(−0.3%).

배경

ViT는 모든 패치를 토큰으로 self-attention하지만, 모든 토큰이 예측에 기여하는 건 아니다. 배경·distractive 영역 토큰은 기여가 작다.

(Fig.1) 배경 토큰을 지워도 ViT 예측은 멀쩡하지만, 객체 토큰을 지우면 틀린다.
그렇다고 학습된 ViT에서 inattentive 토큰을 그냥 제거하면 정확도가 급락한다 (DeiT-S, 4·7·10층에서 CLS attention 기준 제거): keep 0.7 → 78.5(−1.3), keep 0.5 → 73.8(−6.0).

즉, “덜 중요한 토큰”을 단순히 버리는 건 위험하다. 버리되 정보를 어떻게 남길 것인가가 관건.

핵심 아이디어

학습 과정에 token reorganization(재구성)을 통합한다. 매 forward마다 CLS attention으로 중요 토큰을 식별해 attentive는 보존, inattentive는 하나로 fusion → 네트워크가 깊어질수록 토큰 수가 줄어 연산 감소. 추가 파라미터가 전혀 없고, DynamicViT와 달리 pretrained ViT(teacher)가 필요 없다.

Figure 2. 한 Transformer 인코더 안의 token reorganization — MHSA 뒤 CLS attentiveness를 계산해, top-k attentive token을 식별·보존하고 inattentive token들을 하나로 fuse한 뒤 FFN으로 보낸다.

방법

1) Attentive token identification (중요 토큰 식별)

self-attention의 attention map 첫 행 = CLS → 모든 토큰 attention. 이 값 aᵢ가 i번째 토큰의 중요도.
여러 head의 CLS attention을 평균(ā)내고, 상위 k개(top-k) 를 attentive token으로 유지. keep rate κ = k/n (하이퍼파라미터).
(DINO 관찰처럼 CLS는 객체 영역 토큰에 더 큰 attention을 준다 → 합리적 기준.)

2) Inattentive token fusion (덜 중요한 토큰 융합)

나머지 inattentive 토큰 N개를 attention 가중 평균으로 단 하나의 토큰으로 합침:

\[x_{\text{fused}} = \sum_{i \in \mathcal{N}} a_i\, x_i\]

이 fused token을 attentive token들 뒤에 붙여 다음 layer로 전달. 연산 비용은 무시할 수준.
단순 제거 대비 장점: ① 배경/객체 일부 정보 보존, ② inattentive 토큰으로도 gradient가 흘러 식별 학습이 좋아지고 ③ 학습이 더 안정적(정확도 분산↓).

3) 학습 설정

삽입 위치: DeiT-S/B는 4·7·10층, LV-ViT-S는 5·9·13층. keep rate를 1 → target으로 cosine warmup.
추가 파라미터 없음, 원 DeiT/LV-ViT 학습 레시피 그대로, scratch 300 epoch. (DynamicViT와 달리 teacher·distillation 불필요.)

결과

Table 2. DeiT-S에서 inattentive token fusion 유무 비교 (keep rate별 top-1/top-5/throughput/MACs).

DeiT-S, keep 0.7: throughput +50%, MACs −35%, top-1 −0.3% (79.5).
fusion이 약간 더 좋고 더 안정적 — 제거만 한 버전보다 정확도·표준편차 모두 우수, 추가 연산 없음.
반대 방향 활용: 같은 연산으로 더 많은 토큰(고해상도) 입력 → 정확도 향상. DeiT-S +1%, LV-ViT-S@448 finetune → 84.7(@384보다 +0.3).
DynamicViT 대비 우위: 같은 연산에서 정확도↑·파라미터↓. 특히 scratch 학습 시 격차 큼 (keep 0.7, 300ep: EViT 79.5 vs DynamicViT 77.6, +1.9) — DynamicViT는 teacher 의존.
더 길게 학습하면 계속 향상(300→600ep: 79.5 → 81.0), oracle ViT 사용 시 80.8.

Figure 4. 정확도–throughput trade-off. EViT(원 표시)가 여러 ViT 대비 더 나은 균형.

시각화 & ablation

Figure 3. 층(4·7·10)이 깊어질수록 inattentive 토큰(마스크 영역)이 점차 융합·제거되고 객체 토큰에 집중 — 해석 가능성도 확보.

마스크 품질 = 정확도 (keep 0.5): top-k 78.5 > random 77.3 > min-k 75.2. CLS attention 기준 선택이 핵심.
삽입 위치: 얕은 층에 두면 정확도 급락(초기 attention은 신뢰도 낮음), 깊은 층들 사이에선 영향 미미.

한 줄 정리 & 의의

“버릴 건 버리되 나머지를 1개로 fuse“라는, pruning에 merging 감각을 더한 접근. 추가 파라미터·teacher 없이 CLS attention만으로 동작하는 게 DynamicViT와의 결정적 차이.
한계 / 이후. CLS attention에 의존(없는 ViT는 token-to-token 변형 필요). 이후 연구는 fusion/merging을 더 정교화(ToMe·MCTF)하거나 pruning·merging을 통합(ToFu·DiffRate)하는 쪽으로 발전. → Token Reduction 개요