Vision Transformers (ViTs) achieve state-of-the-art segmentation accuracy but require large training datasets because each layer has unique parameters that must be learned independently. We present RD-ViT, a Recurrent-Depth Vision Transformer that adapts the Recurrent-Depth Transformer (RDT) architecture to dense prediction tasks, supporting both 2D and 3D inputs. RD-ViT replaces the deep stack of unique transformer blocks with a single shared block looped T times, augmented with LTI-stable state injection for guaranteed convergence, Adaptive Computation Time (ACT) for spatial compute allocation, depth-wise LoRA adaptation, and optional Mixture-of-Experts (MoE) feed-forward networks for category-specific specialization. We evaluate on the ACDC cardiac MRI segmentation benchmark in both 2D slice-level and 3D volumetric settings with exclusively real experiments executed in Google Colab. In 2D, RD-ViT outperforms standard ViT at 10% training data (Dice 0.774 vs 0.762) and at full data (0.882 vs 0.872). In 3D, RD-ViT with MoE achieves Dice 0.812 with 3.0M parameters, reaching 99.4% of standard ViT performance (0.817) at 53% of the parameter count. MoE expert utilization analysis reveals that different experts spontaneously specialize for different cardiac structures (RV, MYO, LV) without explicit routing supervision. ACT halting maps show higher compute allocation at cardiac boundaries, and the mean ponder time decreases from 2.6 to 1.4 iterations during training, demonstrating learned computational efficiency. Depth extrapolation enables inference with more loops than training without degradation. All code, notebooks, and results are publicly released.
翻译:[translated abstract in Chinese]
视觉Transformer(ViT)在语义分割任务中达到了最先进的精度,但由于每层具有独立学习的独特参数,其训练需要大规模数据集。我们提出RD-ViT——一种递归深度视觉Transformer,将递归深度Transformer(RDT)架构适配至稠密预测任务,支持2D与3D输入。RD-ViT用单一共享模块循环T次替代原有的深层独特Transformer模块堆叠,并融入LTI稳定状态注入以保证收敛性、自适应计算时间(ACT)实现空间计算分配、深度级LoRA适配以及可选的混合专家(MoE)前馈网络实现类别特定专业化。我们在ACDC心脏MRI分割基准上分别于2D切片级和3D体素级场景进行评估,所有实验均在Google Colab中实际执行。在2D任务中,RD-ViT在10%训练数据下以Dice系数0.774优于标准ViT的0.762,在全数据下以0.882优于0.872。在3D任务中,使用MoE的RD-ViT以3.0M参数量达到Dice 0.812,以标准ViT 53%的参数量实现了其99.4%的性能(0.817)。MoE专家利用率分析表明,不同专家在没有显式路由监督的情况下自发针对不同心脏结构(右心室、心肌、左心室)实现专业化。ACT中止图显示心脏边界处计算分配更高,且训练过程中平均驻留时间从2.6次迭代降至1.4次迭代,展现了学习到的计算效率。深度外推能力使推理时可使用比训练更多的循环次数且无性能退化。所有代码、笔记与结果均已公开开源。