Two-Fold Patch Perturbation for Efficient Self-Supervised Learning in 3D Medical Imaging

1 Introduction

Self-supervised learning (SSL) has become a key strategy for reducing annotation costs in 3D medical imaging, where dense voxel-level labels are expensive and scarce. This is relevant especially for volumetric segmentation and classification, where model performance typically scales with both data size and label quality, yet curated annotations remain a major bottleneck.

Prior works explore proxy objectives that encourage the development of transferable volumetric representations without labels. Broadly, these objectives emphasize either appearance-based reconstruction, such as masked modeling [Xie et al., 2022b] and reconstruction, or geometric pretexts, such as rotation [Chen et al., 2023], jigsaw [Chen et al., 2021a], and Rubik-style transformations [Tao et al., 2020], that induce spatial reasoning and orientation sensitivity. In parallel, recent methods also exploit volumetric priors such as anatomical layout and contextual position to inject higherlevel structure into SSL representations [Wu et al., 2024]. Despite strong progress, these approaches are often developed in isolation, and their combination is typically ad hoc or entails substantial computational cost.

Designing SSL for 3D volumes also introduces practical constraints. Compared to 2D images, volumetric inputs are high-dimensional and memory-intensive. This makes it difficult to scale sophisticated pretext pipelines. At the same time, reconstruction-only objectives can under-emphasize geometric consistency. Also, transformation-prediction objectives may become unreliable when local appearance is ambiguous, such as near-isotropic texture, and the applied transformation cannot be inferred without a broader spatial context. These observations motivate a simple question: can we unify appearance and geometry supervision in a way that remains explicit, tunable, and computationally lightweight?

We instead view augmentation as an explicit perturbation process, where each component can be controlled and ablated. For volumetric data, patch-wise perturbation is particularly appealing. As it is natural for 3D inputs, it aligns with transformer-style tokenization and provides a direct handle on how much local information is removed or perturbed. If perturbation is explicitly defined, the resulting SSL task can be systematically harder or easier, and the model can be supervised not only to reconstruct, but also to reason about where perturbation occurred and what was applied.

In this work, we propose a simple and principled SSL framework that unifies appearance and geometric supervision through a two-fold perturbation process. Specifically, we model perturbation as the composition of two distributions: (i) a Bernoulli selection over 3D patches and (ii) a discrete rotation applied only to the selected patches. This formulation is controlled by the mask (selection) rate p, the rotation set size K, and the patch grid resolution n, allowing the perturbation strength to be systematically varied and ablated from weak to strong.

To exploit this structured perturbation, we introduce a three-head pre-training objective that jointly supervises reconstruction, patch selection, and rotation prediction. Beyond reconstruction alone, predicting where perturbations occur and how they are applied encourages the encoder to integrate local and contextual information, yielding spatially aware representations. This is particularly beneficial for medical volumetric data, with its rich anatomical semantics and localized information. Experiments across diverse segmentation and classification benchmarks, including modality-shift scenarios, show consistent gains over general SSL baselines and competitive or superior performance to recent medical SSL methods, while substantially reducing memory, computation, and training time compared to heavier state-of-the-art pre-training pipelines.

Contributions. (1) We propose a simple and tunable two-fold patch perturbation algorithm for volumetric SSL, parameterized by (n, p,K) to control pre-training difficulty systematically; (2) we introduce a lightweight three-head objective (reconstruction, mask prediction, rotation prediction) that promotes spatially aware and transferable representations with minimal additional decoder capacity; and (3) we provide extensive evaluation on diverse medical imaging benchmarks (segmentation and classification), including cross-modality transfer, demonstrating strong downstream performance while substantially reducing pre-training memory, FLOPs, and wall-clock time compared to recent state-ofthe-art SSL methods.

2 RelatedWork

