AdaGradSelect: An adaptive gradient-guided layer selection method for efficient fine-tuning of SLMs

2025-12-12

AI Summary

To address the trade-off between high computational overhead in fine-tuning small language models (SLMs) and the performance degradation of parameter-efficient methods (e.g., LoRA) due to restrictive low-rank factor spaces, this paper proposes a gradient-driven adaptive layer selection fine-tuning method. Our approach innovatively integrates gradient norm evaluation, Dirichlet prior sampling, and an ε-greedy exploration strategy, enabling broad-layer exploration early in training and progressive focusing on critical Transformer layers as training proceeds. Only parameters of selected sparse layers are updated, while all others remain frozen. Experiments demonstrate that, compared to LoRA (rank 256), our method achieves a 3-percentage-point accuracy gain on GSM8K and matches full-parameter fine-tuning performance on MATH. Moreover, it accelerates training by 12% and reduces GPU memory consumption by 35%.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large Language Models (LLMs) can perform many NLP tasks well, but fully fine-tuning them is expensive and requires a lot of memory. Parameter-Efficient Fine-Tuning (PEFT) methods such as LoRA reduce this cost by adding small low-rank updates to frozen model weights. However, these methods restrict the training to a limited subspace, which can sometimes reduce performance. For Small Language Models (SLMs), where efficiency gains matter even more, we introduce AdaGradSelect, an adaptive method that selects which transformer blocks to update based on gradients. Early observations showed that updating only the transformer blocks with the highest gradient norms can achieve performance close to full fine-tuning. Building on this insight, AdaGradSelect adaptively chooses which blocks to train. It uses a combination of Dirichlet-based sampling, which depends on how frequently blocks were updated in the past, and an epsilon-greedy exploration strategy. This lets the method explore different blocks in early training and gradually focus on the most important ones in later epochs. Experiments show that AdaGradSelect trains about 12 percent faster and uses 35 percent less GPU memory while delivering performance very close to full fine-tuning. On the GSM8K dataset, it outperforms LoRA (rank 256) by about 3 percent on average across models such as Qwen2.5-0.5B, LLaMA3.2-1B, and Phi4-mini-3.8B. It also achieves similar accuracy on the MATH dataset. Overall, AdaGradSelect provides a more effective and resource-efficient alternative to traditional fine-tuning methods.

Problem

• Reduces computational cost and memory usage for fine-tuning small language models
• Improves performance over parameter-efficient methods by selecting optimal layers
• Adaptively chooses transformer blocks using gradient-guided exploration strategies

Innovation

• Adaptive gradient-guided transformer block selection
• Dirichlet sampling with epsilon-greedy exploration strategy
• Faster training with less GPU memory usage