Flexora: Flexible Low Rank Adaptation for Large Language Models

TL;DR

Flexora addresses overfitting and inefficiency in LoRA-based fine-tuning of large language models by automatically selecting the most impactful layers for adaptation. It frames layer selection as a hyperparameter optimization problem and solves it with unrolled differentiation, followed by a fine-tuning stage that updates only the chosen layers. Empirical results across multiple models and tasks show that this approach reduces the number of trainable parameters while achieving superior performance compared to LoRA and other baselines, with minimal additional computational overhead. Theoretically, Flexora links layer sparsity to reduced network smoothness and better generalization, providing a principled justification for focusing fine-tuning on a subset of layers and offering a practical, scalable strategy for PEFT in diverse downstream tasks.

Abstract

Large Language Models (LLMs) are driving advancements in artificial intelligence by increasing the scale of model parameters, which has significantly enhanced generalization ability and unlocked new capabilities in practice. However, their performance in specific downstream tasks is usually hindered by their knowledge boundaries on these tasks. Thus, fine-tuning techniques, especially the widely used Low-Rank Adaptation (LoRA) method, have been introduced to expand the boundaries on these tasks, whereas LoRA would underperform on certain tasks owing to its potential overfitting on these tasks. To overcome this overfitting and improve the performance of LoRA, we propose the flexible low rank adaptation (Flexora) method to automatically and flexibly select the most important layers needing to be fine-tuned to achieve the best performance on different downstream tasks. Specifically, Flexora firstly frames this layer selection problem as a well-defined hyperparameter optimization (HPO) problem, then addresses it using the unrolled differentiation (UD) method, and finally selects the most useful layers based on the optimized hyperparameters. Our extensive experiments on many pretrained models and natural language tasks show that Flexora is able to consistently improve over the existing baselines, indicating the effectiveness of our Flexora in practice. We additionally provide insightful theoretical results and many ablation studies to deliver a comprehensive understanding of our Flexora.

Figures (15)

• Figure 1: An overview of Flexora: (a) Initialization of hyperparameters $\widehat{\alpha}$ and their integration with LoRA parameters to produce the Trainable Model. (b) Simultaneous training of LoRA parameters and hyperparameters $\widehat{\alpha}$ using different datasets, minimizing empirical risk for both validation and training datasets. The hyperparameter vector $\widehat{\alpha}$ is then ranked based on magnitude. (c) Flexible selection of layers to be trained, where higher-ranked layers are activated for training while others remain frozen.
• Figure 2: This figure depicts the relationship between the number of LoRA fine-tuning layers and model accuracy across four distinct datasets: Hellaswag, PIQA, Winogrande, and RACE, with the latter including two separate tasks, RACE-mid and RACE-high, which vary in difficulty. Results for LoRA rank 8 are shown here. The $x$-axis represents the number of fine-tuned layers, ranging from 0 to 32, where 0 corresponds to the base model without fine-tuning. Selected configurations include 6, 12, 18, 24, and 32 randomly fine-tuned layers. The full 32-layer configuration, representing the vanilla LoRA setup, is shown as a horizontal dashed line in the plots. The $y$-axis indicates model accuracy as a percentage.
• Figure 3: Comparison of the accuracy of various models (Llama-3-8B, ChatGLM3-6B, Mistral-7B-v0.1, and Gemma-7B) across different tasks. Bars with green diagonal stripes represent LoRA accuracy, while blue circles indicate Flexora accuracy, and the red dotted line represents the improvement ratio of Flexora over LoRA. Notably, Flexora generally outperforms LoRA in most tasks and models, demonstrating its effectiveness.
• Figure 4: We present LLM as a hierarchical network. In this context, all parameters of a Decoder layer are represented as a weight matrix $W$ for subsequent analysis.
• Figure 5: Training and validation loss during the flexible layer selection phase. The figure shows the training and validation loss over 20,000 steps for four different datasets (Hellaswag, PIQA, RACE, and Winogrande), where the batch size at each step is 1. The blue line shows the validation loss and the orange line shows the training loss. These plots visually compare how the performance of the models changes during the flexible layer selection phase, highlighting the convergence behavior.

Theorems & Definitions (6)

• Proposition 1
• Theorem 1: Theorem 3.8 in c:84
• Proposition 2
• proof
• Definition 1
• proof