Self-supervised learning for 3D medical imaging. Most SSL methods for volumetric medical images can be broadly grouped into two families relevant to this work: (i) masked reconstruction (masked modeling) and (ii) geometric/spatial pretext tasks. Our method draws on both directions while emphasizing an explicit, tunable perturbation process.

(i) Masked reconstruction and modeling. Masked reconstruction is a dominant paradigm for learning 3D representations without labels, where models reconstruct missing or perturbed voxels from partial observations, as in MAEstyle pre-training and its medical variants [He et al., 2022; Chen et al., 2023; Zhuang et al., 2025]. GL-MAE, in particular, tailors masked autoencoding to volumetric data and highlights the role of global context under masking. While effective, reconstruction-centric objectives primarily target appearance recovery and may provide weaker supervision for geometric consistency.

(ii) Geometric and spatial pretext tasks. A complementary line of work [Taleb et al., 2020; Chen et al., 2021a] uses transformation-prediction objectives, such as rotation classification and patch arrangement, to encourage spatial reasoning and geometric awareness. Rubik-style objectives [Tao et al., 2020] extend this idea by applying discrete permutations and training the network to predict the applied discrete transformation. Although conceptually related, these pretexts typically apply constrained transformations and do not explicitly couple where perturbation occurs with what transformation is applied at the patch level, thus making the proxy task less informative.

Position-aware 3D SSL and relation to VoCo. Recent work exploits anatomical regularities in 3D medical volumes by learning representations sensitive to spatial context. VoCo [Wu et al., 2024] introduces a volume-contrast formulation that leverages contextual position prediction using global-local views and prototypes. In contrast, our approach does not rely on contrastive sampling or prototype construction; instead, it injects spatial and geometric supervision directly through structured patch-wise perturbation and corresponding prediction heads.

Relation to our work. Our approach differs from prior work in how reconstruction and transformation supervision are integrated. Rather than combining objectives in a loosely coupled manner, we define a structured two-fold patch-wise perturbation process that explicitly models both where perturbation occurs and how it is applied. This yields a unified, tunable pre-training objective that encourages spatial awareness and orientation sensitivity, while remaining computationally lightweight compared to contrastive or similar SSL methods.

3 Methodology

3.1 Overview

We aim to learn transferable representations from unlabeled 3D medical volumes using complementary self-supervised pretext tasks. Given a fixed-size ROI volume, we generate a two-fold perturbed view via patch-wise masking and rotation, encode it with a Swin-UNETR-style backbone, and jointly optimize reconstruction, mask localization, and rotation prediction objectives (Figure 1) . The resulting encoder is transferred and finetuned for downstream supervised tasks.

3.2 Input Preprocessing and Notation

After standard preprocessing, each input volume is cropped or padded to a fixed ROI of size 64×64×64. We denote the resulting tensor as

where B is the batch size and C the number of input channels (C = 1 in our experiments).

3.3 Two-Fold Patch-wise Perturbation

We partition the ROI volume into a regular grid of n × n × n non-overlapping cubic patches ("patchification"), yielding P = n3 patches, each of spatial size s×s×s with s = 64/n. Let Sn(·) denote the patchification operator:

and let Cn(·) denote the inverse operation that reassembles patches into a full volume.

Fold 1: Patch selection (masking). We sample a Bernoulli mask over patches:

where pmask is the masking rate. Intuitively, mb,i = 1 indicates that patch i in sample b will be perturbed.

Fold 2: Patch-wise rotation. For each selected patch, we sample a discrete rotation label from a predefined set of K 3D rotation transformations:

where r = 0 denotes the identity (no rotation). Let {Rk}K-1 k=0 be the set of 3D rotation operators implemented by axis permutations and flips (90◦ -multiple rotations without interpolation). We perturb patches as:

and reassemble the perturbed patches into perturbed volume:

This two-fold perturbation (Figure 2) introduces both spatial uncertainty and orientation inconsistency, requiring the encoder to localize perturbed regions and leverage surrounding anatomical context to resolve ambiguities.

