토큰 퓨전: 효율적인 비전 트랜스포머를 위한 토큰 프루닝과 병합의 연결

1. 서론 및 개요

비전 트랜스포머(ViTs)는 컴퓨터 비전 분야에 혁명을 일으켰지만, 입력 토큰 수에 대한 셀프 어텐션의 2차 복잡도로 인해 높은 계산 비용이 발생합니다. 본 논문, Token Fusion: Bridging the Gap between Token Pruning and Token Merging은 토큰 퓨전(ToFu)을 소개합니다. 이는 효율성과 정확도의 균형을 최적화하기 위해 모델 동작을 기반으로 토큰 프루닝과 병합 중 동적으로 선택하는 하이브리드 방법입니다.

핵심 통찰은 프루닝(토큰 폐기)과 병합(토큰 평균화) 중 어느 것도 보편적으로 최적이 아니라는 점입니다. 본 논문은 레이어별 적절한 연산을 선택하는 원칙적인 방법과, 표준 평균 병합에서 발생하는 분포 변화 문제를 해결하기 위한 MLERP(Multi-token Linear intERPolation)이라는 새로운 병합 기술을 제안합니다.

2. 핵심 방법론: 토큰 퓨전 (ToFu)

ToFu는 보간된 입력에 대한 모델의 반응을 분석하여 병합 또는 프루닝에 대한 적합성을 결정하는 데 기반을 둡니다.

2.1. 프루닝 대 병합 딜레마

저자들은 핵심 기준인 모델 선형성을 확인했습니다. 모델 레이어가 보간된 입력에 대해 거의 선형적으로 반응한다면(예: $f(\alpha x_1 + (1-\alpha)x_2) \approx \alpha f(x_1) + (1-\alpha)f(x_2)$), 평균화를 통한 유사 토큰 병합은 효과적이며 정보를 보존합니다. 그러나 초기/후기 비선형 레이어(그림 1에서 시각화됨)에서는 입력 공간의 선형 보간이 높은 비선형 출력을 초래하여 평균 병합에 문제를 일으키고 잠재적으로 분포 변화를 유발할 수 있습니다. 이러한 경우, 중요도가 낮은 토큰을 프루닝하는 것이 정보 손실은 있더라도 더 안전한 대안입니다.

2.2. ToFu 프레임워크

ToFu는 트랜스포머 블록별로 작동합니다:

토큰 중요도 점수화: 각 토큰에 중요도 점수를 할당합니다(예: 어텐션 노름 또는 그래디언트 기반).
선형성 평가: 레이어의 근사 선형성을 평가하며, 경험적으로 또는 경량 프로브를 통해 도출됩니다.
적응형 연산: 목표 토큰 감소 비율에 대해:
- 고선형성 영역: 중요도가 가장 낮은 토큰을 가장 유사한 중요도 높은 이웃 토큰과 병합합니다.
- 저선형성 영역: 중요도가 가장 낮은 토큰을 직접 프루닝합니다.

이는 동적이고 문맥을 인지하는 압축 파이프라인을 생성합니다.

2.3. MLERP: 노름 보존 병합

단순 평균화를 개선하기 위해 저자들은 $K$개의 토큰을 병합하기 위한 구면 선형 보간(SLERP)을 적용한 MLERP를 제안합니다. 노름 $n_i = ||t_i||$를 가진 토큰 $t_1, t_2, ..., t_K$에 대해, MLERP는 먼저 단위 구면에서 방향을 보간한 다음 원래 노름의 가중 평균으로 스케일링합니다:

$t_{\text{merged}} = \left( \frac{\sum_{i=1}^K w_i n_i}{\| \sum_{i=1}^K w_i \frac{t_i}{n_i} \|} \right) \left( \sum_{i=1}^K w_i \frac{t_i}{n_i} \right)$

여기서 $w_i$는 중요도 기반 가중치입니다. 이는 특징의 통계적 노름 분포를 보존하여, 순진한 평균화로 인한 분포 변화를 완화하고 특히 비선형 영역에서 더 안정적인 성능으로 이어집니다.

3. 기술적 세부사항 및 수학적 공식화

본 논문은 토큰 감소 문제를 공식화합니다. 레이어에 $N$개의 입력 토큰 $T = \{t_1, ..., t_N\}$이 있다고 가정합니다. 목표는 $M < N$개의 토큰을 가진 감소된 집합 $T'$을 생성하는 것입니다.

