GeLoRA: Geometric Adaptive Ranks For Efficient LoRA Fine-tuning

TL;DR

Fine-tuning large language models (LLMs) is computationally expensive. While Low-Rank Adaptation (LoRA) offers efficiency by updating fewer parameters, choosing the right "rank" is tricky. This paper introduces GeLoRA, a principled method that uses the intrinsic dimensionality of data representations to adaptively set LoRA ranks, achieving better performance and efficiency than existing methods.

What Problem Does This Paper Address?

Fine-tuning massive LLMs for specific tasks is crucial for personalization and performance, but it demands significant computational resources (time and memory). Parameter-Efficient Fine-Tuning (PEFT) methods, like LoRA, aim to reduce this burden by only updating a subset of model parameters. LoRA achieves this by decomposing weight updates into low-rank matrices. However, a key challenge with LoRA is determining the optimal rank for these matrices. A low rank boosts efficiency but can limit the model's ability to capture complex patterns (expressivity), while a high rank improves expressivity at the cost of increased computation. Existing adaptive LoRA methods often lack a strong theoretical foundation for this rank selection, relying on trial-and-error. This paper addresses the need for a principled, theoretically grounded approach to adaptively select LoRA ranks that balances efficiency and expressivity.

Methods

The core of GeLoRA lies in its novel approach to determining LoRA ranks based on the intrinsic dimensionality (ID) of data representations within the LLM.

• Theoretical Framework: The authors establish a theoretical link between the intrinsic dimensionality of data representations and the optimal rank for LoRA matrices. They derive a lower bound for the optimal rank, showing it depends on the ID of the input and output data manifolds of transformer blocks.
• Intrinsic Dimension Estimation: GeLoRA uses the Two Nearest Neighbors (TwoNN) method to estimate the intrinsic dimensionality of hidden representations. TwoNN is chosen for its efficiency and robustness in high-dimensional settings.
• Adaptive Rank Allocation: Based on the theoretical framework and estimated IDs, GeLoRA adaptively assigns LoRA ranks to each transformer block. The minimum rank for a block is determined by the difference in intrinsic dimensions between its input and output representations, plus a small offset to avoid zero ranks. This ensures that blocks processing data with increasing complexity receive higher ranks.
• Uniform Allocation within Blocks: While ranks are adapted per block, within each block, the ranks for the key, query, value, and output projection matrices are set uniformly to the block's computed rank.

Datasets

The paper evaluates GeLoRA on several standard NLP benchmarks:

• GLUE: A collection of diverse natural language understanding tasks.
• SQuAD (v1.1 and v2.0): Datasets for question answering.
• MT-Bench: Used for evaluating instruction-following capabilities, trained on the Airoboros dataset.

The models used include DeBERTaV3-base (encoder-only) and LLaMA 3 (8B) (decoder-only).

Key Insights and Ideas

• Intrinsic Dimension as a Rank Lower Bound: The paper's central insight is that the intrinsic dimensionality of data representations provides a theoretical lower bound for the optimal LoRA rank. This offers a principled way to move beyond heuristic rank selection.
• Geometry-Driven Adaptation: GeLoRA leverages the "geometry" of data representations (how they are structured in lower-dimensional manifolds) to guide parameter-efficient fine-tuning. This geometric perspective is key to its effectiveness.
• Dynamic Rank Allocation: Unlike static LoRA, GeLoRA dynamically adjusts ranks across layers, recognizing that different parts of a model might require varying levels of expressivity.
• Efficiency-Performance Balance: The results consistently show GeLoRA achieving state-of-the-art performance among PEFT methods while using significantly fewer parameters and less computation time.

Key Results

• Performance Gains: GeLoRA consistently outperforms existing adaptive LoRA methods (e.g., AdaLoRA, SaLoRA) by up to +1.0% on GLUE, SQuAD, and MT-Bench benchmarks.
• Computational Efficiency: GeLoRA reduces computational time by 13.5% to 64.2% compared to baselines, while keeping the total number of tunable parameters fixed. For example, on SQuAD, GeLoRA uses 80.34 x 10⁻³M parameters and achieves better EM/F1 scores than LoRA variants with higher parameter counts.
• Adaptive Rank Patterns: The paper visualizes GeLoRA's rank patterns, showing how ranks are assigned differently across layers and tasks, reflecting the varying complexity of data representations. For instance, some layers might consistently receive higher ranks than others.

Limitations and Future Work

• TwoNN Estimator Accuracy: The effectiveness of GeLoRA relies on the TwoNN estimator for intrinsic dimension. The authors acknowledge that TwoNN can underestimate ID in high-dimensional or non-uniform settings. While they mitigate this with an ad-hoc offset, a more robust theoretical justification is lacking.
• Theoretical Bounds Tightness: The derived theoretical bounds on LoRA ranks are lower bounds and might be loose in practice. The tightness of these bounds and their ability to constrain true optimal ranks needs further investigation.
• Scalability: While evaluated on large models like LLaMA 3, GeLoRA's behavior on even larger models, multilingual LLMs, or multi-modal architectures is not yet explored.
• Preprocessing Overhead: The preprocessing step to estimate intrinsic dimensions, though one-time, can introduce a nontrivial overhead, especially for very large datasets.
• Future Work: The authors suggest exploring more robust and scalable intrinsic dimension estimators, extending the framework to multimodal and multilingual tasks, and investigating dynamic or sample-specific rank assignment.