BertsWin: Resolving Topological Sparsity in 3D Masked Autoencoders via Component-Balanced Structural Optimization

TL;DR

BertsWin addresses the challenge of topological sparsity in 3D masked autoencoders for medical CBCT data by preserving a full 3D token grid and applying a Swin-transformer encoder with a light CNN decoder. The authors introduce MVC-based loss and a physics-informed PhysLoss to balance brightness, contrast, and structure, supplemented by the GradientConductor optimizer that blends LARS and LION with bias correction. Empirical results on temporomandibular joint CBCT demonstrate accelerated convergence (up to 15× fewer epochs) without FLOP penalties and improved latent representations, enabling robust 3D reconstruction and competitive downstream segmentation through linear/probing and non-linear decoders. The work underlines the importance of architectural topology, domain-specific loss design, and optimizer choices in achieving efficient, high-quality volumetric representation learning for specialized medical tasks with potential broader applicability.

Abstract

The application of self-supervised learning (SSL) and Vision Transformers (ViTs) approaches demonstrates promising results in the field of 2D medical imaging, but the use of these methods on 3D volumetric images is fraught with difficulties. Standard Masked Autoencoders (MAE), which are state-of-the-art solution for 2D, have a hard time capturing three-dimensional spatial relationships, especially when 75% of tokens are discarded during pre-training. We propose BertsWin, a hybrid architecture combining full BERT-style token masking using Swin Transformer windows, to enhance spatial context learning in 3D during SSL pre-training. Unlike the classic MAE, which processes only visible areas, BertsWin introduces a complete 3D grid of tokens (masked and visible), preserving the spatial topology. And to smooth out the quadratic complexity of ViT, single-level local Swin windows are used. We introduce a structural priority loss function and evaluate the results of cone beam computed tomography of the temporomandibular joints. The subsequent assessment includes TMJ segmentation on 3D CT scans. We demonstrate that the BertsWin architecture, by maintaining a complete three-dimensional spatial topology, inherently accelerates semantic convergence by a factor of 5.8x compared to standard ViT-MAE baselines. Furthermore, when coupled with our proposed GradientConductor optimizer, the full BertsWin framework achieves a 15-fold reduction in training epochs (44 vs 660) required to reach state-of-the-art reconstruction fidelity. Analysis reveals that BertsWin achieves this acceleration without the computational penalty typically associated with dense volumetric processing. At canonical input resolutions, the architecture maintains theoretical FLOP parity with sparse ViT baselines, resulting in a significant net reduction in total computational resources due to faster convergence.

Figures (13)

• Figure 1: The architecture of BertsWin.
• Figure 2: Loss Function Components and Rationale for Multi-Component Variance (MVC) Regularization. Our gradient analysis during the MAE training process on head and neck CT images confirmed an inherent conflict within the $L2$-loss. The optimization process consistently prioritizes the minimization of the mean intensity component (Brightness Error).
• Figure 3: Validation Convergence Dynamics and Optimization Efficiency. Comparison of BertsWin models with various loss functions and optimizers (PhysLoss GCond, L2 GCond, L2 AdamW) and the basic MAE ViT model (L2 AdamW). (a) Dynamics of the MVC metric for soft tissues; (b) MVC metric for surfaces; (c) Dynamics of L2 losses on the validation set. The x-axis shows the number of learning steps, and the y-axis shows the metric values on a logarithmic scale.
• Figure 4: Comparative Density of Latent Representations. Visual analysis of feature magnitude distributions. Random networks (dashed) display high-entropy Gaussian-like profiles. Solid curves correspond to trained models, while dashed curves denote randomly initialized networks.
• Figure 5: Feature Space Discriminability and Invariance Analysis. Violin plots visualize the density distribution of cosine similarity scores across three evaluation protocols. The internal box plots represent the median and interquartile range (IQR). Inter-Patient serves as the baseline noise floor (lower is better), while Intra-Patient and Geometric Invariance represent signal preservation and robustness (higher is better).