Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

기술 요약: 레지스터의 활용: Vision Transformer를 위한 토큰 프루닝용 태스크 레지스터

문제 정의

Vision Transformer(ViT)는 토큰 수에 따라 계산량이 이차 함수적으로 증가하므로, 추론 비용을 줄이기 위한 토큰 프루닝(token pruning)이 필요합니다. 프루닝 기준(예: 어텐션 점수, 공간적 커버리지)은 주로 이미지 분류를 위해 개발되었으나, 사전 학습된 ViT 백본은 시맨틱 세그멘테이션 및 객체 탐지와 같은 밀집 예측(dense prediction) 작업에 점점 더 많이 적응되고 있습니다.

저자들은 다음과 같은 핵심적인 격차를 식별했습니다: 한 작업에 설계된 프루닝 정책이 다른 작업으로 효과적으로 전이되지 않을 수 있다는 점입니다. 통제된 "프로브(probe)" 실험을 통해 논문은 다음을 입증했습니다:

1. 기준 순위의 불일치(Criterion Ranking Divergence): 세그멘테이션과 탐지 작업은 프루닝 기준의 순위를 다르게 매깁니다. 예를 들어, 어텐션 기반 프루닝은 분류 작업의 초기 레이어에서 상당한 정확도 저하를 일으키는 반면, 공간적 커버리지는 더 안정적입니다. 반대로, 탐지 오차를 최소화하는 기준은 세그멘테이션에서는 성능이 저하되는 경우가 많습니다.
2. 깊이 민감도(Depth Sensitivity): 프루닝 기준의 효과는 네트워크 깊이에 따라 크게 달라지며, 작업 간에도 차이가 있습니다.
3. 복구 갈등(Recovery Conflicts): 밀집 예측 작업은 제거된 위치에서의 특징 복구(feature reconstruction)를 필요로 합니다. 본 논문은 세그멘테이션과 탐지가 서로 반대되는 복구 엔드포인트를 선호한다는 것을 발견했습니다. 세그멘테이션은 전체 특징 오프셋(제거된 토큰과 유지된 토큰 사이의 차이)을 복원할 때 이득을 얻는 반면, 탐지는 유지된 "스탠드인(stand-in)" 토큰만을 사용할 때 더 나은 성능을 보입니다.

기존 방법들은 일반적으로 단일한 정적 프루닝 정책이나 작업 불가지론적(task-agnostic) 선택기를 사용하며, 이러한 작업별 및 깊이 의존적 변동을 고려하지 못합니다.

방법론: 태스크 적응형 프루닝 (Task-Adaptive Pruning, TAP)

논문은 **태스크 레지스터(Task Registers)**를 사용하여 토큰 프루닝, 예산 할당 및 특징 복구를 동적으로 제어하는 태스크 적응형 프루닝(TAP) 프레임워크를 제안합니다.

핵심 메커니즘:

• 태스크 레지스터: 레지스터를 단순히 (Darcet et al., 2024와 같은 이전 연구처럼) 특징 아티팩트의 작업 불가지론적 저장소로 사용하는 대신, TAP는 작업당 하나씩 학습 가능한 초기 레지스터를 도입합니다. 순전파(forward pass) 동안 현재 작업과 관련된 레지스터만 활성화됩니다.
• 진화하는 상태(Evolving State): 활성화된 레지스터는 백본을 통과하며 이미지 토큰과 함께 진화합니다. 임의의 레이어에서 레지스터의 상태는 세 가지 핵심 결정 사항을 조건화합니다:
  1. 토큰 선택: 레지스터는 블록의 기존 쿼리(query) 및 키(key) 프로젝션을 재사용하여 각 후보 패치에 대한 점수를 계산합니다. 토큰의 순위가 매겨지고, 일부 하위 집합이 셀프 어텐션을 위해 유지됩니다.
  2. 정확한 예산 할당: 레지스터 상태에 따라 조건화된 공유 선형 리드아웃(linear readout)이 각 프루닝 레이어에서 제거할 정확한 정수 개수의 토큰을 결정합니다. 이를 통해 레이트 손실(rate loss) 없이 글로벌 제거 예산을 정확히 충족합니다.
  3. 복구 스케일링: 밀집 작업을 위해, 레지스터는 리드아웃 시 저장된 특징 오프셋(제거된 토큰과 스탠드인 사이의 차이)을 얼마나 복원할지 결정하는 스케일링 인자($\alpha$)를 제어합니다.

구현 변형:

• TAP-J (Joint Fine-tuning): 각 작업은 별도의 백본, 태스크 레지스터 및 리드아웃을 가집니다. 이를 통해 완전한 특성화를 구현할 수 있습니다.
• TAP-F (Frozen Base): 하나의 공유된 프리즈(frozen) 백본을 모든 작업이 공유합니다. 작업별 레지스터, 저계수(low-rank) 업데이트(LoRA) 및 헤드가 공동으로 학습됩니다. 이는 적응성을 유지하면서 파라미터 저장 용량을 줄여줍니다.

