Spectral Bridge Variational Inference: Dynamic LoRA via Bures-Wasserstein Gradient Flows
Abstract
Parameter-Efficient Fine-Tuning (PEFT) is essential for adapting Large Language Models, yet existing methods struggle to balance capacity with computational efficiency. Standard approaches enforce rigid low-rank constraints, while dynamic alternatives incur significant memory overheads. To resolve this, we propose Spectral Bridge Variational Inference (SBVI), a geometric framework reformulating LoRA as a continuous Wasserstein gradient flow on the manifold of Gaussian measures. Instead of fixing ranks at initialization, SBVI governs singular value evolution via a stochastic differential equation driven by thermodynamic competition between task gradients and adaptive entropic friction. This induces a spectral bifurcation that automatically prunes noise modes while amplifying signal-rich components, discovering an optimal layer-wise rank distribution. We derive a scalable algorithm with linear complexity using factorized Riemannian retractions and Empirical Bayes friction updates. Experiments on reasoning and coding benchmarks show SBVI achieves state-of-the-art performance, offering superior accuracy and memory efficiency over existing static and dynamic methods.
Lay Summary
Fine-tuning large language models to perform specific tasks is highly resource-intensive. A common approach, Low-Rank Adaptation (LoRA), reduces this burden by updating only a small, fixed-size subset of parameters. However, finding the right size is challenging: too small restricts the model's capabilities, while too large wastes memory. To address this, we introduce a method that dynamically adjusts this capacity during training. By modeling the learning process as a fluid flow, our approach automatically expands capacity for important task-related information while pruning inactive components. This self-adjusting mechanism reduces memory usage and maintains training efficiency. It also naturally directs more computational resources to the model's middle layers, which handle complex reasoning.