Geometry aware method improves fine tuning and uncertainty in AI models
Geometry-Aware Bayesian Parameter-Efficient Fine-Tuning on the Stiefel Manifold via Stein Variational Gradient Descent
Machine Learning
Summary
Fine-tuning large AI models often involves adjusting many connected parts efficiently. The authors studied a way to improve this process by using math that respects the special shape those parts naturally form, called the Stiefel manifold. By applying a technique called Stein variational gradient descent within this shape, their method not only enhances the model's predictions but also gives a better sense of how confident those predictions are. This helps make AI systems both more accurate and more reliable in estimating uncertainty.
Parameter-efficient fine-tuningLow-rank adaptationStiefel manifoldOrthogonality constraintsSingular value decompositionStein variational gradient descentGeometric structureUncertainty quantificationModel calibrationPre-trained models
Authors
Quang-Duy Tran, Trung Le, Bao Duong, Phuoc Nguyen, Thin Nguyen
Abstract
Several geometry-aware approaches to low-rank adaptation have emerged for parameter-efficient fine-tuning of large pre-trained models. These methods aim to take full advantage of the geometric structure of low-rank manifolds for improving the efficiency in subspace utilization and reducing redundancy by enforcing orthogonality constraints during optimization. The strong empirical results of these techniques have motivated further study into whether predictions from such geometry-based adaptation methods could be overconfident. In this paper, we build on the singular value decomposition factorization of adapters to develop a framework based on Stein variational gradient descent (SVGD). In this formulation, the low-rank matrices are transported along the Stiefel manifold to match the targeted distributions while retaining their crucial geometric structure. Since this geometry-aware SVGD approach provides multiple solutions during inference, it supports uncertainty quantification and produces better-calibrated adapters on the Stiefel manifold. Extensive experiments show that our method delivers strong model calibration and attains higher prediction accuracy than SVGD and related uncertainty estimation methods that are formulated in Euclidean space.