특징 복구:
제거된 토큰은 키 공간(key-space) 코사인 유사도가 가장 높은 유지된 토큰(즉, "스탠드인")과 결정론적으로 매칭됩니다. 특징 차이(오프셋)가 저장됩니다. 밀집 리드아웃 지점에서 복구는 $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$로 계산되며, 여기서 $\alpha$는 작업 및 깊이별로 학습됩니다. 분류 작업은 이 복구 단계를 건너뜁니다.

주요 기여

1. 통제된 프레임워크: 저자들은 프루닝이 없는 체크포인트를 동결하고, 분류, 세그멘테이션, 탐지 전반에 걸쳐 프루닝 선택의 영향을 격리하기 위해 파라미터가 없는 감소 기준을 레이어별로 적용하는 프레임워크를 개발했습니다. 이를 통해 기준 순위의 불일치를 밝혀냈습니다.
2. 컨트롤러로서의 태스크 레지스터: 본 논문은 레지스터의 역할을 수동적인 저장소에서 능동적이고 작업 특화적인 희소 계산(sparse computation)의 컨트롤러로 재정의합니다. 레지스터는 토큰 선택, 깊이에 따른 예산 배분, 복구 스케일링을 조정합니다.
3. 통합된 프루닝 및 복구: TAP는 토큰 제거와 특징 복구를 하나의 통합된 연산으로 취급합니다. 이는 버려진 토큰을 스탠드인과 매칭하고, 오프셋의 복원을 태스크 레지스터에 따라 스케일링하며, 인코더 외부의 희소성을 유지하기 위해 복구된 특징을 인코더 내부에 두지 않습니다.

실험 결과

저자들은 $\rho = 0.5$ (50% 토큰 유지)인 ViT-B 백본을 사용하여 ImageNet-1K (분류), ADE20K (세그멘테이션), COCO (탐지/인스턴스 세그멘테이션)에서 TAP를 평가했습니다.

• 성능:
  • TAP-J는 완전 미세 조정(fully fine-tuned) 방식들 중 가장 작은 정확도 저하를 달로 성취했습니다. ADE20K에서 47.0 mIoU(프루닝 미적용 대비 0.2 하락)와 1.30× 인코더 처리량을 달성했습니다. COCO에서는 53.7 box AP(0.3 하락)와 1.32× 처리량을 달성했습니다. ImageNet-1K에서는 83.2 top-1(0.4 하락)에 도달했습니다.
  • TAP-F (공유 백본)는 TAP-J보다 약간 뒤처지지만, 정적 규칙 및 다른 베이스라인보다 유의미하게 우수한 성능을 보이며 ADE20K에서 45.8 mIoU, COCO에서 52.5 box AP를 달성하면서도 파라미터 오버헤드는 약 1.2M 수준으로 매우 낮았습니다.
• 효율성: TAP-J는 경쟁력 있는 정확도를 유지하면서 1.26×에서 1.32×의 인코더 속도 향상을 제공합니다. 컨트롤러 오버헤드(레지스터 쿼리, 예산 리드아웃, 북키핑)는 인코더 런타임의 약 0.7%만을 차지합니다.
• 메모리: TAP는 밀집 베이스라인 대비 피크 추론 메모리를 9~18% 감소시키며, 중간 블록의 활성화 크기 감소 덕분에 세그멘테이션에서 가장 큰 이득을 보았습니다.
• 어블레이션 연구: 태스크 레지스터를 제거하거나(공유된 것 사용) 할당/복구를 정적 값으로 고정하면 성능이 크게 저하되어, 태스크 및 깊이 적응형 제어의 필요성을 확인시켜 주었습니다.

의의 및 주장

본 논문은 토큰 감소를 단일한 전이 가능한 규칙으로 취급해서는 안 된다고 주장합니다. ViT 표현이 국소적 구조에서 고수준 의미론(semantics)으로 진화함에 따라, 토큰의 가치는 깊이와 백본을 읽어들이는 특정 작업 모두에 따라 변화합니다.

TAP의 의의는 공유된 백본을 통해 서로 다른 시각적 작업의 특정 요구 사항에 희소 계산을 적응시킬 수 있는 능력에 있습니다. 진화하는 태스크 레지스터를 활용함으로써 TAP는:

• 작업 간 프루닝 기준 순위의 갈등을 해결합니다.
• 네트워크 깊이에 따른 제거 예산 배분을 최적화합니다.
• 밀집 vs 희소 리드아웃의 특정 요구 사항에 맞춰 특징 복구를 조건화합니다.

저자들은 TAP를 최소한의 작업별 저장 공간을 사용하여 여러 작업에 대응할 수 있는 통합된 비전 시스템을 향한 단계로 포지셔닝하며, 정적 프루닝 정책보다 예측 성능과 인코더 처리량 사이의 균형을 더 효과적으로 맞춥니다. 저자들은 분류 작업이 클래스 토큰에 의존하고 밀집 복구를 건너뛰기 때문에, TAP의 메커니즘 중 하나가 해당 작업에서 충분히 활용되지 못해 TAP-J가 최고의 분류 전용 방법(DiffRate)보다 0.1점 뒤처진다는 점을 언급하며 겸허한 입장을 유지했습니다. 향후 연구로는 이미지 조건부 예산(image-conditioned budgets) 및 다중 스탠드인 탐색이 제안되었습니다.