[Paper Review] Adapting Pre-trained Vision Transformers from 2D to 3D through Weight Inflation Improves Medical Image Segmentation

Yuhui Zhang, Shih-Cheng Huang+3 more|arXiv (Cornell University)|Feb 8, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes a simple yet effective weight inflation strategy to adapt pre-trained 2D Vision Transformers for 3D medical image segmentation, enabling transfer learning while preserving depth information. By inflating 2D weights to 3D via centering initialization and using small 3D windows with center-slice prediction, the method achieves state-of-the-art performance across 12 diverse 3D medical imaging datasets with only a 0.75% increase in computational cost.

ABSTRACT

Given the prevalence of 3D medical imaging technologies such as MRI and CT that are widely used in diagnosing and treating diverse diseases, 3D segmentation is one of the fundamental tasks of medical image analysis. Recently, Transformer-based models have started to achieve state-of-the-art performances across many vision tasks, through pre-training on large-scale natural image benchmark datasets. While works on medical image analysis have also begun to explore Transformer-based models, there is currently no optimal strategy to effectively leverage pre-trained Transformers, primarily due to the difference in dimensionality between 2D natural images and 3D medical images. Existing solutions either split 3D images into 2D slices and predict each slice independently, thereby losing crucial depth-wise information, or modify the Transformer architecture to support 3D inputs without leveraging pre-trained weights. In this work, we use a simple yet effective weight inflation strategy to adapt pre-trained Transformers from 2D to 3D, retaining the benefit of both transfer learning and depth information. We further investigate the effectiveness of transfer from different pre-training sources and objectives. Our approach achieves state-of-the-art performances across a broad range of 3D medical image datasets, and can become a standard strategy easily utilized by all work on Transformer-based models for 3D medical images, to maximize performance.

Motivation & Objective

• To address the challenge of adapting pre-trained 2D Vision Transformers to 3D medical image segmentation while preserving transfer learning benefits and depth information.
• To investigate whether weight inflation can effectively bridge the domain gap between 2D natural images and 3D medical scans.
• To establish a generalizable, standardizable strategy for leveraging pre-trained 2D models in 3D medical vision tasks.
• To evaluate the impact of different pre-training sources and objectives on 3D segmentation performance.
• To demonstrate consistent performance gains across diverse 3D medical imaging datasets with minimal hyperparameter tuning.

Proposed method

• Adapting pre-trained 2D Vision Transformers to 3D by inflating their learned weights using a centering strategy, where the center-most slice's weights are copied and others are initialized to zero.
• Parsing 3D input volumes into small 3D windows containing a few neighboring slices to preserve spatial context.
• Training the model to predict segmentation only for the center slice within each window, mimicking residual learning and enabling gradual depth integration.
• Using a standard 3D U-Net-style architecture with the inflated 2D Transformer as the encoder for end-to-end segmentation.
• Applying ablation studies to compare different inflation strategies, including prediction target, weight initialization, stride, and number of input slices.
• Evaluating generalization across 11 additional 3D medical imaging datasets with varying anatomical regions and imaging modalities (CT/MRI).

Figure 1: Approach overview. Large-scale pre-trained Transformers are used as the encoder in the segmentation model for transfer learning, in which weights are adapted using the inflation strategy to support 3D inputs. Each 3D image is split into windows, which contain a small number of neighbor sli

Experimental results

Research questions

• RQ1 Can weight inflation effectively adapt pre-trained 2D Vision Transformers to 3D medical image segmentation without architectural modification?
• RQ2 How does the choice of pre-training source (self-supervised vs. supervised) affect transfer performance in 3D medical segmentation?
• RQ3 What is the optimal window size and prediction strategy (e.g., center-slice only) for 3D adaptation using weight inflation?
• RQ4 Does the proposed method generalize across diverse 3D medical imaging datasets with different anatomical structures and imaging protocols?
• RQ5 To what extent does weight inflation improve performance compared to standard 2D or 3D random initialization?

Key findings

• The proposed weight inflation strategy improves mean DSC by 1.95% over strong baselines with only a 0.75% increase in computational cost.
• The model initialized with pre-trained weights outperforms a randomly initialized 3D model by 11.18% on the multi-organ segmentation dataset, highlighting the importance of transfer learning.
• The centering inflation strategy—copying 2D weights to the center slice and zero-initializing others—achieves the best performance, suggesting gradual depth integration is effective.
• Including only a few neighboring slices (e.g., 3–5) is sufficient for performance, as adjacent slices provide more relevant context than distant ones.
• The method achieves state-of-the-art performance on 12 datasets, including 10 out of 11 additional datasets beyond the primary evaluation, demonstrating strong generalization.
• Supervised pre-training sources yield better transfer performance than self-supervised ones, indicating that label-rich pre-training is more beneficial for 3D medical segmentation.

Figure 2: Transfer learning reduces noise in local predictions, and depth information further improves prediction consistency across slices. We show five consecutive input images, ground-truth labels, predictions from Ours , Ours w/o D , Ours w/o T&D .