[논문 리뷰] Vision Tiny Recursion Model (ViTRM): Parameter-Efficient Image Classification via Recursive State Refinement

Ange-Clément Akazan, Abdoulaye Koroko외 4명|arXiv (Cornell University)|2026. 03. 19.

Advanced Neural Network Applications인용 수 0

한 줄 요약

ViTRM은 공유 가중치를 가진 아주 작은 재귀 인코더가 CIFAR-10/100에서 ViTs나 CNN에 비해 훨씬 적은 매개변수로도 경쟁력 있는 이미지 분류 정확도를 달성할 수 있음을 보여줍니다. 잠재 메모리와 예측 토큰을 깊은 감독과 정지 메커니즘으로 반복적으로 정제함으로써.

ABSTRACT

The success of deep learning in computer vision has been driven by models of increasing scale, from deep Convolutional Neural Networks (CNN) to large Vision Transformers (ViT). While effective, these architectures are parameter-intensive and demand significant computational resources, limiting deployment in resource-constrained environments. Inspired by Tiny Recursive Models (TRM), which show that small recursive networks can solve complex reasoning tasks through iterative state refinement, we introduce the extbf{Vision Tiny Recursion Model (ViTRM)}: a parameter-efficient architecture that replaces the $L$-layer ViT encoder with a single tiny $k$-layer block ($k{=}3$) applied recursively $N$ times. Despite using up to $6 imes $ and $84 imes$ fewer parameters than CNN based models and ViT respectively, ViTRM maintains competitive performance on CIFAR-10 and CIFAR-100. This demonstrates that recursive computation is a viable, parameter-efficient alternative to architectural depth in vision.

연구 동기 및 목표

• 비전에서의 딥 아키텍처 확장에 대한 매개변수 효율적 대안을 제시한다.
• 공유 가중치를 갖는 작은 재귀 인코더 ViTRM을 제안하여 깊은 ViT 인코더를 대체한다.
• CIFAR-10/100에서 매개변수가 크게 줄어든 상태에서 경쟁력 있는 정확도를 보여준다.
• 감독 깊이와 잠재 추론 깊이가 성능 및 학습 역학에 미치는 영향을 분석한다.

제안 방법

• 고정된 이미지 특징을 가진 ViT 스타일 패치 임베딩을 사용하여 이미지를 패치 토큰으로 임베딩한다.
• 잠재 메모리 z(K 토큰)와 예측 토큰 y(단일 벡터)라는 두 개의 재귀 상태를 유지한다.
• 내부 M 반복으로 z를 정제하고 T 스텝에 걸쳐 y를 업데이트하기 위해 아주 작은 공유 Transformer 블록을 재귀적으로 적용한다.
• 가볍게 구성된 헤드를 통해 y에서 각 단계의 클래스 로짓과 정지 확률 q를 읽어 들여 조기 중지 가능하게 한다.
• 여러 스텝에서 깊은 감독으로 학습하고 각 스텝 이후 가중치를 업데이트하며, 정지 조건이 충족되거나 고정 스텝 이후에 학습을 중지한다.

실험 결과

연구 질문

• RQ1 재귀적이고 가중치 공유된 계산이 시각 작업의 아키텍처 깊이를 대체하면서 정확도를 유지할 수 있을까?
• RQ2 감독 깊이와 잠재적 추론 깊이가 최종 성능과 학습 안정성에 어떤 영향을 미치는가?
• RQ3 ViTRM이 CIFAR-10에서 매개변수 효율성과 배치 크기 강건성 면에서 ViT 및 ResNet 베이스라인과 어떤 차이를 보이는가?
• RQ4 잠재 상태에서 이미지 토큰으로 가는 교차 주의가 순수하게 피드포워드 잠재 정제보다 시각 작업에 유용한가?

주요 결과

모델	매개변수	BS=64	BS=128	BS=256	BS=512	BS=1024
ViTRM (ours)	3.6M	92.6	93.1	92.8	92.1	91.2
ViT-Small	16.9M	94.0	93.9	93.3	92.8	91.1
ViT-Base	85.1M	93.7	93.8	94.3	94.1	94.0
ViT-Large	302.4M	80.8	73.4	82.0	83.2	OOM
ResNet-18	11.2M	90.5	89.7	88.8	87.8	86.1
ResNet-34	21.3M	90.8	90.1	89.7	87.9	86.3
ResNet-50	23.5M	92.0	90.9	89.7	87.9	84.2
ViTRM (ours) - CIFAR-100	3.7M	71.8	72.1	72.0	71.3	69.8
ViT-Small - CIFAR-100	16.9M	76.1	75.2	74.8	72.7	71.2
ViT-Base - CIFAR-100	85.2M	72.4	74.6	75.4	76.4	75.1
ViT-Large - CIFAR-100	302.5M	61.1	58.9	61.7	64.7	OOM
ResNet-18 - CIFAR-100	11.2M	64.6	63.4	62.1	60.5	58.2
ResNet-34 - CIFAR-100	21.3M	65.4	63.8	62.0	58.8	56.7
ResNet-50 - CIFAR-100	23.7M	63.8	63.4	61.5	58.0	54.5

• ViTRM은 CIFAR-10/100에서 ViT 변형 및 ResNet보다 훨씬 적은 매개변수 3.6–3.7M으로 경쟁력 있는 Top-1 정확도를 달성한다.
• CIFAR-10에서 ViTRM은 최적 배치 크기에서 93.1%에 도달하며 ViT-Small(94%)와 근접하나 매개변수는 4.7x 더 적고 ViT-Base보다 훨씬 작다.
• CIFAR-100에서 ViTRM(72.1%)은 ViT-Small(76.1%)와의 격차를 좁히고 소형 배치 크기에서 ViT-Base와 일치하며, 매개변수는 훨씬 적다.
• ViTRM은 모든 데이터셋과 배치 크기에서 ResNet 변형들보다 우수하며, CIFAR-100에서 특히 큰 이점이 있다(예: 72.1% vs. 63.8% for ResNet-50).
• 배치 크기에 대한 견고성이 ViTRM에서 관찰되며 배치 크기가 증가함에 따라 ResNet 변형들보다 감소가 작다; 최적 배치 크기는 약 128.
• 적용 제거 연구에 따르면 최소 감독 깊이(N_supervision=1)에서 가장 좋은 성능이 나오고, 중간 잠재 추론 깊이(n_latent_steps≈3)가 정확도와 계산의 최적 균형을 제공한다.