AdaRank: Disagreement Based Module Rank Prediction for Low-rank Adaptation

TL;DR

AdaRank addresses the inefficiency of applying a uniform low-rank adaptation across all layers in large pretrained models. It introduces a disagreement-based, two-step method to predict per-module ranks by perturbing individual modules and measuring output disagreement, then mapping these importances to absolute ranks with a target average rank $r$. The approach preserves pretraining and finetuning stages without extra objectives, and empirical results on TREC, Yelp, and AG News show AdaRank improves generalization over uniform-rank baselines, both when adapting individual modules and when adapting all modules jointly. The findings highlight that later transformer modules are more critical and that lower disagreement correlates with better generalization, offering a practical and transferable strategy for parameter-efficient adaptation with potential runtime and memory benefits. The work suggests future directions including broader input protocols, task-specific vs generic prompts for rank estimation, and theoretical grounding of the disagreement-to-rank link.

Abstract

With the rise of language and multimodal models of ever-increasing size, pretraining a general-purpose foundational model and adapting it to downstream tasks has become common practice. To this end, adaptation efficiency can be a critical bottleneck given the large model sizes, hence efficient finetuning methods such as LoRA have become prevalent. However, LoRA is typically applied with the same rank across all model layers, despite mounting evidence from transfer learning literature that during finetuning, later layers diverge more from pretrained weights. Inspired by the theory and observations around feature learning and module criticality, we develop a simple model disagreement based technique to predict the rank of a given module relative to the other modules. Empirically, AdaRank generalizes notably better on unseen data than using uniform ranks with the same number of parameters. Compared to prior work, AdaRank has the unique advantage of leaving the pretraining and adaptation stages completely intact: no need for any additional objectives or regularizers, which can hinder adaptation accuracy and performance. Our code is publicly available at https://github.com/google-research/google-research/tree/master/adaptive_low_rank.

Figures (2)

• Figure 1: Disagreement-based module relative importance prediction in AdaRank. For a module $\mathfrak{m}$ (in orange), to approximate its importance score, we 1) add a random noise tensor of the same shape to $\mathfrak{m}$'s pretrained weights, while keeping all other modules frozen (in blue), 2) do this twice, and 3) take the $\ell_1$ difference between model logits produced by these two perturbed model instances.
• Figure 2: Model disagreement rates when perturbing different modules, perturbed one layer at a time. As outlined in Algorithm \ref{['alg:predict-importance']}, random noise is added to a given module, while the remaining modules retain their pretrained weights. The model disagreement is calculated as the $\ell_1$ difference in model logits between two perturbed model instances. Each data point is the average of five runs. These perturbation rates are used as module importance scores for predicting AdaRank ranks in Algorithm \ref{['alg:disagreement-to-rank']}. The upward trends in these model disagreements mirror prior findings that later layers deviate more from pretrained weights during adaptation moduleCriticalityChatterji2020. Note the varying levels and ranges of disagreement rates for different modules, indicating varying levels of module criticality.