Deep Fusion: Efficient Network Training via Pre-trained Initializations

TL;DR

This work presents Deep Fusion, an efficient approach to network training that leverages pre-trained initializations of smaller networks, and proposes a theoretical framework using backward error analysis to illustrate the dynamics of mid-training network growth.

Abstract

In recent years, deep learning has made remarkable progress in a wide range of domains, with a particularly notable impact on natural language processing tasks. One of the challenges associated with training deep neural networks in the context of LLMs is the need for large amounts of computational resources and time. To mitigate this, network growing algorithms offer potential cost savings, but their underlying mechanisms are poorly understood. We present two notable contributions in this paper. First, we present Deep Fusion, an efficient approach to network training that leverages pre-trained initializations of smaller networks. Second, we propose a theoretical framework using backward error analysis to illustrate the dynamics of mid-training network growth. Our experiments show how Deep Fusion is a practical and effective approach that not only accelerates the training process but also reduces computational requirements, maintaining or surpassing traditional training methods' performance in various NLP tasks and T5 model sizes. Finally, we validate our theoretical framework, which guides the optimal use of Deep Fusion, showing that with carefully optimized training dynamics, it significantly reduces both training time and resource consumption.

Figures (8)

• Figure 1: Accuracy and loss on validation data. The x-axis of the graph is scaled in millions of steps.
• Figure 2: Settings for final T5-L fusion: yellow signifies fused models, white indicates regular training, and links represent fusion (double link signifies self fusion). Every node in the graph is trained 1M steps (as an example - algorithm 4 is trained a total of 7M steps).
• Figure 3: Performance (Glue average - an average over many NLP tasks that score between 0 and 100) of the various models.
• Figure 4: Performance on T5 when applying large offsets to the learning rate schedule.
• Figure 5: Heat maps of the first layers' feed forward kernel every 50K steps (left to right till 400K steps training). The upper row is for high learning rates (highly negative offset: -50K), while the lower row is displaying the heat maps for low learning rates (highly positive offset: +50K). Figure \ref{['fig:zoomin_heatmap']} in the Appendix zooms in on the last heat map in the first row showcasing how the replicas of the smaller model diverge.

Theorems & Definitions (17)

• Theorem 1: Modified loss
• Theorem 2: Modified equation
• Corollary 1
• Lemma 1: Scaled gradient
• Lemma 2: Same gradient
• Lemma 3: Non-zero Lie bracket
• Corollary 2
• Theorem 2: Modified loss
• Theorem 2: Modified equation
• proof