GTR이 게이트 선형 순환으로 고해상도 밀집 예측 비용을 줄인다
GTR: Gated Token Recurrence for Efficient Dense Prediction
무엇인가
자기어텐션 기반 비전 백본은 밀집 예측에서 좋은 성능을 내지만, 이미지 토큰 수에 대해 전역 소프트맥스 어텐션이 제곱 비용을 갖기 때문에 해상도가 올라갈수록 비용이 급격히 커진다. 특히 자원이 제한된 엣지 장치에서 이 문제가 크다. EfficientViT나 ViG 같은 선형 어텐션 계열 연구가 있었지만, 모델 크기·연산·메모리·지연 시간 제약을 만족하면서 여러 밀집 예측 과제에서 일관되게 정확한 범용 백본은 아직 없다는 것이 이 논문의 출발점이다. 저자들은 선형 어텐션이 이론적 효율을 실제 엣지 하드웨어 이득으로 바꿀 수 있는지 묻는다.
어떻게 동작하나
제안하는 GTR은 토큰 수에 선형인 단일 스케일 순환 백본이다. 이미지를 stride-16 패치 그리드로 잘라 12개 블록 동안 같은 해상도를 유지하고, 각 블록이 게이트 선형 어텐션(GLA)과 Spatial SwiGLU를 결합한다. GLA는 키 전용 변형을 쓴다. 토큰별로 상태 S_t = Diag(γ_t)·S_{t-1} + k_t^T v_t 를 갱신하고 o_t = d_k^(-1/2)·q_t·S_t 로 출력을 낸다. 감쇠 게이트 γ_t는 저랭크 투영(rank 16, temperature 16)에 시그모이드를 적용해 (0,1) 범위로 만들며, 키 쪽 상태만 감쇠시키고 값 쪽 감쇠는 학습·추론 모두 1로 고정한다. RMSNorm을 거친 출력에는 별도의 데이터 의존 출력 게이트 r_t = SiLU(x_t·W_r)를 곱한다. 순환은 스캔 순서를 따라 인과적이므로 한 방향만으로는 2차원 문맥을 담지 못한다. 그래서 블록마다 좌→우, 우→좌, 상→하, 하→상 네 방향을 번갈아 적용하고 스캔 후 원래 공간 순서를 복원한다. 지역 상호작용은 Spatial SwiGLU가 맡는다. Y = [SiLU(XW_g) ⊙ DWConv3x3(XW_v)]·W_o 형태로, 값 분기를 2차원 패치 그리드로 되돌려 3x3 depthwise 컨볼루션을 적용한 뒤 다시 시퀀스 순서로 편다. 게이팅 분기는 토큰 단위로 남아 있고, GTR은 별도의 위치 임베딩을 쓰지 않는다.
무엇과 다른가
증류는 의도적으로 단순하다. 검출용으로 적응시킨 DINOv3 백본을 고정 교사로 두고, 교사와 학생이 같은 전체 이미지를 받아 공간적으로 대응하는 패치 토큰을 만든다. 학생의 최종 패치 토큰에만 학습된 선형 투영 φ를 적용해 교사의 최종 표현과 제곱 L2 손실로 맞춘다. 마스킹도, 중간 레이어 감독도, 어텐션 모듈 정렬도 없다. GTR과 투영만 최적화되며, 이렇게 얻은 가중치가 이후 검출기 사전학습과 과제별 학습의 백본 초기값이 된다. 추론 커널도 별도로 설계했다. GLA의 청크 단위 계산을 세 단계로 나눠 청크 요약을 병렬로 구하고, 청크 경계 상태만 순차 스캔으로 전파한 뒤, 모든 청크 출력을 다시 독립적으로 계산한다. 이 연산자는 추론 전용이며 FP16, 청크 크기 64, 헤드당 d_k=32·d_v=64, 초기 상태 0이라는 GTR 설정에 특화되어 있고, 학습은 기존 FLA 커널을 그대로 쓴다. RTX 4090에서 FP16·CUDA Graph 기준 FLA v0.5.0 대비 256토큰에서 2.6배, 1.6K에서 4.0배, 16.4K에서 6.4배 빠르다.
어떻게 쓰나
검출 실험은 Objects365에서 36 에폭 사전학습 후 COCO train2017에서 30 에폭 학습하고 640x640에서 평가했다. GTR-S/M/L/X가 각각 53.6/57.3/58.9/59.4 box AP로 네 스케일 모두에서 비교 대상 중 최고 AP를 기록했다. GTR-L은 58.9 AP를 1.908ms 중앙값 지연으로 내며 L 스케일 베이스라인을 AP와 지연 모두에서 앞섰고, GTR-X는 59.4 AP를 2.115ms에 냈는데 RF-DETR-X는 58.6 AP에 3.243ms였다. 같은 백본을 과제별 헤드만 바꿔 전이한 결과도 제시된다. 인스턴스 분할은 GTR-S/X가 45.0/49.8 mask AP(SAM2 의사 마스크 없이)로 GTR-X가 RF-DETR-Seg-X보다 25.3% 낮은 지연을 보였고, 자세 추정은 70.1/75.5 keypoint AP로 ECPose를 S·X 양쪽에서 앞섰다. DOTA-v1.0 방향 검출은 GTR-X가 81.3 AP50을 3.960ms에 내 YOLO26x-obb의 5.085ms보다 빨랐고, Cityscapes 의미 분할은 경량 FCN 헤드로 81.5/83.6 mIoU를 기록해 GTR-X가 YOLO26-sem과 같은 83.6 mIoU를 19.5% 낮은 지연으로 달성했다. NYU Depth V2 단안 깊이 추정에서도 δ1, AbsRel, RMSE가 YOLO26 깊이 변형보다 개선됐다.
전제와 한계
절제 실험은 설계 선택의 근거를 수치로 보여준다. 같은 교사·GTR-S 학생·COCO 학습 조건에서 최종 출력 정렬만 쓰는 방식이 50.7 AP로 ViT-Linearizer(50.0)와 ViT-AdaLA(48.2)를 앞섰고, 증류 없이 처음부터 학습하면 31.1 AP에 그쳤다. 학습된 위치 임베딩을 S-SwiGLU로 바꾸면 768²/1024²에서 1.2/1.3 AP 오르는데 추가 연산은 0.2/0.5 GFLOPs뿐이다. 스캔 방향은 GTR-L에서 네 방향이 55.5 AP로 수평·수직 양방향(55.1/55.0)과 단일 방향(54.4)보다 좋았고, 단일 방향에서 S-SwiGLU를 위치 임베딩으로 되돌리면 45.6 AP까지 떨어진다. 12개 블록 중 3개를 소프트맥스 어텐션으로 바꾸면 정확도는 55.4~55.6 AP로 사실상 같지만 지연이 1.908ms에서 2.094~2.117ms로, 연산이 106에서 119 GFLOPs로 늘어난다. 해상도를 640²에서 1280²로 올리면 AP가 53.6에서 57.0으로 오르고 작은 객체 AP는 36.4에서 42.4로 6.0 오르는데, 토큰이 4배 늘어도 지연은 1.225→2.529ms로 2.06배만 증가하고 피크 메모리는 0.096→0.249GB가 된다. 배포는 NVIDIA DRIVE AGX Thor에서 TensorRT로 그래프·커널 공동 최적화를 적용해 6개 과제·4개 스케일 전 조합에서 2.282~8.769ms 중앙값 FP16 배치1 지연을 냈고, 48개 과제·스케일·배치 구성 모두 PyTorch FP32 참조 대비 코사인 유사도 0.9989 이상을 유지했다.
실무 관점에서 이 논문이 주는 신호는 두 가지다. 첫째, 고해상도 입력에서 지연과 메모리가 문제라면 소프트맥스 어텐션을 일부 섞는 하이브리드가 정확도 이득 없이 비용만 늘린다는 절제 결과(55.5 AP 대 55.4~55.6 AP, 1.908ms 대 2.094~2.117ms)는 구조 선택에 직접 참고할 만하다. 둘째, 선형 어텐션의 이론적 이득은 커널 구현에서 나온다. 저자들이 FLA 대비 4배 이득을 낸 지점은 청크 요약을 병렬화하고 경계 상태만 순차 전파하는 실행 스케줄 변경이며, 이는 계산 자체를 바꾸지 않는다. 다만 이득이 토큰 길이에 따라 2.6배(256)에서 6.4배(16.4K)로 달라지므로, 실제 서비스 해상도에서 자기 토큰 길이 기준으로 다시 벤치마크해야 한다. 또한 공개된 지연 수치는 RTX 4090에서 FP16·배치1·torch.compile·CUDA Graphs, 워밍업 후 CUDA 이벤트 중앙값이라는 조건이고, 베이스라인 정확도는 원 출처에서 가져오되 지연은 저자 환경에서 다시 측정한 값이라는 점을 감안해야 한다.
저자가 명시한 전제와 한계도 분명하다. 전용 청크 CUDA 연산자는 추론 전용이고 FP16, 청크 크기 64, d_k=32, d_v=64, 초기 상태 0이라는 특정 구성에 맞춰져 있어 다른 설정에 그대로 옮겨지지 않는다. 학습은 참조 FLA 커널을 계속 사용한다. 순환 상태에서는 값 쪽 감쇠를 1로 고정하는 등 키 쪽 감쇠만 쓴다. 백본 자체는 단일 스케일이고, 다중 스케일 특징은 검출 넥에서 블록 4·8·12의 특징으로 만들어지므로 다른 과제에서는 이 부분을 다시 설계해야 한다. 단안 깊이 추정 결과는 NYU 전용 미세조정 없이 측정·의사 라벨 8개 소스 혼합으로 학습했고, 다중 뷰 테스트타임 증강과 이미지별 log-affine 정렬을 적용한 값이다.