PaperlayerView exact revision on arXiv

Source-linked AI summary

Geometry-Aware Bayesian Parameter-Efficient Fine-Tuning on the Stiefel Manifold via Stein Variational Gradient Descent

Quang-Duy Tran, Trung Le, Bao Duong, Phuoc Nguyen, Thin Nguyen

arXiv:2609.08354v1cs.LG

TL;DR

Fine-tuning can disrupt calibration and produce overconfident predictions, while geometry-aware low-rank methods have lacked calibration and uncertainty evaluations. StePS applies SVGD to SVD-based adapters on the Stiefel manifold, and evaluations report favorable accuracy and calibration results across in- and out-of-distribution settings.

  • Problem

    Fine-tuning can break calibration and cause overconfident predictions, while geometry-aware low-rank adaptations lacked evaluations of calibration and epistemic uncertainty.

    Source: §I · p. 1

  • Method

    StePS uses Bayesian SVGD to transport SVD-based adapter particles on the Stiefel manifold toward target distributions, retaining the factors' geometric structure.

    Source: §I · p. 2, §VI · p. 8

  • Results

    Across in- and out-of-distribution evaluations, StePS reports promising accuracy with comparable calibration, and favorable accuracy and expected calibration errors versus related methods on commonsense reasoning benchmarks.

    Source: §I · p. 2, §VI · p. 8

  • Takeaways & Limitations

    The Stiefel-manifold Bayesian posterior predictive distribution provides explicit epistemic uncertainty estimates and improves prediction calibration for geometry-aware PEFT.

    Source: §I · p. 2

    On WG-S, StePS runtime is substantially longer than model-based variational baselines because it evaluates multiple solutions and adds manifold operations, though these costs are minimal against ensemble evaluations.

    Source: A. Runtime · p. 10

Abstract

from arXiv · show

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.

I. INTRODUCTION

The introduction connects calibration and uncertainty concerns in fine-tuning with geometry-aware low-rank adaptation. It presents StePS, a Bayesian inference framework on the Stiefel manifold, and reports favorable accuracy and calibration results.

I. INTRODUCTION

  • Fine-tuning can break model calibration and produce overconfident predictions, while reduced fine-tuning data heightens epistemic uncertainty.Bayesian posterior predictive distributions can estimate epistemic uncertainty and support better-calibrated predictions.Source: §I · p. 1

A. Geometry-Aware Parameter-Efficient Fine-Tuning

  • Geometry-aware LoRA methods impose orthogonality constraints on adapter factors to use rank subspace dimensions, but their calibration and uncertainty remain unevaluated.Stiefel-LoRA constrains the up-projection basis, while StelLA constrains both up- and down-projection bases.Source: §I · p. 1, A. Geometry-Aware Parameter-Efficient… · p. 2

I. INTRODUCTION

  • StePS shows favorable accuracy and expected calibration errors versus related methods on in- and out-of-distribution commonsense reasoning benchmarks.The introduction also reports improvements in predictive accuracy and calibration compared with related Bayesian PEFT frameworks.Source: §I · p. 2
  • StePS performs Bayesian inference on the Stiefel manifold to estimate epistemic uncertainty and improve calibration through the posterior predictive distribution.The framework uses SVGD to transport particles on the manifold toward target distributions.Source: §I · pp. 1–2
  • StePS derives a closed-form steepest-descent update from gradient flow to transport particles toward target distributions on the Stiefel manifold.The introduction frames this as an optimal transport function for iterative particle updates.Source: §I · p. 2

B. Uncertainty Quantification of Large Models

Existing Bayesian uncertainty methods for large models offer different calibration and stability trade-offs, but largely treat low-rank parameters as Euclidean. The section introduces a tri-factor adapter whose basis matrices are optimized on the Stiefel manifold.

B. Uncertainty Quantification of Large Models

  • Existing Bayesian PEFT methods can improve uncertainty estimates and calibration, but predominantly use Euclidean parameters and miss low-rank geometry.Monte Carlo Dropout and Deep Ensemble generally yield limited calibration gains; Laplace-LoRA is posthoc, while BLoB and ScalaBL can require stability techniques.Source: B. Uncertainty Quantification of Large… · p. 2
  • Riemannian SVGD transports posterior particles on manifold-constrained parameter spaces, but its use in large-scale deep learning and PEFT had remained unexplored.Unlike model-based variational inference, SVGD avoids restrictive distributional assumptions at a computational cost.Source: B. Uncertainty Quantification of Large… · p. 2

