SmartFRZ: An Efficient Training Framework using Attention-Based Layer Freezing

TL;DR

Training deep networks efficiently remains challenging due to high compute and energy costs. SmartFRZ addresses this by using an attention-based predictor to automatically decide which layers to freeze and when during training, with the predictor trained offline on layer representations labeled by representational similarity. The framework demonstrates significant reductions in training time and memory while maintaining or improving accuracy across computer vision and natural language tasks, outperforming prior freezing methods. This approach provides a generic, scalable path to faster, greener model training applicable to diverse architectures and datasets.

Abstract

There has been a proliferation of artificial intelligence applications, where model training is key to promising high-quality services for these applications. However, the model training process is both time-intensive and energy-intensive, inevitably affecting the user's demand for application efficiency. Layer freezing, an efficient model training technique, has been proposed to improve training efficiency. Although existing layer freezing methods demonstrate the great potential to reduce model training costs, they still remain shortcomings such as lacking generalizability and compromised accuracy. For instance, existing layer freezing methods either require the freeze configurations to be manually defined before training, which does not apply to different networks, or use heuristic freezing criteria that is hard to guarantee decent accuracy in different scenarios. Therefore, there lacks a generic and smart layer freezing method that can automatically perform ``in-situation'' layer freezing for different networks during training processes. To this end, we propose a generic and efficient training framework (SmartFRZ). The core proposed technique in SmartFRZ is attention-guided layer freezing, which can automatically select the appropriate layers to freeze without compromising accuracy. Experimental results show that SmartFRZ effectively reduces the amount of computation in training and achieves significant training acceleration, and outperforms the state-of-the-art layer freezing approaches.

Figures (5)

• Figure 1: Overview of our attention-based layer freezing framework SmartFRZ. SmartFRZ continuously records the weights of all actively being-trained layers during the training process. Periodically, there will be a freezing stage, when SmartFRZ applies an attention-based predictor to determine whether each active layer can be frozen at this moment, based on the collected historical data. After the freezing stage, the model training process continues, and the frozen layer will not be updated and its weights will not be recorded anymore. In this figure, we assume that the attention-based predictor predicts to freeze layer 1 only.
• Figure 2: (a) An illustrated example of noisy weight history in which a darker shade within the circle indicates a lower loss. The training loss decreases at iteration $\{1,2,5\}$, yet increases at iterations $\{3,4\}$. So the predictor needs to pay more attention to the weight history from iteration $\{1,2,5\}$ and less attention to the weight history from iterations $\{3,4\}$ since the latter converge towards less optimal directions. (b) Detailed workflow of the attention-based predictor. At timestamp $t$, the predictor decides whether to freeze a specific layer $l$ in three steps. First, it encodes the training history independently into feature vectors (i.e., $K^j_l$, $Q^j_l$, $V^j_l$ ). Second, it computes attention scores $\alpha^j_l$ and aggregates the historical state feature into a context vector $C^t_l$. Third, it predicts the confidence scores of freezing or not and selects the decision with higher confidence.
• Figure 3: The histogram of the gradient value frequency distribution of a network layer at a certain point during the training. Figure (a) counts all the parameters of the layer, while Figure (b) counts the randomly selected 1024 parameters. These gradient values belong to the 30th CONV layer of ResNet50 and were sampled during the training of ResNet50 on CIFAR-10.
• Figure 4: CKA variation curve as training proceeds. Result is obtained while training ResNet50 on CIFAR-10 and the ResNet50 is pre-trained on ImageNet.
• Figure 5: (a) Memory cost as training proceeds. Results are obtained by training VGG11, ResNet50, and MobileNetV2 on CIFAR-100 dataset. (b) The frozen moment for each layer while training ResNet50 on CIFAR-100.