Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

기술 요약: Reg4Pru – 토큰 프루닝을 위한 무작위 토큰 라우팅 기반 정규화

문제 정의
비전 트랜스포머(ViT)는 그 확장성과 일반화 능력 덕분에 현대 비전 모델의 근간이 되었습니다. 그러나 셀프 어텐션 메커니즘에 대한 의존도는 토큰 수 $N$에 따라 제곱($O(N^2)$)으로 증가하는 계산 복잡도를 초래합니다. 이러한 병목 현상은 특히 결정적인 진단 세부 정보를 보존하기 위해 고해선 입력을 요구하는 의료 영상 분할 분야에서 더욱 심각하게 나타납니다.

토큰 프루닝(중복된 토큰 제거) 방식과 재활성화(추론 시 밀집된 픽셀 단위 예측을 위해 프루닝된 토큰을 재도입하는 전략)는 효율성을 위한 경로를 제공하지만, 저자들은 깊은 레이어에서 상당한 성능 저하가 관찰됨을 확인했습니다. 구체적으로, 추론 중에 프루닝된 토큰을 재활성화할 때, 이들은 훈련 단계에서의 상태와 판이하게 다른 표현(representation)을 보입니다. 이러한 훈련-테스트 분포 변화(train-test distribution shift)는 불안정성을 유발하며, 모델의 깊이가 깊어질수록 성능이 오히려 하락하는 역설적인 결과를 초래합니다.

방법론
본 논문은 토큰 프루닝 및 재활성화와 관련된 불안정성을 완화하기 위해 설계된 훈련 시 정규화 기술인 Reg4Pru를 제안합니다. 핵심 메커니즘은 토큰을 영구적으로 폐기하지 않고도 프루닝의 효과를 시뮬레이션하기 위해 훈련 단계에서 **무작위 토큰 라우팅(random token routing)**을 사용하는 것입니다.

1. 무작위 토큰 라우팅: 훈련 시 순전파(forward pass) 과정에서, 모델은 중간 트랜스포머 블록 사이에서 무작위 라우팅 범위 $[l, n]$을 선택합니다. 무작위로 선택된 토큰 부분 집합($\tau_r$)은 $l$부터 $n$까지의 블록 계산을 건너뛰도록 선택되며, 나머지 토큰($\tau_k$)은 정상적으로 처리됩니다.
2. 그래디언트 흐름 및 표현: 라우팅된 토큰들은 특정 블록들을 우회하며, 이는 해당 레이어들에 대한 그래디언트를 효과적으로 차단합니다. 그러나 매 단계마다 라우팅되는 토큰의 부분 집합과 라우팅 범위 $[l, n]$을 무작위화함으로써, 서로 다른 토큰들이 여러 블록에 걸쳐 서로 다른 그래디언트에 기여하게 됩니다. 이는 모델이 깊이 불변적 표현(depth-invariant representations)을 학습하도록 강제하여, 추론 시 어떤 깊이에서도 토큰이 표현 불일치 없이 재활성화될 수 있도록 보장합니다.
3. 추론 전략: 추론 시에는 별도의 정책 네트워크(policy network)가 계층적 유지 비율(keep-ratio) 스케줄에 따라 어떤 토큰을 프루닝할지 예측합니다. 프루닝된 토큰은 계산량을 줄이기 위해 일시적으로 제거되었다가, 마지막 블록에서 재활성화되어 밀집된 세그멘테이션 마스크를 생성합니다.
4. 손실 함수: 훈련 목적 함수는 표준 비율 기반 손실(토큰 예산을 강제하기 위함)과, 선택적으로 프루닝 결정을 가이드하기 위해 정답 세그멘테이션 마스크를 사용하는 정보 기반 정책 손실($L_{pol}$)을 결합합니다. 라우팅 메커니즘은 이러한 손실들과 함께 정규화 도구로서 작용합니다.

주요 기여