A. Tri-Factor Factorization of Adapters and Optimization on the Stiefel Manifold

  • The adapter uses truncated SVD as ΔW = USV⊤, with U and V as orthonormal output- and input-subspace bases and S = diag(s).Constraining U and V to the Stiefel manifold preserves their orthonormality during optimization.Source: A. Tri-Factor Factorization of… · pp. 2–3
  • Stiefel optimization projects updates into the tangent space and retracts them back onto the manifold after the optimizer step.The retraction uses the orthogonal matrix from a polar decomposition.Source: A. Tri-Factor Factorization of… · p. 3

B. Stein Variational Gradient Descent for Bayesian Inference

The section formulates inference as gradient flow over distributions and derives a steepest-descent particle update. An RKHS assumption yields the closed-form SVGD update, which can be approximated with particles.

B. Stein Variational Gradient Descent for Bayesian Inference

  • Sampling targets a distribution p(θ) ∝ exp(−βΨ(θ)) over Θ, with Ψ as the energy function, and describes its gradient flow in Wasserstein space.The flow is expressed through a continuity equation.Source: B. Stein Variational Gradient Descent… · p. 3
  • A manifold update first projects a Euclidean step into the local tangent space, then retracts it to obtain the next point on the manifold.The figure depicts projection in blue and retraction in red.Source: Figure 1 · p. 3
  • The steepest-descent solution seeks a velocity field T(θ) = θ + ηϕ*(θ) to push the current distribution toward the target.The derivation considers the push-forward distribution T♯ρt and assumes ϕ belongs to an RKHS defined by a positive semidefinite kernel.Source: B. Stein Variational Gradient Descent… · p. 3
  • With the RKHS assumption, SVGD gives a closed-form optimal update that is approximated in practice using M particles.The update combines a kernel-weighted energy-gradient term and a kernel-gradient term.Source: B. Stein Variational Gradient Descent… · p. 3

IV. STEPS: RIEMANNIAN SVGD FOR BAYESIAN PEFT ON THE STIEFEL MANIFOLD

StePS adapts SVGD to geometry-aware low-rank parameters by transporting distributions on a Riemannian manifold rather than an unconstrained Euclidean space. Its Stiefel-manifold formulation uses tangent-space projection and retraction to preserve the manifold structure while seeking a steepest-descent transport.

IV. STEPS: RIEMANNIAN SVGD FOR BAYESIAN PEFT ON THE STIEFEL MANIFOLD

  • StePS formulates Bayesian adaptation over distributions on a Riemannian manifold, addressing the fact that Euclidean SVGD cannot directly handle parameters on the Stiefel manifold.The framework targets the SVD-based adapter factors U and V, whose orthonormality supports low-rank subspace utilization.Source: §IV · p. 3, A. Gradient Flow on the Riemannian… · p. 4

A. Gradient Flow on the Riemannian Manifold

B. Steepest Descent Solution for the Stiefel Manifold

  • The formulation assumes each parameter θ lies on St(k, r), while more complex model parameters can be represented as collections of such matrices.The stated setting is sufficient for the Stiefel-manifold development.Source: B. Steepest Descent Solution for the… · p. 4
  • The steepest-descent solution is the negative of the derived vector field, with tangent-space projection and retraction used in the Stiefel-manifold formulation.The theorem presents the optimal transportation-map solution for the optimization problem.Source: B. Steepest Descent Solution for the… · pp. 4–5

C. Practical Implementation

The practical implementation approximates the expectation in the optimal solution with particles and uses a shared kernel formulation for manifold and Euclidean adapter parameters. It drives optimization with empirical training loss.

C. Practical Implementation

  • The expectation in the optimal solution is approximated with M particles, following SVGD.The implementation uses a radial basis function kernel.Source: C. Practical Implementation · p. 5
  • The RBF kernel is equivalent up to a multiplicative constant to a von Mises–Fisher kernel on the Stiefel manifold, enabling a unified formulation for U, S, and V.The method computes a joint kernel for the adapter parameters and projects and retracts U and V on the Stiefel manifold; bandwidth is set adaptively from pairwise particle distances.Source: C. Practical Implementation · p. 5
  • The energy function is the empirical loss over the training set, with the task-specific loss applied to model predictions.The loss may be cross-entropy for classification.Source: C. Practical Implementation · p. 5

A. Experimental Settings

The experiments fine-tune Llama2-7B on commonsense reasoning tasks and evaluate in-distribution performance as well as generalization to shifted datasets. Accuracy, ECE, and NLL are compared across deterministic, Bayesian, and particle-based baselines.

