ViT 특징 공간의 과제 유도 리만 계량을 진단하고 저랭크로 학습하는 방법

무엇인가

이 논문은 Vision Transformer(ViT) 특징 공간에서 유클리드 거리나 코사인 유사도를 쓰는 관행이 왜 과제에 맞지 않을 수 있는지에서 출발한다. DPT 깊이, CLS 임베딩, VGGT 카메라 포즈 같은 디코더는 특징의 어떤 방향 변화에는 민감하지만 다른 방향에는 거의 반응하지 않는다. 저자들은 디코더 출력의 특징에 대한 야코비안 J를 이용해 당겨진 계량(pullback metric) g(F)=J(F)^T J(F)를 정의한다. v^T g(F) v = ||J v||^2는 특징을 v 방향으로 미소하게 움직였을 때 과제 출력이 얼마나 변하는지를 나타낸다. 문제는 ViT-B/14 특징 맵에서 J^T J가 약 4×10^10개 항목을 가져 저장이 불가능하고, 깊이 맵처럼 조밀한 출력에서는 J 자체가 약 10^10개 항목이라 만들 수 없다는 점이다. 그래서 논문은 두 질문을 던진다. 과제 민감도가 저랭크로 집중되는가, 그리고 추론 시 역전파 없이 그 계량을 학습할 수 있는가.

어떻게 동작하나

첫 기여는 저랭크 계량 학습 가능성을 미리 판단하는 진단법이다. κ_cap(r)은 J의 상위 r개 특이값 제곱 합을 전체 프로베니우스 노름 제곱으로 나눈 값으로, rank-r 계량이 잡을 수 있는 과제 민감도 비율이다. 또 상위 r개 오른쪽 특이벡터를 토큰 블록으로 나눠 토큰별 에너지가 얼마나 고르게 퍼졌는지를 CV로 잰다. CV가 낮으면 민감 방향이 모든 토큰에 퍼져 있어 토큰별 점수화가 무의미해지고(정리 2), CV가 높으면 소수 토큰에 집중돼 토큰별 점수화가 가능하다. 이 값들은 J를 만들지 않고 Hutchinson 추정, Stochastic Lanczos Quadrature, 랜덤 파워 이터레이션으로 계산되며 기본 설정 m=100, r=20, q=2에서 약 140쌍의 JVP/VJP만 쓴다. 논문은 상위 r개 방향이 잡을 수 있는 민감도 상한(정리 1), 토큰별 인수분해 계량의 불가능성(정리 2), 파워 이터레이션 수렴률(정리 3)을 제시한다. DPT, DINOv2, CLIP, VGGT를 포함한 여섯 조합을 네 체제로 나눈다. VGGT 깊이·포인트클라우드는 유효 랭크가 282와 212로 예산을 크게 넘고, Depth Anything V2의 DPT L2/L10은 CV≈0.026으로 저랭크지만 탈국소화되어 토큰별 점수가 ρ≈0.20에 그친다. DINOv2 CLS는 rank-20이 민감도의 3분의 1만 잡지만 공간 집중이 커서 토큰별 점수가 ρ=0.998에 이르고, VGGT 카메라 헤드는 출력이 9차원이라 κ_cap(20)≈100%가 구조적으로 보장된다.

무엇과 다른가

진단이 허용하는 경우를 위해 Spectral Pullback Network(SPN)를 제안한다. SPN은 특징 F만 입력받아 저랭크 계량을 예측한다. v^T gθ(F)v = Σ_{j=1}^r λ_j(F,θ)(u_j(F,θ)^T v)^2 + ε||v||^2 형태로, u_j는 상위 r개 오른쪽 특이벡터를, λ_j는 σ_j^2를 추정하며 U는 직교정규화된다. 구조는 D→H 선형 투영, 두 개의 pre-norm 다중헤드 어텐션 블록, 토큰별 Dr 좌표 투영을 거쳐 ND×r로 재배열한 뒤 QR 직교정규화하는 방식이고, 에너지는 전역 풀링과 MLP, Softplus로 만든다. 추론 시 gθ(F)v 계산은 O(NDr)로, 전체 행렬을 쓰는 O((ND)^2)를 피한다. 학습은 RandNLA 방식이다. 가우시안 Ω를 뽑고 Ω←J^T(JΩ)를 q번 반복하며 JVP와 VJP를 번갈아 쓰고 재직교화해 상위 r개 목표 부분공간 V_r을 얻는다. 손실은 부분공간 각도를 재는 1 - (1/r)||U^T V||_F^2와 log λ와 log S^2의 MSE로 구성되며, S_j^2=||J v_j||^2이다. 인접 특이값이 거의 겹치면(σ_j/σ_{j+1}<1.05) 모드가 이미지마다 바뀔 수 있어 에너지 손실에서 제외한다. 토큰 중요도만 필요한 응용을 위해 약 310K 파라미터 중요도 헤드로 증류한다. 이 헤드는 한 개의 교차 토큰 MHA 블록(4헤드), 토큰별 MLP, Softplus 출력을 쓰고, 목표는 imp*(F)_t = sqrt(Σ_j σ_j^2 ||v_j^(t)||_2^2)이다. 손실은 로그 MSE에 순위 손실 0.5를 더한다. DINOv2 CLS L10에서 2,000장 ImageNet-val로 학습하고 200장 홀드아웃에서 평가한 중요도 헤드는 목표와 Spearman ρ=0.998±0.001을 기록했고, 탈국소화된 DPT 깊이에서는 같은 헤드가 ρ≈0.20, 단순 블록 노름 목표는 ρ≈0.09에 그쳤다.

