Klemens Flöge
All publications

Gaussian Stochastic Weight Averaging for Bayesian Low-Rank Adaptation of Large Language Models

Emre Onal, Klemens Flöge, Emma Caldwell, Arsen Sheverdin, Vincent Fortuin

6th Symposium on Advances in Approximate Bayesian Inference (non-archival workshop)·workshop

The SWAG-LoRA training and inference process. Left: the LLM with LoRA fine-tuning. Middle and top right: SWAG training, collecting weight samples across SGD iterations to estimate the mean and an approximate covariance of the posterior. Bottom right: forming the inference ensemble by sampling from the learned SWAG posterior.
The SWAG-LoRA training and inference process. Left: the LLM with LoRA fine-tuning. Middle and top right: SWAG training, collecting weight samples across SGD iterations to estimate the mean and an approximate covariance of the posterior. Bottom right: forming the inference ensemble by sampling from the learned SWAG posterior.

Abstract

Fine-tuned Large Language Models (LLMs) often suffer from overconfidence and poor calibration, particularly when fine-tuned on small datasets. To address these challenges, we propose a simple combination of Low-Rank Adaptation (LoRA) with Gaussian Stochastic Weight Averaging (SWAG), facilitating approximate Bayesian inference in LLMs. Through extensive testing across several Natural Language Processing (NLP) benchmarks, we demonstrate that our straightforward and computationally efficient approach improves model generalization and calibration competitively with comparable, more sophisticated methods for Bayesian inference in LLMs. We further show that our method exhibits greater robustness against distribution shift, as reflected in its improved performance on out-of-distribution tasks.

Fine-tuned LLMs are badly calibrated, especially on small datasets — they are confidently wrong. Full Bayesian treatment is out of reach at this scale, so the question is how much you can recover cheaply.

The answer here is: quite a lot. Stacking Gaussian SWAG on top of LoRA yields an approximate posterior over the adapter weights at almost no extra cost, and matches far more sophisticated Bayesian methods on NLP benchmarks. The gap widens under distribution shift, where the ensemble’s calibration advantage shows up most clearly.