3.4 Model Architecture

Our network consists of (i) a Swin-UNETR-style encoder, (ii) a UNETR decoder for volumetric reconstruction, and (iii) two lightweight MLP heads for mask and rotation prediction.

Swin-UNETR-style Encoder

We use a 3D Swin Transformer backbone with hierarchical representations. Denote the encoder by E(·); given the perturbed volume ˜x it produces multi-scale feature maps:

where h(ℓ) corresponds to a spatial scale in the Swin hierarchy. Following the Swin-UNETR wiring, selected transformer features are refined using convolutional blocks to produce multi-resolution encoder feature maps

where e(0) captures low-level appearance and e(4) provides the highest-level spatial representation.

Multi-scale pooled representation. We construct a global representation by applying global average pooling to each encoder scale and concatenating the results:

where dg denotes the concatenated feature dimension. This compact representation is used to drive patch-level prediction heads.

Reconstruction Decoder

We attach a UNETR-style decoder Drec(·), reusing the SwinUNETR decoder pathway that combines the highest-level feature map with skip connections from intermediate encoder resolutions:

Mask and Rotation Heads

We predict (i) which patches were perturbed and (ii) the rotation applied to each perturbed patch, using lightweight MLP heads on the pooled multi-scale representation zglobal:

where ˆm are patch-wise logits for perturbation localization and ˆ r are patch-wise logits over K discrete rotation classes. Using lightweight heads biases representational capacity toward the encoder rather than task-specific decoders.

3.5 Pre-training Objective

We jointly optimize reconstruction, mask prediction, and rotation prediction losses:

with scalar weights λrec, λmask, λrot.

Masked-only Charbonnier Reconstruction Loss

We use the Charbonnier penalty (a smooth, robust variant of ℓ1 [Barron, 2019]):

To align reconstruction learning with our perturbation process, we compute reconstruction error only over voxels in perturbed patches. To do so, we convert the patch mask m ∈ {0, 1}B×P into a voxel-level mask w ∈ {0, 1}B×1×64×64×64 by reshaping m into a n × n × n grid and upsampling with nearest-neighbor interpolation (preserving exact patch boundaries). The masked-only reconstruction loss is:

where v indexes voxels (and channels). This term encourages the model to "undo" local patch perturbation by leveraging the surrounding anatomical context.

Mask Prediction Loss

The mask head predicts which patches were perturbed. We use a class-balanced binary cross-entropy on patch logits:

where α is a positive-class reweighting factor to counter class imbalance induced by sparse masking. In practice, we estimate α per mini-batch, with a safe fallback to unweighted BCE when the number of positives = 0.

Rotation Prediction Loss (Selected-only)

For masked patches, the rotation head predicts the applied rotation label. We compute cross-entropy only over the masked subset Ω = {(b, i) | mb,i = 1}:

where ˆ rb,i ∈ RK are logits over K rotation classes. We employ label smoothing to reduce overconfident predictions and improve the generalization of the learned representation. If |Ω| = 0 for a mini-batch, we set Lrot = 0 for that update.

3.6 Training Procedure

During pre-training, a perturbed view (˜x,m, r) is generated for each unlabeled input x using the two-fold perturbation in Sec. 3.3. The network predicts (ˆx, ˆm, ˆ r), and all parameters are jointly optimized using the combined objective in Eq. (13). The encoder is subsequently finetuned on downstream tasks.

4 Experiments

4.1 Experimental Protocol

We adopt the experimental protocol of recent state-of-the-art methods to ensure a fair and directly comparable evaluation against a broad range of prior self-supervised learning methods. Unless otherwise stated, we follow VoCo for preprocessing, training budget, backbone family, and finetuning recipe, and we compare against the results reported in [Wu et al., 2024; Zhuang et al., 2025], which re-implements most baselines under a unified setup.

