Reg4Pru: 토큰 프루닝을 위한 무작위 라우팅 정규화

초록

Reg4Pru는 훈련 시 무작위 토큰 라우팅을 적용해 토큰 프루닝의 깊이별 불안정을 완화하고, 혈관 영상 분할 데이터셋(FIVES)에서 평균 정밀도를 46% 절대 상승시키면서 29%의 실시간 속도 향상을 달성한 정규화 기법이다.

상세 분석

본 논문은 고해상도 의료 영상에서 Transformer 기반 모델이 토큰 수에 따라 O(N²) 연산 비용을 갖는 문제를 해결하고자 한다. 기존 토큰 프루닝 기법은 불필요한 토큰을 제거해 연산량을 감소시키지만, 프루닝된 토큰을 추후 재활성화할 때 깊은 층에서 표현이 크게 변형돼 성능 저하가 발생한다는 점을 지적한다. 저자들은 이 현상을 “깊이별 불안정성”이라 명명하고, 훈련 단계에서 실제 프루닝과 유사한 효과를 모방하기 위해 무작위 라우팅(Random Token Routing)을 도입한다. 구체적으로, 매 훈련 포워드 패스마다 라우팅 구간