Flat-LoRA: Low-Rank Adaptation over a Flat Loss Landscape

TL;DR

Flat-LoRA introduces a novel low-rank adaptation method that finds flatter regions in the full parameter space of large models, enhancing generalization without high computational costs, by using a Bayesian loss and refined perturbation strategies.

Contribution

It proposes Flat-LoRA, a new approach that improves low-rank adaptation by ensuring flatness in the full parameter space, avoiding costly sharpness-aware minimization techniques.

Findings

Improves in-domain and out-of-domain generalization across tasks.

Achieves better performance without high computational overhead.

Demonstrates effectiveness on diverse tasks like reasoning, coding, and image generation.

Abstract

Fine-tuning large-scale pre-trained models is prohibitively expensive in terms of computation and memory costs. Low-Rank Adaptation (LoRA), a popular Parameter-Efficient Fine-Tuning (PEFT) method, offers an efficient solution by optimizing only low-rank matrices. Despite recent progress in improving LoRA's performance, the relationship between the LoRA optimization space and the full parameter space is often overlooked. A solution that appears flat in the loss landscape of the LoRA space may still exhibit sharp directions in the full parameter space, potentially compromising generalization. We introduce Flat-LoRA, which aims to identify a low-rank adaptation situated in a flat region of the full parameter space. Instead of adopting the well-established sharpness-aware minimization approach, which incurs significant computation and memory overheads, we employ a Bayesian expectation loss objective to preserve training efficiency. Further, we design a refined random perturbation generation strategy for improved performance and carefully manage memory overhead using random seeds. Experiments across diverse tasks-including mathematical reasoning, coding abilities, dialogue generation, instruction following, and text-to-image generation-demonstrate that Flat-LoRA improves both in-domain and out-of-domain generalization. Code is available at https://github.com/nblt/Flat-LoRA.