Implementation details We use a Swin-UNETR [Hatamizadeh et al., 2021] architecture as the model backbone for both self-supervised pre-training and downstream finetuning, in accordance with the previous baselines. For training, we use AdamW [Loshchilov and Hutter, 2017] optimization and a warmup-cosine learning rate schedule over 500 epochs to maintain consistency with previous methods. The finetuning setup is mostly consistent with the pre-training, differing mainly in the number of training epochs across datasets. We apply slicing-window inference to compare with SOTA methods. We do not use foundation models or post-processing, aiming to evaluate our method's pure effectiveness. Details of the training settings and preprocessing are provided in our GitHub repository.

4.2 Pre-training Data

We pre-train on three publicly available CT datasets: BTCV [Landman et al., 2015], TCIA COVID-19 [Clark et al., 2013], and LUNA [Setio et al., 2017], which totals approximately 1.6k CT scans. For downstream experiments on BTCV [Landman et al., 2015], WORD [Luo et al., 2022], and AMOS22 [Ji et al., 2022], we pre-train only on BTCV and TCIA COVID-19, consistent with the protocol of [Wu et al., 2024; Zhuang et al., 2025; Chen et al., 2023]. For the other 4 downstream tasks on MM-WHS [Zhuang, 2018], MSD-spleen [Antonelli et al., 2022], CC-CCII [Zhang et al., 2020], and BraTS [Baid et al., 2021], we pre-train on the union of all three datasets.

The two-fold perturbation parameters. Unless otherwise specified, we use a patch grid of n × n × n with n = 4 (thus P = n3 = 64 patches), masking rate pmask = 0.3, and a rotation set of size K = 9 for the pre-training task.

4.3 DownstreamTasks and Datasets

We evaluate transfer performance on seven downstream benchmarks spanning segmentation and classification across CT and MRI data. Details on the source, modality, size, and split of the dataset are provided in the supplementary materials in the GitHub repository and align with the standard SOTA methods reported in previous works. We will also provide the JSON files containing the sample lists for each dataset to ensure reproducibility and fair comparison.

Segmentation on CT. We finetune on BTCV, AMOS, WORD, MSD-Spleen, and MM-WHS to benchmark against all medical segmentation baselines, and report the Dice score (%) on the test set.

Segmentation on MRI. To evaluate robustness under modality shift, we finetune on BraTS (MRI data). Since pretraining is conducted exclusively on CT, BraTS constitutes an out-of-distribution (OOD) modality transfer setting.

Classification on CT. We further evaluate on CC-CCII, a CT-based classification benchmark. We report classification accuracy (%) following the evaluation protocol of previous methods.

4.4 Baselines

We compare against three baseline categories, using results reported in [Wu et al., 2024; Zhuang et al., 2025] under a matched setup. Not all methods are reported for every dataset; we follow the available comparisons per benchmark.

Training from scratch. We include randomly initialized baselines using standard 3D architectures, including 3D UNet [Ronneberger et al., 2015], UNETR [Hatamizadeh et al., 2022], and Swin-UNETR [Hatamizadeh et al., 2021].

General-purpose SSL. We compare with representative general SSL methods adapted to 3D data, including MAE3D [He et al., 2022], SimCLR [Chen et al., 2020], SimMIM [Xie et al., 2022b], MoCo v3 [Chen et al., 2021b], Jigsaw [Chen et al., 2021a], Position-Label [Zhang and Gong, 2023], and DINO [Caron et al., 2021].

Medical SSL. We compare with medical SSL methods, including Model Genesis [Zhou et al., 2019], Rotation [Taleb et al., 2020], VICRegK [Bardes et al., 2022], UniMiSS [Xie et al., 2022a], Rubik++ [Tao et al., 2020], PCRLv2 [Zhou et al., 2023], Med-MAE [Chen et al., 2023], SwinMM [Wang et al., 2023], JointSSL [Nguyen et al., 2023], GVSL [He et al., 2023], GLMAE [Zhuang et al., 2025], and VoCo [Wu et al., 2024].

