Self-Supervised Weight Templates for Scalable Vision Model Initialization
Abstract
The increasing scale and complexity of modern model parameters underscore the importance of pre-trained models. However, deployment often demands architectures of varying sizes, exposing limitations of conventional pre-training and fine-tuning. To address this, we propose SWEET, a self-supervised framework that performs constraint-based pre-training to enable scalable initialization in vision tasks. Instead of pre-training a fixed-size model, we learn a shared weight template and size-specific weight scalers under Tucker-based factorization, which promotes modularity and supports flexible adaptation to architectures with varying depths and widths. Target models are subsequently initialized by composing and reweighting the template through lightweight weight scalers, whose parameters can be efficiently learned from minimal training data. To further enhance flexibility in width expansion, we introduce width-wise stochastic scaling, which regularizes the template along width-related dimensions and encourages robust, width-invariant representations for improved cross-width generalization. Extensive experiments on classification, detection, segmentation and generation tasks demonstrate the state-of-the-art performance of SWEET for initializing variable-sized vision models.
Lay Summary
Modern deep learning models have become increasingly large and expensive, making it difficult to deploy a single model across different devices and resource constraints. In practice, we often need models of different sizes depending on the application, but standard training methods do not easily support this flexibility. In this work, we propose SWEET, a method that learns a reusable set of model weights that can be adapted to different network sizes. Instead of training separate models from scratch, SWEET learns a shared weight template together with simple scaling components that adjust it to different architectures. This allows the same learned knowledge to be efficiently reused across models with different depths and widths.