COBRA: Contribution-Based Bayesian Rank Allocation for Parameter-Efficient Fine-Tuning
Abstract
Full fine-tuning of large language models (LLMs) incurs prohibitive computational and storage costs. Parameter-efficient fine-tuning (PEFT) addresses this limitation, with Low-Rank Adaptation (LoRA) gaining widespread adoption due to its simplicity and zero inference overhead. However, LoRA and its variants typically rely on uniform rank allocation or a single importance metric such as gradient magnitude or output sensitivity to guide rank distribution. This approach fails to recognize that gradient magnitude and output contribution are decoupled properties, leading to suboptimal allocation where critical layers are under-provisioned while less important ones waste capacity. To address this challenge, we propose COBRA, a principled framework integrating dual importance factors for adaptive rank allocation. COBRA operates in three stages: (1) layer conductance attribution quantifies each layer's contribution via path-integral attribution; (2) dual-factor aggregation combines contribution with adaptation demand, producing the Task-Adaptive Layer Conductance (TA-LC) distribution; and (3) Bayesian rank allocation translates this distribution into optimal heterogeneous ranks via variational optimization. Layer conductance provides layer-level interpretability by explicitly quantifying how much each layer contributes to predictions without redundancy, directly aligning with the granularity of rank allocation decisions and enabling principled cross-layer comparison for rank distribution. Experiments across diverse architectures and tasks demonstrate that COBRA consistently outperforms existing methods, achieving up to 1.6 points improvement on GLUE and a 6.6\% average MSE reduction in high-rank regression regimes under comparable parameter budgets.
Lay Summary
Modern AI language models are generalists, and turning one into a specialist, say to answer medical questions or write code, is famously expensive. To avoid retraining the whole model, engineers freeze it and add a small set of adjustable "tuning knobs." But most methods give every layer of the model the same number of knobs, even though some layers matter far more than others. The tempting fix is to follow a single clue, such as which layers change most while learning. That sounds reasonable, and it is wrong. We show that two separate things should guide the decision: how much a layer needs to change to learn the task, and how much it actually shapes the model's final answers. Like a fixed renovation budget that should go to the rooms that both need repair and get the most use, these two clues often disagree, so methods watching only one spend effort in the wrong places. Our method, COBRA, measures both, blends them into a single score, and directs a fixed budget of knobs where they count most. The result is cheaper, sharper, and less energy-hungry AI customization that beats existing methods at the same cost.