4.5 DownstreamResults

We report downstream transfer performance in Tables 1-5, covering challenging multi-organ CT segmentation (BTCV, AMOS22, WORD), cardiac whole-body CT segmentation (MM-WHS), single-organ CT segmentation (MSDSpleen), MRI brain tumor segmentation under modality shift (BraTS21), and CT-based classification (CC-CCII).

Large gains on AMOS22 and WORD. As shown in Table 1, our method achieves the best results on both AMOS22 and WORD, reaching 89.78 and 84.20 Dice, respectively. These scores substantially exceed the strongest prior baseline in the table (GL-MAE) by +2.64 Dice on AMOS22 (89.78 vs. 87.14) and +4.76 Dice on WORD (84.20 vs. 79.44), indicating that the proposed method's objective transfers effectively to complex multi-organ segmentation.

Robust transfer under CT→MRI modality shift. Table 2 reports BraTS21 results, where pre-training is performed on CT while finetuning is conducted on MRI, making this an out-of-distribution modality transfer setting. Our method achieves the best overall performance with an average Dice of 88.50, improving upon VoCo by +9.97 Dice (88.50 vs. 78.53). Notably, the gain is driven by a large improvement on the enhancing tumor (ET) region (83.56 vs. 59.87), suggesting that the learned representation captures transferable cues that remain useful under substantial appearance shifts.

Improved CT classification on CC-CCII. On CC-CCII classification (Table 3), our method achieves 94.25% accuracy, outperforming VoCo (90.83%) and all other reported SSL baselines. This result indicates that the learned encoder features generalize beyond segmentation and remain effective for image-level decision making.

1 'From Scratch' denotes the supervised baseline

Competitiveness on BTCV organ-wise segmentation. Table 4 shows organ-wise results on BTCV. Our method achieves an average Dice of 83.57, which is close to the bestperforming method (VoCo at 83.85) while surpassing other strong baselines such as GL-MAE (82.01) and SwinMM (81.81). Qualitatively, our approach yields particularly significant improvements on smaller or shape-variable structures such as the gallbladder and the esophagus, while remaining comparable on large organs.

Strong performance on MM-WHS and MSD-Spleen. Table 5 summarizes results on MM-WHS and MSD-Spleen. Our method achieves the best performance on MM-WHS with a Dice score of 90.67, slightly improving over VoCo (90.54). On MSD-Spleen, our method achieves 96.11 Dice, remaining close to VoCo (96.34) and outperforming other SSL baselines. These results suggest that the proposed pretraining is effective both for challenging multi-structure settings and for high-accuracy single-organ segmentation.

Across datasets. The proposed two-fold SSL objective delivers consistent gains over training from scratch and general SSL methods, and it is competitive with, and often superior to, specialized medical SSL pre-training. The greatest improvements are observed on AMOS22, WORD, BraTS21, and CC-CCII, while performance on BTCV and MSD-Spleen remains close to the best reported baseline.

4.6 Computational Efficiency

We evaluate the computational efficiency of the proposed pretraining strategy against the current state-of-the-art baseline, VoCo. To ensure a controlled comparison, we pre-train both methods on a single GPU (AMD Instinct MI300X) each, under identical experimental protocol. We report three practical efficiency indicators: (i) average wall-clock time per epoch, (ii) peak GPU memory allocation, and (iii) forwardpass FLOPs (reported per sample).

As summarized in Table 6, our method is consistently more efficient across all metrics while maintaining strong downstream transfer performance (Sec. 4.5). Specifically, our approach reduces peak GPU memory usage from 53.11GiB to 10.41GiB (an ∼80% reduction), decreases forward FLOPs from 1499.03 to 104.59 GFLOPs per sample (over 14× fewer FLOPs), and shortens average epoch time from 723.19 s to 489.85 s (approximately 1.5× faster). These gains make the proposed two-fold SSL method substantially more practical for large-scale 3D pre-training, enabling faster iteration and deployment under constrained compute budgets.