어떻게 쓰나

저랭크 근사가 불가능한 조밀 출력에는 VAE 병목을 쓴다. VGGT 깊이와 포인트클라우드는 유효 랭크가 282와 212라 어떤 실용적 rank-r SPN도 충실한 근사가 아니다. 랭크는 min(M,ND) 이하이므로 출력 차원 M을 바꿀 수 없을 때 입력 차원 ND를 줄이는 접근을 택한다. 네 층(4, 11, 17, 23)의 패치 토큰을 이어 붙인 x_vis를 사전학습된 S^2-VAE에 넣어 32개 레지스터 토큰, 차원 128, M'=4096의 잠재 z로 압축하고, 디코더 G가 특징을 복원한 뒤 고정된 과제 헤드가 복원 특징을 받는다. 모든 경로가 z를 지나므로 rank(J)≤M'가 되고, VGGT 깊이의 유효 랭크는 [30,53], r_0.90≈40으로 약 6배 줄어든다. VAE 압축 VGGT 깊이에서 잠재 z를 프로브하면 RandNLA q=1이 진짜 rank-r 과제 민감 부분공간과의 정렬 V1=0.713을 달성해, 결정론적 감독 상한 0.710과 사실상 같고, 랜덤 초기화 0.005와 큰 차이를 보인다. 다만 이때 학습하는 계량은 원래 특징이 아니라 VAE 복원 특징을 통과한 근사 pullback이다.

전제와 한계

토큰 병합·프루닝 실험이 이 중요도 점수를 검증한다. 프루닝의 1차 오차 한계(정리 4)는 제거 토큰의 야코비안 블록 노름과 잔차의 곱에 비례하므로, 중요도가 낮은 토큰을 제거하는 것이 유리하다. CLS 디코더에서는 ToMe의 코사인 유사도를 단일 토큰에 제한한 rank-r pullback 계량으로 바꾸고, 중요도가 낮은 토큰을 중요도 가중 평균으로 병합한다. DPT 같은 조밀 디코더에서는 중요도가 낮은 토큰을 하드 프루닝해 특징을 고정하고 원래 위치에 재삽입한 뒤, 마지막 ViT 블록 전에 다시 넣어 전체 시퀀스로 어텐션하게 하는 Last-Layer Fusion을 쓴다. DPT 깊이에서 prune ratio 0.5일 때 중요도 기반 선택은 ToMe 기반 토큰 선택의 추가 깊이 오차를 25% 줄였고, ViT는 미세조정하지 않았다. 다만 가장 낮은 비율에서는 ToMe가 더 낫다. DINOv2 CLS에서는 중요도 기반 병합이 모든 비율에서 가장 낮은 성능 저하를 보였고 낮은 비율에서는 ToMe 점수보다 한 자릿수 이상 좋았다. ImageNet top-1은 모든 방법과 비율에서 0.66~0.68에 머물렀고 미프루닝은 0.674였다. 조밀 깊이에서는 2단계 [4,8] 스케줄이 NYU-Depth-V2의 모든 비율과 해상도에서 ToMe 점수를 25~35% 이겼지만, 단일 단계 결과는 데이터셋에 의존한다. η=0.20에서 이론 FLOPs의 12.1%를 줄이지만 실제 벽시계 속도 향상은 448×448에서 1.12배에 그친다.

개발자 관점에서 이 논문은 ViT 특징을 코사인 유사도로만 비교하는 파이프라인에 대한 진단 도구를 준다. 고정된 ViT와 디코더를 쓸 때, 먼저 κ_cap(r)과 CV를 오프라인으로 재서 저랭크 계량이 의미 있는지, 토큰별 점수화가 통할지 판단할 수 있다. 유효 랭크가 예산을 크게 넘으면 VAE 병목으로 입력 차원을 줄이는 선택지가 있고, 저랭크지만 탈국소화된 경우에는 토큰별 헤드보다 교차 토큰 어텐션이 있는 SPN이 필요하다. 반대로 CV가 높으면 작은 중요도 헤드로도 충분할 수 있다. 토큰 프루닝·병합을 넣을 때는 과제 민감도 점수를 쓰되, 프루닝 비율, 해상도, 다단계 스케줄, 실제 속도 향상을 함께 확인해야 한다. 출력 거리도 유클리드로 고정된 것이 아니라 양의 정부호 H를 쓰면 J를 LJ로 바꿔 같은 방법을 적용할 수 있다. VGGT 카메라 헤드에서 SO(3)를 고려한 H를 써도 κ_cap=0.99로 추적 가능성이 유지되고 두 토큰 순위의 Spearman ρ=0.995로 중요도가 거의 같았다.

한계도 분명하다. g는 1차 국소 계량이라 멀리 떨어진 두 특징 맵 사이의 전역 거리를 말하지 않으며, 전역 측지선은 부록 K에서 근사할 뿐이다. g는 양의 준정부호이고 퇴화할 수 있어 디코더가 무시하는 방향은 커널에 들어간다. 진단은 어떤 아키텍처가 실행 가능한지 예측하지만 모든 백본-디코더 쌍을 추적 가능하게 만들지는 않는다. SPN 하나가 모든 이미지에 정확히 맞지는 않는다. VAE 압축은 원래 pullback이 아니라 복원 특징을 통과한 근사 pullback이며, 복원 품질에 의존한다. 탈국소화 영역에서는 토큰별 점수화가 약하고, 프루닝 이득은 비율에 따라 비단조적이며 벽시계 속도 향상도 제한적이다. 저자들은 부록 M과 O에서 한계와 향후 방향을 더 다룬다.