핵심 방정식:

중요도 점수: $I(t_i) = ||\text{Attn}(t_i)||_1$ 또는 그래디언트 기반 측정법.
유사도 측정법: 일반적으로 코사인 유사도 $S(t_i, t_j) = \frac{t_i \cdot t_j}{||t_i|| \, ||t_j||}$.
선형성 측정법 ($\mathcal{L}$): 입력의 선형 보간으로부터 레이어 출력의 편차로 측정됩니다. 낮은 $\mathcal{L}$은 병합을 선호하고, 높은 $\mathcal{L}$은 프루닝을 선호합니다.

ToFu 알고리즘은 파인튜닝 없이(제로샷) 사전 훈련된 모델에 적용하거나 경량 훈련으로 향상시킬 수 있습니다.

4. 실험 결과 및 성능

저자들은 ToFu를 이미지 분류(ImageNet with ViT-B/16, DeiT) 및 이미지 생성(잠재 확산 모델) 작업에서 평가합니다.

주요 성능 하이라이트

분류: ToFu는 독립적인 프루닝(예: DynamicViT) 또는 병합(ToMe) 방법보다 더 나은 정확도 대 FLOPs 균형을 달성합니다. 예를 들어, FLOPs 40% 감소 시 ToFu는 ImageNet에서 <0.5%의 top-1 정확도 손실을 보이며, ToMe보다 약 0.3% 우수합니다.
이미지 생성: Stable Diffusion에서 ToFu는 ToMe와 비교하여 감소된 계산 비용으로 더 높은 시각적 충실도(FID 측정)를 유지하며, 특히 많은 수의 토큰을 감소시킬 때 두드러집니다. MLERP 병합은 출력 분포가 중요한 생성 작업에서 더 명확한 이점을 보입니다.
어블레이션: 적응형 전략(병합/프루닝 선택)은 모든 레이어에서 단일 연산만 사용하는 것보다 우수함이 입증되었습니다. MLERP는 일관되게 평균 병합을 능가합니다.

차트 설명 (논문의 그림 1 기반): 이 그림은 ViT 레이어의 비선형성을 보여줍니다. 두 입력 특징점(x1, x2)이 선형으로 보간됩니다(색상 선). ViT 내부의 네 가지 다른 MLP 레이어에서의 출력(f1-f4)이 표시됩니다. 초기 및 후기 MLP 출력(f1, f4)은 직선에서 상당한 편차를 보여 강한 비선형성을 나타냅니다. 두 입력의 평균(보라색 별)은 출력 평균에서 멀리 떨어진 출력 점에 매핑되며, 평균 병합이 비선형 레이어에서 실패할 수 있는 이유를 시각적으로 보여줍니다.

5. 분석 프레임워크 및 사례 예시

사례: 엣지 배포를 위해 사전 훈련된 ViT에 ToFu 적용

시나리오: 개발자가 실시간 이미지 분류를 위해 모바일 장치에서 ViT-B 모델을 실행해야 합니다. 전체 모델은 너무 느립니다.

프레임워크 적용:

프로파일링: 작은 캘리브레이션 데이터셋을 모델에 통과시킵니다. 각 트랜스포머 블록에 대해 토큰 쌍을 샘플링하고 출력 보간 오류를 확인하여 선형성 측정법 $\mathcal{L}$을 계산합니다.
전략 맵: 프로파일 생성: 블록 1-3(저선형성) → 프루닝 선호. 블록 4-8(고선형성) → MLERP 병합 선호. 최종 블록(저선형성) → 프루닝 선호.
구성: 전역 토큰 감소 예산 설정(예: 35%). 중요도 점수에서 도출된 블록별 예산을 준수하면서 저선형성 블록에서는 프루닝을, 고선형성 블록에서는 MLERP 병합을 적용합니다.
평가: 압축된 모델을 배포합니다. 적응형 접근 방식은 민감한 비선형 레이어에서 공격적인 병합을 피함으로써, 일괄 적용 방법에 비해 최소한의 정확도 하락을 보장합니다.

이 예시는 ToFu가 단일 알고리즘이 아닌 구조화된 압축 프레임워크로서의 실용적 유용성을 보여줍니다.

6. 향후 응용 및 연구 방향

