[論文レビュー] Vision Tiny Recursion Model (ViTRM): Parameter-Efficient Image Classification via Recursive State Refinement
ViTRM は tiny な共有重み再帰的エンコーダを用いて、 latent memory と prediction token を深い supervision と halting 機構で反復的に refine することで、ViTs や CNNs よりはるかに少ないパラメータで CIFAR-10/100 において競争力のある画像分類精度を達成する。
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 \textbf{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 \times $ and $84 \times$ 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.
研究の動機と目的
- vision における深層的なアーキテクチャスケーリングの parameter-efficient な代替を提案する。
- 共有重みを用いた tiny な再帰的エンコーダ ViTRM を提案し、深い ViT エンコーダを置換する。
- CIFAR-10/100 で substantially fewer parameters で競争的な精度を示す。
- supervision の深さと latent reasoning の深さ が性能と学習ダイナミクスに与える影響を分析する。
提案手法
- ViT スタイルのパッチ埋め込みで画像をパッチトークンへ埋め込む(固定画像特徴を使用)。
- 二つの再帰状態を維持する: latent memory z(K トークン)と prediction token y(単一ベクトル)。
- tiny な共有 Transformers ブロックを再帰的に適用し、z を refine しつつ T ステップで y を更新する。z には内側の M イテレーションを使用。
- 各ステップで y からクラスロジットと halting 確率 q を読み出し、軽量ヘッドで early stopping を可能にする。
- 複数ステップで深層 supervision を用いて訓練し、各ステップ後に重みを更新、 halting 条件を満たすか固定ステップ後に訓練を停止。
実験結果
リサーチクエスチョン
- RQ1 再帰的で重みを共有する計算は、視覚タスクにおけるアーキテクチャ的深さを置き換え、精度を保てるか。
- RQ2 supervision の深さと latent reasoning の深さ が最終性能と訓練安定性に与える影響はどうか。
- RQ3 パラメータ効率とミニバッチサイズに対する頑健性の点で、 ViTRM は ViT や ResNet のベースラインとどう比較されるか。
- RQ4 latent state から image tokens へのクロスアテンションは、純粋なフィードフォワードの latent refinement と比較して視覚タスクに有益か。
主な発見
| Model | Params | 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.6M–3.7M で競争力の Top-1 精度を達成する。
- CIFAR-10 では、最適バッチサイズで 93.1% に達し、4.7 倍少ないパラメータで ViT-Small の 94% に近く、ViT-Base の 23 倍も少ない。
- CIFAR-100 では、ViTRM(72.1%)は ViT-Small(76.1%)との差を縮め、小さなバッチサイズで ViT-Base に匹敵し、はるかに少ないパラメータを使用。
- ViTRM はすべての ResNet 系統を上回り、特に CIFAR-100 での利得が大きく、例:72.1% 対 63.8%(ResNet-50)」。
- バッチサイズに対する頑健性が ViTRM に観測され、ResNet 系より劣化が小さい。最適バッチサイズは約 128。
- アブレーションにより最小の supervision depth(N_supervision=1)で最良の性能となり、中程度の latent reasoning depth(n_latent_steps≈3)が精度と計算の良いトレードオフを提供する。
より良い研究を、今すぐ始めましょう
論文設計から論文執筆まで、研究時間を劇的に削減しましょう。
クレジットカード登録不要
このレビューはAIが作成し、人間の編集者が確認しました。