WAVE: Weight Templates for Adaptive Initialization of Variable-sized Models

TL;DR

WAVE introduces a novel method for initializing variable-sized models using shared weight templates and size-specific scalers, enabling efficient, adaptable, and transfer-friendly model deployment across different sizes and tasks.

Contribution

WAVE reformulates variable-sized model initialization as a multi-task problem using shared templates and scalers, incorporating knowledge distillation for consistent, adaptable, and transferable initialization.

Findings

Achieves state-of-the-art performance in model initialization across various sizes.

Templates encapsulate task-agnostic knowledge transferable across datasets.

Efficient initialization with minimal training data.

Abstract

The growing complexity of model parameters underscores the significance of pre-trained models. However, deployment constraints often necessitate models of varying sizes, exposing limitations in the conventional pre-training and fine-tuning paradigm, particularly when target model sizes are incompatible with pre-trained ones. To address this challenge, we propose WAVE, a novel approach that reformulates variable-sized model initialization from a multi-task perspective, where initializing each model size is treated as a distinct task. WAVE employs shared, size-agnostic weight templates alongside size-specific weight scalers to achieve consistent initialization across various model sizes. These weight templates, constructed within the Learngene framework, integrate knowledge from pre-trained models through a distillation process constrained by Kronecker-based rules. Target models are then initialized by concatenating and weighting these templates, with adaptive connection rules established by lightweight weight scalers, whose parameters are learned from minimal training data. Extensive experiments demonstrate the efficiency of WAVE, achieving state-of-the-art performance in initializing models of various depth and width. The knowledge encapsulated in weight templates is also task-agnostic, allowing for seamless transfer across diverse downstream datasets. Code will be made available at https://github.com/fu-feng/WAVE.