멀티모달 트랜스포머: 토큰 역학이 더 복잡한 비디오, 오디오 또는 멀티모달(예: CLIP, Flamingo) 트랜스포머로 ToFu 확장.
하드웨어 인지 공동 설계: 특정 AI 가속기(NPU, GPU)에 대한 ToFu 결정 알고리즘(프루닝/병합) 및 MLERP 구현 최적화로 실제 속도 향상 극대화.
다른 기술과의 통합: ToFu를 양자화, 지식 증류 또는 효율적인 어텐션 메커니즘(Linformer 등)과 결합하여 복합적인 효율성 향상 도모.
자동화된 하이퍼파라미터 탐색: 신경망 구조 탐색(NAS) 또는 강화 학습을 사용하여 최적의 레이어별 프루닝/병합 비율 및 선형성 임계값 자동 결정.
비전 이상: 토큰 의미론이 크게 다르지만, 시퀀스 압축을 위한 대규모 언어 모델(LLM)에서의 효능 탐구.

7. 참고문헌

Dosovitskiy, A., et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR 2021.
Bolya, D., et al. "Token Merging: Your ViT But Faster." ICLR 2023 (ToMe).
Wang, Y., et al. "DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification." NeurIPS 2021.
Rombach, R., et al. "High-Resolution Image Synthesis with Latent Diffusion Models." CVPR 2022.
Zhu, J.Y., et al. "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks." ICCV 2017 (CycleGAN).
Vaswani, A., et al. "Attention Is All You Need." NeurIPS 2017.

8. 전문가 분석 및 비판적 통찰

핵심 통찰: ToFu는 단순한 또 다른 압축 도구가 아닙니다. 이는 트랜스포머 레이어가 이질적이라는 점에 대한 공식적인 인식입니다. 모든 레이어를 동일한 압축 원시 연산으로 처리하는 것은 순진합니다. 본 논문의 탁월함은 진단적 접근법—레이어 선형성을 측정하여 올바른 "수술"(프루닝 또는 병합)을 처방하는—에 있습니다. 이는 현대 컴파일러가 최적화를 적용하기 위해 코드를 프로파일링하는 방식과 유사하며, ML 효율성 연구에서 종종 부족한 수준의 정교함입니다.

논리적 흐름: 주장은 설득력이 있습니다: 1) 평균 병합이 비선형 레이어에서 실패함을 보임(그림 1). 2) 이 실패 모드를 감지하는 측정법(선형성) 제안. 3) 측정법을 사용하여 토큰을 라우팅. 4) MLERP로 실패하는 연산(평균 병합) 수정. 문제 식별부터 다중 구성 요소 솔루션까지의 흐름이 깔끔하고 논리적입니다.

강점과 약점:
강점: 하이브리드 접근법은 이론적으로 타당하며 작업 전반에 걸쳐 경험적으로 검증되었습니다. MLERP는 실제 문제(노름 붕괴)에 대한 간단하면서도 영리한 해결책입니다. 제로샷 적용 가능성은 기존 모델 배포에 있어 주요 실용적 이점입니다.
약점: 논문은 "선형성 평가"의 오버헤드를 약간 과소평가합니다. 사전 계산된 프로파일(정적)인가, 아니면 실시간 계산(동적 오버헤드)인가? MLERP의 이점은 분명하지만 분류 작업에서는 미미해 보입니다. 그 진정한 가치는 출력 분포가 최우선인 확산 모델 문헌의 발견과 일치하며, 생성 작업에서 더 두드러지는 것으로 보입니다. 비교는 공정하지만, 직교적 이점을 제공하는 최신 사후 훈련 양자화 방법에 대해 더 공격적일 수 있습니다.

실행 가능한 통찰: 실무자에게: 특히 생성 작업을 위한 ViT의 1차 토큰 감소 방법으로 ToFu/MLERP를 즉시 채택하십시오. 이는 기본 병합 전략으로서 ToMe를 대체합니다. 연구자에게: "레이어 인지 압축" 패러다임이 핵심 시사점입니다. 향후 연구는 CNN의 네트워크 프루닝 또는 CycleGAN과 같은 GAN의 모드 붕괴 분석 작업에서 영감을 얻어, 압축에 친화적인 모델 영역과 압축에 민감한 영역의 자동 감지에 초점을 맞춰야 합니다. 다음 프론티어는 사후 압축을 넘어, 이러한 진단 연구의 통찰력을 활용하여 구조 탐색에 정보를 제공함으로써 본질적으로 효율적인 모델을 설계하는 것입니다.