• 깊이별 불안정성 분석: 저자들은 밀집 예측을 위한 토큰 프루닝에서 발생하는 특정한 불안정성을 식별하고 규명했습니다. 이는 프루닝된 토큰이 훈련 과정에서 네트워크의 전체 깊이에 노출되지 않았기 때문에 더 깊은 레이어에서 성능이 저하되는 현상입니다.
• Reg4Pru 프레임워크: 저자들은 무작위 경계(random bounds)를 가진 훈련 시 토큰 라우팅 전략을 도입했습니다. 고정된 경계나 학습된 라우팅 정책을 사용하는 기존의 라우팅 방식과 달리, Reg4Pin은 라우팅 범위를 무작위화하여 모든 프루닝 깊이에 걸친 특징 분포를 근사함으로써 모델을 정규화합니다.
• 실증적 검증: 본 방법론은 FIVES 망막 혈관 세그멘테이션 데이터셋을 통해 검증되었으며, 무작위 라우팅이 고정 라우팅이나 라우팅을 적용하지 않은 경우보다 성능을 유의미하게 안정화시킨다는 것을 입증했습니다.

결과
실험은 DINOv2-small으로 초기화된 EoMT 아키텍처를 사용하여 FIVES 데이터셋(2048$\times$2048 이미지를 처리량 테스트를 위해 쿼드런트로 리사이즈함)에서 수행되었습니다.

• 성능 향상: 라우팅 없이 훈련된 프룬드 모델과 비교했을 때, Reg4Pru는 ($L_{ratio}$ 구성에서 25.60%에서 71.67%로) 평균 정밀도(AP)를 절대치 기준 46% 개선했습니다.
• 효율성: Reg4Pru 구성은 높은 세그멘테이션 품질을 유지하면서 비프룬드 베이스라인 대비 29%의 상대적 속도 향상을 달s성했습니다.
• 안정성: 깊이별 분석(그림 5) 결과, 라우팅되지 않은 프룬드 모델은 더 깊은 블록에서 AP가 급격히 하락하는 반과 달리, Reg4Pru는 안정적인 성능을 유지했습니다. 고정 경계 라우팅은 일부 불안정성을 완화했으나, 고정된 범위가 끝난 후의 성능 저하를 막지는 못했으며, 이는 무작위화의 필요성을 확인시켜 주었습니다.
• 비교: 유사한 처리량 구성에서 비교했을 때, Reg4Pru는 특히 ToMe가 MHSA와 MLP 레이어 모두에 적용되었을 때보다 AP 측면에서 우수한 성능을 보였습니다.

의의 및 주장
본 논문은 Reg4Pru가 밀집 예측 작업에서 토큰 감소 전략을 위한 가치 있는 정규화 도구라고 주장합니다. 저자들은 프루닝에서 관찰되는 불안정성이 토큰 표현에 관한 훈련-테스트 분포 변화의 결과라고 설명합니다. 훈련 중에 라우팅 범위를 무작위화함으로써, Reg4Pru는 모델이 다양한 깊이에서 스킵되는 것에 견고한 표현을 학습하도록 강제하며, 결과적으로 훈련과 추론 사이의 간극을 메웁니다.

저자들은 실험이 단일 고해상도 의료 데이터셋에서 수행되었다는 점을 언급하며, 연구 결과의 일반화 가능성에 대해 신중한 태도를 유지했습니다. 이들은 본 방법론이 세그멘테이션에는 효과적이지만, 다른 밀집 예측 작업(예: 깊이 추정)에 대한 적용 가능성과 드롭아웃(dropout) 또는 확률적 깊이(stochastic depth)와 같은 기법과 결합되었을 때의 일반적인 정규화 메커즘으로서의 잠재력을 평가하기 위해 추가적인 검증이 필요하다고 제안합니다. 또한, 정보 기반 정책을 탐색했음에도 현재의 설정에서는 비율 기반 정책이 더 우수한 성능을 낸다는 점을 인정했습니다.