A. Experimental Settings

  • The evaluation fine-tunes Llama2-7B on multiple-choice commonsense reasoning benchmarks using negative log-likelihood of the correct answer tokens.Candidate answers are evaluated using their next-token logits.Source: A. Experimental Settings · p. 6
  • The in-distribution setting covers six tasks, while the out-of-distribution setting fine-tunes on OBQA and tests on ARC-C, ARC-E, MMLU-Chemistry, and MMLU-Physics.The shifted tests represent smaller and larger distribution shifts.Source: A. Experimental Settings · p. 6
  • Accuracy measures correct answers, while ECE and NLL assess uncertainty; higher accuracy and lower ECE and NLL are preferred.Source: A. Experimental Settings · p. 6
  • Baselines include deterministic MLE and MAP, Bayesian dropout and ensembles, LoRA uncertainty methods, and Euclidean-space SVGD for LoRA.The implementation uses adapters in attention query/value projections and the softmax output head; the backbone is quantized to 8-bit and stochastic methods use four inference samples or particles.Source: A. Experimental Settings · p. 6

B. In-Distribution Inference Results

On the in-distribution benchmarks, StePS combines strong accuracy with relatively low calibration error. It has the highest accuracy on five of six datasets and improves all reported metrics over deterministic StelLA variants.

B. In-Distribution Inference Results

  • Compared with StelLA MLE and MAP, StePS improves all evaluation metrics; Laplace-LoRA and ScalaBL show relatively minimal calibration gains, while BLoB’s calibration gains are not accompanied by meaningful predictive accuracy.StePS is reported to deliver simultaneous improvements in accuracy and calibration.Source: B. In-Distribution Inference Results · p. 6
  • StePS achieves the highest accuracy on five of six datasets while attaining relatively low ECE and NLL, with the lowest calibration values on several datasets and second-lowest on most others.The in-distribution results are presented in Table II.Source: B. In-Distribution Inference Results · p. 6, Table II · p. 7

C. Out-of-Distribution Inference Results

On out-of-distribution benchmarks, StePS combines competitive accuracy with consistently strong calibration, while StelLA-based methods outperform their LoRA variants on most benchmarks. The conclusion describes StePS as a Stiefel-manifold Bayesian PEFT framework with a closed-form update and reports promising accuracy with comparable calibration.

C. Out-of-Distribution Inference Results

  • Accuracy and calibration degrade on MMLU-Chemistry and MMLU-Physics, which have larger distribution shifts than ARC-C and ARC-E.The passage relates the degradation to the challenge specialized-knowledge domain shifts pose for fine-tuned LLMs.Source: C. Out-of-Distribution Inference Results · p. 7
  • StePS maintains competitive accuracy and second-lowest ECE and NLL across out-of-distribution datasets, while StelLA-based methods outperform LoRA variants on most benchmarks.BLoB has the lowest ECE and NLL, but its accuracy is lower than other baselines on ARC-C and ARC-E; Stiefel formulations appear especially robust under wider shifts.Source: C. Out-of-Distribution Inference Results · pp. 7–8

VI. CONCLUSION

  • StePS estimates epistemic uncertainty by transporting particles on the Stiefel manifold and uses a closed-form update derived from the manifold gradient flow.The conclusion reports promising accuracy and comparable calibration on in- and out-of-distribution settings; it identifies model-based variational inference on Riemannian manifolds as future work to improve computational efficiency.Source: §VI · p. 8

APPENDIX A PROOF OF THEOREM 3

The proof derives the stated form of equation (18) by working through its terms, including the tangent-projection derivative, and combining intermediate trace expressions.

APPENDIX A PROOF OF THEOREM 3

  • The proof recalls equation (18) and derives its first term before treating the tangent-projection derivative.It identifies the RKHS with kernel κ and dimension d = k × r when introducing the first-term derivation.Source: Appendix A · pp. 9–10
  • The tangent-projection calculation formulates trace expressions for its first part and for the derivative involving the flattened matrix αϕ(θ).The proof defines α = θθT and δp,q(θ) in this calculation.Source: Appendix A · p. 10
  • Combining the intermediate trace results yields the stated form of equation (18), completing the proof.Source: Appendix A · p. 10

APPENDIX B ADDITIONAL EXPERIMENTAL RESULTS

The runtime comparison reports added costs for particle-based inference and Stiefel-manifold operations, while the particle ablation identifies four particles as the best accuracy-calibration trade-off.

A. Runtime

  • On WG-S, Deep Ensemble, SVGD, and StePS require substantially longer runtime because they evaluate multiple solutions per forward pass.Stiefel tangent-space projection and retraction add computation, but those costs are described as minimal compared with ensemble evaluations.Source: A. Runtime · p. 10, Table IV · p. 10

B. Ablation Study on the Number of Particles

Loading 2609.08354v1…