Efficient Training with Denoised Neural Weights

TL;DR

The paper addresses the inefficiency of weight initialization in training deep networks by proposing a diffusion-based weight generator that synthesizes initialized weights for unseen concepts. It uses image-to-image translation with GANs as a testbed, leveraging LoRA to reduce parameter count and a block-wise weight representation conditioned on text and block indices, enabling fast one-step weight generation and subsequent fine-tuning. The authors build a large ground-truth dataset of LoRA weights via diffusion-edited paired images, train a UNet-based 1D diffusion model to predict weight blocks, and achieve about 15x faster training with improved FID compared to training from scratch and other efficient GAN methods. This approach offers a general framework for rapid, cross-task weight initialization that could substantially lower training costs across diverse architectures.

Abstract

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.

Figures (11)

• Figure 1: The framework overview of our weight generator design. The standard diffusion process turns an image into noise in the forward pass and reverses a clean image from pure noise in the reverse process. Our weight generator is designed to turn a noise to weight initializations for efficient training purposes. Given the text information and block index, the weight generator provides the corresponding weight values.
• Figure 2: The UNet Weight Generator. The weight generator is composed of 1-d ResBlocks and 1-d Transformer blocks. The block embedding $emb_n$ is combined with the time step embedding $emd_t$ to be leveraged in each ResBlock.
• Figure 3: Qualitative comparisons across different concept domains. The leftmost column shows two original images and the remaining columns present the corresponding synthesized images in the target concept domain, where target prompts are shown at the top/bottom row. We provide images generated by various models.
• Figure 4: The FID performance comparison between our method and baseline methods along with the training process on the test dataset for different concepts/styles.
• Figure A1: Examples of collected text prompts of concepts/styles.