PolicyFRZ: Adaptive layer freezing via policy networks for accelerating DNN training in edge environments

Layer freezing can significantly accelerate the training of deep neural networks (DNNs) in edge environments. Existing freezing strategies typically rely on heuristic rules or intermediate model states such as activations, gradients, or loss values to assess parameter convergence. However, these methods often exhibit poor robustness to training fluctuations, which can result in degraded model accuracy. This paper proposes PolicyFRZ, a policy network based dynamic layer-freezing algorithm for accelerating DNN training. PolicyFRZ leverages a policy network within a reinforcement learning framework to dynamically determine effective freezing strategies during training. The policy network is trained using contextual information gathered throughout the training process, enabling it to make informed decisions and learn generalizable freezing patterns. Furthermore, a two-stage training scheme is designed to transfer the learned freezing patterns across different models and datasets. Extensive experiments on computer vision (CV) and natural language processing (NLP) tasks demonstrate that PolicyFRZ accelerates model training without compromising accuracy and outperforms existing layer freezing methods.