4.7 Ablation Studies

We conduct ablation studies to analyze the contribution of individual objective components and to assess the effect of the perturbation strength of our proposed method. All ablations are performed using the same pre-training budget and finetuning protocol as the main experiments, ensuring that performance differences are attributable solely to the ablated factors.

Effect of Objective Components

We first examine the role of each loss component by selectively removing the rotation loss Lrot, the mask prediction loss Lmask, and both jointly, while keeping the masked-only reconstruction loss Lrec unchanged. Results are summarized in Table 7.

Removing either auxiliary objective leads to a consistent drop in performance, indicating that both patch localization and rotation prediction provide complementary supervision signals beyond reconstruction alone. The full objective achieves the best results across all reported datasets, confirming that jointly optimizing reconstruction, mask prediction, and rotation classification is critical to learning transferable representations. When both auxiliary losses are removed, performance degrades further, demonstrating that reconstruction alone is insufficient to capture the full benefits of the proposed two-fold pre-training strategy.

Effect of perturbation Strength

Next, we study the effect of perturbation strength by varying the patch grid size n and the masking rate pmask, while keeping the rotation set size fixed at K = 9. Increasing n produces smaller patches and thus higher spatial granularity, whereas increasing pmask increases the fraction of perturbed patches. Both factors control the difficulty of the pretext task.

As shown in Table 8, moderate perturbation strength yields the best downstream performance. Weak perturbations (small n or low pmask) result in suboptimal performance, likely because the pretext task is insufficiently challenging and fails to enforce robust feature learning. Conversely, over-strong perturbations can introduce excessive noise and degrade representation quality. The default setting (n = 4, pmask = 0.3) strikes a favorable balance and is used throughout the main experiments.

Discussion. These ablations confirm that the proposed twofold SSL framework benefits from both auxiliary objectives and appropriately strong perturbation. The combination of masked-only reconstruction with patch-wise localization and rotation prediction provides richer supervision than any of the individual components. Also, carefully balancing perturbation strength is essential for maximizing downstream transfer performance.

5 Conclusion

We proposed a simple and efficient self-supervised pretraining framework for 3D medical imaging based on a twofold patch-wise perturbation strategy that combines masking and discrete rotations. By jointly optimizing reconstruction, patch localization, and rotation prediction, the method learns spatially aware and transferable representations without annotations. Experiments across diverse segmentation and classification benchmarks show consistent gains over training from scratch and general SSL baselines, competitive or superior performance compared to recent medical SSL methods, and robust transfer under modality shift, while substantially reducing memory, computation, and training time relative to the SOTA method. Ablation studies further support the contribution of both auxiliary objectives and the role of balancing perturbation strength. Overall, this work offers a practical, scalable approach to learning 3D medical image representations with limited annotation and compute budget.

The explicit parameterization of perturbation by (n, p,K) provides a transparent handle for adjusting pre-training difficulty and systematically studying robustness. Beyond the tasks evaluated here, the same formulation can be extended to other volumetric modalities and objectives, such as richer transformation sets or additional auxiliary signals, without changing the core training methodology.

Limitations. Results use a fixed training budget. Evaluation with broader hyperparameter sweeps and adaptive loss weighting [Kendall et al., 2018] remains future work.

Acknowledgments

We thank collaborators and domain experts for their helpful feedback. Large language models (LLMs) were used only for minor text refinement and block diagram visualization. All experiments and analyses were conceived and verified by the authors.

Code and Data Availability

The datasets used in this study are publicly available. All preprocessing, training scripts, and JSON files are available in our https://github.com/kabir2505/Two-Fold repository to ensure full reproducibility.

Ethics Statement

This study used only publicly available, de-identified data. No human subjects or patient interaction were involved, and no institutional approval was required.