ノイズ除去されたニューラル重みを用いた効率的な訓練

Efficient Training with Denoised Neural Weights

July 16, 2024

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

cs.AI

要旨

適切な重み初期化は、ディープニューラルネットワーク（DNN）モデルの学習コストを削減するための効果的な手段です。パラメータの初期化方法の選択は難しく、手動での調整が必要となる場合があり、時間がかかる上に人的ミスが発生しやすいという課題があります。このような制限を克服するため、本研究ではニューラルネットワークの重みを初期化するための重み生成器を構築するという新たなアプローチを提案します。広範なモデル重みを容易に収集できることから、生成敵対ネットワーク（GAN）を用いた画像変換タスクを例として取り上げます。具体的には、まずさまざまな画像編集コンセプトとそれに対応する学習済み重みを含むデータセットを収集し、これを重み生成器の学習に使用します。層ごとの異なる特性と予測すべき重みの膨大な数を考慮し、重みを等サイズのブロックに分割し、各ブロックにインデックスを割り当てます。その後、コンセプトのテキスト条件とブロックインデックスを用いて、このデータセットで拡散モデルを学習させます。提案する拡散モデルによって予測されたノイズ除去済み重みで画像変換モデルを初期化することで、学習時間はわずか43.3秒で済みます。スクラッチからの学習（例：Pix2pix）と比較すると、新しいコンセプトに対して15倍の学習時間短縮を実現しつつ、さらに優れた画像生成品質を達成しています。

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.