Self-Supervised Weight Templates for Scalable Vision Model Initialization

TL;DR

SWEET tackles the problem of deploying vision models across varied architectures without costly full-model pre-training. It learns a size-agnostic weight template under Tucker-based constraints and augments it with width-wise stochastic scaling to ensure cross-width generalization, enabling fast, low-resource initialization for downstream tasks. The approach yields consistent gains across classification, detection, segmentation, and generation, and ablations confirm the importance of the Tucker bottleneck and width-robust training. This framework offers a scalable, task-agnostic pathway to transfer knowledge across model sizes, reducing computational burden while preserving performance across domains.

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 \textsc{classification}, \textsc{detection}, \textsc{segmentation} and \textsc{generation} tasks demonstrate the state-of-the-art performance of SWEET for initializing variable-sized vision models.

Figures (4)

• Figure 1: (a) Traditional pre-training paradigms produce fixed-size, task-specific models, which are difficult to adapt to downstream architectures with varying scales and task requirements. (b) SWEET adopts a constraint-based pre-training paradigm that extracts weight templates under structured constraints in a self-supervised manner, enabling flexible cross-scale and cross-task model initialization and efficient knowledge transfer.
• Figure 2: Overview of SWEET.(a) Constraint-based Pre-training of the weight template $\mathcal{G}$ with weight scalers $(U, V, X)$, where Tucker- and low-rank constraints are applied to condense size-agnostic knowledge within the template. A self-supervised objective guides pre-training to learn generalizable visual representations, enhancing the template’s universality across diverse vision tasks. (b) Width-wise Stochastic Scaling randomly masks weight scalers along the width during pre-training, discouraging overfitting to a specific width and promoting the organization of template knowledge for flexible adaptation across models of varying widths.
• Figure 3: Selected visualizations of Object Detection and Semantic Segmentation for SWEET-initialized models.
• Figure 4: Reconstructions of ImageNet validation images using SWEET pre-trained weight templates with a masking ratio of 75%.