[Paper Review] Efficient Fine-Tuning of BERT Models on the Edge

Danilo Vucetic, Mohammadreza Tayaranian+4 more|arXiv (Cornell University)|May 3, 2022

Distributed and Parallel Computing Systems4 citations

TL;DR

This paper proposes Freeze And Reconfigure (FAR), a memory-efficient fine-tuning method for BERT-like models on edge devices by selectively freezing parameters and reconfiguring the model architecture to reduce activation and gradient memory usage. FAR reduces fine-tuning time by 30% and memory access time by 47% on DistilBERT with only a 1% average performance drop on GLUE and SQuAD benchmarks.

ABSTRACT

Resource-constrained devices are increasingly the deployment targets of machine learning applications. Static models, however, do not always suffice for dynamic environments. On-device training of models allows for quick adaptability to new scenarios. With the increasing size of deep neural networks, as noted with the likes of BERT and other natural language processing models, comes increased resource requirements, namely memory, computation, energy, and time. Furthermore, training is far more resource intensive than inference. Resource-constrained on-device learning is thus doubly difficult, especially with large BERT-like models. By reducing the memory usage of fine-tuning, pre-trained BERT models can become efficient enough to fine-tune on resource-constrained devices. We propose Freeze And Reconfigure (FAR), a memory-efficient training regime for BERT-like models that reduces the memory usage of activation maps during fine-tuning by avoiding unnecessary parameter updates. FAR reduces fine-tuning time on the DistilBERT model and CoLA dataset by 30%, and time spent on memory operations by 47%. More broadly, reductions in metric performance on the GLUE and SQuAD datasets are around 1% on average.

Motivation & Objective

• Address the challenge of on-device fine-tuning of large BERT-like models on resource-constrained edge devices.
• Reduce memory and computation overhead during fine-tuning, which is significantly higher than inference due to activation and gradient storage.
• Enable practical on-device adaptation of pre-trained language models in dynamic environments without relying on cloud-based retraining.
• Maintain high model performance while drastically reducing training resource consumption on edge hardware.
• Develop a structured, efficient parameter freezing strategy that avoids unstructured memory access patterns common in pruning-based methods.

Proposed method

• Introduce a novel L1-norm-based metric to identify high-learning potential parameter groups for selective fine-tuning.
• Propose a dynamic architecture reconfiguration that groups frozen and non-frozen parameters together to improve memory access locality.
• Implement a two-stage process: priming (initialization of learning potential nodes) followed by selective weight updates based on retention percentage.
• Use 1% priming to minimize training overhead while maintaining performance, with retention percentages (e.g., 10%, 40%) controlling the trade-off between efficiency and accuracy.
• Apply parameter freezing selectively to non-critical layers or nodes, avoiding full-layer freezing to preserve model capacity.
• Optimize memory access patterns by restructuring the model’s parameter layout to reduce cache misses and main memory bandwidth usage.

Experimental results

Research questions

• RQ1 Can selective parameter freezing combined with architectural reconfiguration significantly reduce memory usage during BERT fine-tuning on edge devices?
• RQ2 To what extent can fine-tuning time and memory access overhead be reduced without incurring a significant drop in model performance?
• RQ3 How does the choice of priming and retention percentages affect the trade-off between efficiency and accuracy in fine-tuning compressed BERT models?
• RQ4 Does the proposed method outperform existing baselines like BitFit and random node selection in terms of both efficiency and performance?
• RQ5 Can the method be generalized to other compressed models like DistilBERT and applied across diverse NLP tasks such as GLUE and SQuAD?

Key findings

• FAR reduces fine-tuning time by 30% and memory access time by 47% when fine-tuning DistilBERT on the CoLA dataset with 10% retention and 1% priming.
• With 40% retention, FAR achieves a 17% reduction in fine-tuning time and a 36% reduction in memory access time compared to the baseline.
• The average performance drop across GLUE and SQuAD benchmarks is only 1%, with the smallest drop (0.74%) observed for FAR with 40% retention.
• FAR outperforms random node selection and BitFit, especially on complex tasks like SQuAD 2.0, where BitFit suffers a 10%+ F1 drop.
• The method maintains high performance even with aggressive parameter freezing, demonstrating that overparameterization allows for selective, efficient fine-tuning.
• The ablation study confirms that priming and structured reconfiguration are essential, as random selection and BitFit yield significantly worse results under the same resource constraints.