Entraînement efficace avec des poids neuronaux débruités

Efficient Training with Denoised Neural Weights

July 16, 2024

Auteurs: Yifan Gong, Zheng Zhan, Yanyu Li, Yerlan Idelbayev, Andrey Zharkov, Kfir Aberman, Sergey Tulyakov, Yanzhi Wang, Jian Ren

cs.AI

Résumé

Une bonne initialisation des poids constitue une mesure efficace pour réduire le coût d'entraînement d'un modèle de réseau de neurones profond (DNN). Le choix de la manière d'initialiser les paramètres est complexe et peut nécessiter un réglage manuel, ce qui peut être chronophage et sujet à des erreurs humaines. Pour surmonter ces limitations, ce travail propose une approche novatrice en construisant un générateur de poids pour synthétiser les poids neuronaux destinés à l'initialisation. Nous utilisons la tâche de traduction d'image à image avec des réseaux antagonistes génératifs (GANs) comme exemple, en raison de la facilité à collecter des poids de modèles couvrant une large gamme. Plus précisément, nous collectons d'abord un ensemble de données contenant divers concepts d'édition d'images et leurs poids entraînés correspondants, qui sont ensuite utilisés pour l'entraînement du générateur de poids. Pour prendre en compte les différentes caractéristiques entre les couches et le nombre important de poids à prédire, nous divisons les poids en blocs de taille égale et attribuons un index à chaque bloc. Par la suite, un modèle de diffusion est entraîné avec cet ensemble de données en utilisant à la fois les conditions textuelles du concept et les index des blocs. En initialisant le modèle de traduction d'image avec les poids débruités prédits par notre modèle de diffusion, l'entraînement ne nécessite que 43,3 secondes. Par rapport à un entraînement à partir de zéro (c'est-à-dire Pix2pix), nous obtenons une accélération du temps d'entraînement de 15x pour un nouveau concept tout en atteignant une meilleure qualité de génération d'images.

English

Good weight initialization serves as an effective measure to reduce the training cost of a deep neural network (DNN) model. The choice of how to initialize parameters is challenging and may require manual tuning, which can be time-consuming and prone to human error. To overcome such limitations, this work takes a novel step towards building a weight generator to synthesize the neural weights for initialization. We use the image-to-image translation task with generative adversarial networks (GANs) as an example due to the ease of collecting model weights spanning a wide range. Specifically, we first collect a dataset with various image editing concepts and their corresponding trained weights, which are later used for the training of the weight generator. To address the different characteristics among layers and the substantial number of weights to be predicted, we divide the weights into equal-sized blocks and assign each block an index. Subsequently, a diffusion model is trained with such a dataset using both text conditions of the concept and the block indexes. By initializing the image translation model with the denoised weights predicted by our diffusion model, the training requires only 43.3 seconds. Compared to training from scratch (i.e., Pix2pix), we achieve a 15x training time acceleration for a new concept while obtaining even better image generation quality.