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

2024-06-25

AI Summary

Pre-trained models struggle to adapt efficiently to diverse deployment sizes (e.g., varying depth or width), necessitating size-specific initialization strategies. Method: This paper proposes a multi-task-inspired framework for variable-size model initialization. Its core innovation is a size-agnostic shared weight template—constructed via Learngene-based knowledge distillation with Kronecker-structured constraints—and lightweight learnable scalers that enable consistent cross-size initialization. The template-scaler co-design supports zero-shot transfer to diverse downstream tasks, while scalers require only minimal data for adaptation. Contribution/Results: Experiments demonstrate state-of-the-art initialization performance across multiple depth- and width-varied architectures. The method significantly enhances few-shot adaptability and cross-task generalization, offering a scalable, data-efficient solution for deploying models of heterogeneous sizes without retraining from scratch.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Large Vision Models
• Search and Optimization: Learning to Search

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

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.

Problem

• Addresses initialization of variable-sized models
• Overcomes limitations of pre-trained model sizes
• Ensures consistent initialization across different model sizes

Innovation

• Uses shared, size-agnostic weight templates
• Employs size-specific weight scalers for initialization
• Integrates pre-trained knowledge via distillation process