PaperlayerView exact revision on arXiv

Source-linked AI summary

Prior Predictive Matching for Bayesian LoRA

Tomasz Kuśmierczyk, Eliezer de Souza da Silva, Marcin Sendera

arXiv:2610.35457v1cs.LGstat.ML

TL;DR

Bayesian LoRA requires specifying a prior over low-rank weight updates, while common selection approaches assess hyperparameters after fitting. This paper derives a closed-form prior-predictive estimator from residual covariance before fine-tuning. Experiments report prior-scale recovery and improved NLL and calibration over MAP and marginal-likelihood selection on RoBERTa-large without validation data.

  • Problem

    Prior hyperparameters for Bayesian LoRA are commonly assessed after fitting, although the prior is meant to encode plausible adaptations before fine-tuning.

    Source: §1 · p. 1

  • Method

    A closed-form prior-predictive estimator matches cross-example residual covariance to estimate the prior scale before fine-tuning.

    Source: §1 · p. 2, §4 · p. 4, §4.2 · p. 4

  • Results

    On RoBERTa-large, the estimated prior improves NLL and calibration over MAP and marginal-likelihood selection without validation data, while providing competitive OOD detection.

    Source: §8 · p. 10

  • Takeaways & Limitations

    Under standard LoRA scaling, the analysis yields λ ∝ r^-1/2, allowing an estimate to transfer across ranks without a new search.

    Source: §8 · p. 9

    The estimator relies on a local second-order approximation and becomes biased when updates induce large network-output changes.

    Source: §8 · p. 10

Abstract

from arXiv · show

Bayesian variants of Low-Rank Adaptation (LoRA) provide principled uncertainty estimation when fine-tuning large foundation models, but require specifying the prior over the low-rank weight updates. We address this hyperparameter selection problem by introducing a closed-form prior-predictive estimator that operates before fine-tuning. We analyze a second-order expansion of the network and quantify how the low-rank prior drives function-space deviations from the pretrained outputs. We also show that these prior-predictive quantities depend on the prior only through the product of the adaptation rank and the fourth power of the prior scale. Because the training targets share a single draw of the latent weight update, we can leverage the covariance between example residuals to derive a moment-matching estimator. Furthermore, our theoretical analysis establishes that the optimal prior precision scales as $r^{-1/2}$ with the rank $r$, allowing a single estimate to generalize across all ranks. Controlled experiments across regression and classification tasks validate our approximation and show that our estimator recovers the true prior scale. On RoBERTa-large, our approach improves predictive negative log-likelihood (NLL) and calibration over both MAP and marginal-likelihood selection without requiring validation data, while achieving competitive out-of-distribution detection.

1 INTRODUCTION

Bayesian LoRA requires choosing a prior scale and rank, but common selection strategies assess them only after fitting. This paper develops a closed-form prior-predictive estimator that uses the pretrained model and fine-tuning data before optimization or validation.

1 INTRODUCTION

  • Existing validation search and approximate marginal-likelihood selection require fitting before assessing prior hyperparameters.The paper frames this as a mismatch: the prior encodes plausible adaptations before fine-tuning, yet is chosen using its outcome.Source: §1 · p. 1
  • The proposed closed-form estimators use the pretrained model and fine-tuning data without LoRA optimization or validation data.They recast prior specification as statistical estimation before fine-tuning.Source: §1 · pp. 1–2
  • The analysis relates LoRA prior scale and rank to moments of induced changes in network outputs.Source: §1 · p. 2
  • The estimator separates shared parameter-induced variation from independent observation noise when estimating prior scale.Source: §1 · p. 2
  • The paper evaluates whether calibration recovers appropriate prior scales and yields competitive uncertainty quantification and predictive performance after Bayesian LoRA fine-tuning.Source: §1 · p. 2

2 BACKGROUND

Bayesian LoRA places a prior on low-rank factors of weight updates, making rank and prior scale central modeling choices. Prior-predictive matching selects hyperparameters by comparing predictive moments with empirical quantities.

2 BACKGROUND

  • LoRA adapts selected weights through a low-rank update, with the Bayesian model placing a prior on its factors.The update is written as ΔW = UV, where K is the rank and the factor dimensions depend on the layer input and output dimensions.Source: §2 · p. 2
  • The standard factor prior uses zero means and a shared scale σ for both factors.The induced prior on ΔW follows from the factor prior and the product ΔW = UV.Source: §2 · p. 2
  • Table 1 compares approaches to prior selection for Bayesian LoRA.Source: Table 1 · p. 2
  • Prior-predictive matching chooses hyperparameters by aligning moments of outputs under the prior with corresponding empirical quantities.The predictive distribution describes how outputs scatter around the pretrained function under the chosen prior.Source: §2 · p. 2

3 PRIOR PREDICTIVE MOMENTS

The prior-predictive analysis connects low-rank update moments to network-output moments through a second-order expansion. Its key scalar is θ, which combines rank and factor scale and determines how prior effects propagate.

3.1 Moments of the update

  • With zero-mean factors, update entries are mutually uncorrelated and share variance θ = Kσ^4.Thus second-order quantities cannot distinguish ranks when their factor scales yield the same θ; distinguishing rank requires fourth-order information involving the 6/K kurtosis term.Source: §3.1 · pp. 2–3

3.2 Propagation through the network

  • The second-order expansion propagates update moments to network outputs, with the mean shift determined by curvature and output covariance by gradient inner products.The derivation considers a fixed input and scalar output, with an additional output dimension for multiple outputs.Source: §3.2 · p. 3
  • The prior-predictive distribution depends on rank and factor scale only through θ.The expansion’s accuracy is studied in Section 7.1, while Appendix A treats omitted higher-order terms.Source: §3.2 · p. 3

3.3 Practical consequences

  • Under practical LoRA scaling with fixed α, the optimal prior precision scales as r^-1/2.An estimate of θ transfers across ranks without a new search because θ does not depend on rank.Source: §3.3 · pp. 3–4

4 MOMENT MATCHING

Moment matching must account for the fact that examples share one random LoRA update, rather than providing independent prior draws. The section develops residual-based estimators for binary classification, regression, and multiclass tasks.

4 MOMENT MATCHING

  • A dataset shares one draw of the LoRA update, so averaging examples does not directly average over the prior distribution.Examples instead vary in their inputs and label noise, which makes prior-predictive moment estimation more challenging.Source: §4 · p. 4

4.1 Basics

  • For binary classification, observable residuals compare labels with pretrained probabilities, while outputs under the realized update are unobserved.The residual is defined as r_n = y_n − p_0,n.Source: §4.1 · p. 4

4.2 Binary classification estimator

  • Cross-example residual products estimate θ because examples share the update draw while their label-noise terms are independent and mean-zero conditional on that draw.The method uses pairs of different examples and estimates θ from the resulting equations by least squares.Source: §4.2 · p. 4

4.3 Regression estimator

  • For regression, the residual-based estimators apply without knowing the observation-noise variance; known noise permits an alternative estimator.Source: §4.3 · p. 4

4.4 Multiclass estimator

  • For multiclass estimation, gradients are computed per class and residual products are taken across different examples, avoiding within-example dependent label noise.The construction reduces to the binary estimator when C = 1.Source: §4.4 · p. 5

5 IMPLEMENTATION

The implementation avoids storing costly per-example gradients and estimates the prior precision in one training-data pass using random sketches. The estimate is available before fine-tuning and applies across LoRA ranks.

5 IMPLEMENTATION

  • The estimator avoids storing full per-example gradient matrices, which can be prohibitively expensive for large models.It estimates the required quantities in one pass over the training data.Source: §5 · p. 5
  • Independent priors across adapted modules make their first-order contributions uncorrelated; experiments use a shared prior scale across modules.Separate module scales yield a noise-dominated least-squares problem at typical single-task sample sizes.Source: §5 · p. 5

5.1 Gradient sketches

  • Random projections replace per-example gradients because the estimator depends on them only through inner products preserved in expectation.The sketches can be accumulated from projected token-level factors without forming full gradient matrices.Source: §5.1 · p. 5
  • Two independent random sketches are multiplied to estimate squared inner products without the positive bias caused by squaring one sketch.The resulting estimator of the squared inner product is unbiased.Source: §5.1 · p. 5

5.2 Algorithm

  • One forward-backward pass yields a prior-precision estimate before fine-tuning, and the same computation applies to all LoRA ranks.If the estimated cross-example signal is non-positive, the algorithm sets precision to its maximum value.Source: §5.2 · p. 5

6 RELATED WORK

This section places the method among Bayesian LoRA, prior-predictive matching, and neural-network prior analyses. It distinguishes the paper's pre-fine-tuning, data-based precision estimate and relates it to a rank-scaling rule for preserving prior predictions.

Bayesian LoRA

  • Existing Bayesian LoRA methods use varied posterior approximations, while their prior scales are fixed in advance or tuned after training.The proposed single-task estimate is complementary and can in principle be combined with these posterior approximations.Source: §6 · p. 6
  • The derivation can be adapted to Bayesian LoRA structures that place the posterior on a small matrix within a projected low-rank subspace, but that extension is out of scope.The section also positions the approach relative to function-space methods that specify or match induced functions to a target process.Source: §6 · p. 6

Prior predictive analysis of neural networks

  • Neural-network prior analyses study output distributions under random weights, including Gaussian-process behavior at infinite width and finite-width moment properties.Signal-propagation analyses use output moments to choose initialization scales as depth grows.Source: §6 · p. 6
  • Corollary 1 gives a prior-precision rank rule that keeps Bayesian LoRA's prior predictive fixed, a rule the section says was previously absent.Rank-stabilized LoRA instead scales adapters by α/√r to keep output and gradient size unchanged during fine-tuning.Source: §6 · p. 6
  • The method propagates a factorized update prior through a second-order expansion around pretrained weights and estimates its scale from data.The analysis is at finite width.Source: §6 · p. 6

Prior predictive matching

  • Prior-predictive matching selects hyperparameters by matching predictive statistics to targets; for Bayesian matrix factorization, data-derived moments yield closed-form matching.Bayesian LoRA differs because observations depend on the update through a nonlinear pretrained network and share one draw of it.Source: §6 · p. 6

Precision selection and calibration

  • Standard Laplace prior-precision selection uses marginal likelihood or validation search, both of which require a trained model; this estimate needs neither training nor validation data.With validation data, an optional temperature fit allows it to match validation search in test NLL.Source: §6 · p. 6

7 EXPERIMENTS

Experiments first test the moment approximations and parameter recovery in controlled settings, then evaluate predictive performance, rank dependence, and OOD detection on RoBERTa-large. Across these tests, prior-predictive estimates recover prior scales and provide competitive predictive and uncertainty results, with approximation accuracy and rank-related limitations at stronger priors or ranks.

7.1 Accuracy of the second-order approximation

  • For θ̄s ≲3·10^-3, the second-order approximation has variance error below 2% for K ≥8 and mean-shift error at most 17%, with K = 1 the main exception.Beyond roughly 3·10^-3, nonlinear saturation makes the expansion overestimate both moments, producing upward bias in estimated precision; scale-recovery experiments stay in the accurate regime.Source: §7.1 · p. 7, Figure 1 · p. 7

7.2 Recovery of the prior scale

  • Among positive estimates, the recovered prior scale is within 4% of truth for regression and 6% for classification; cross-covariance estimation matches the second-moment estimator given true noise.Five percent of tasks at the smallest classification scale yield non-positive estimates and are excluded.Source: §7.2 · p. 7, Figure 2 · p. 7

7.3 Rank information in fourth moments

  • The median estimated rank is within 13% of the true rank for linearized outputs, but recovery deteriorates for stronger priors, particularly in classification.At θ̄s = 2·10^-4, the nonlinear network agrees with the linearized estimate within 3%; rank enters through a small fourth-moment correction.Source: Figure 3 · p. 8, §7.3 · p. 8

7.5 Prior precision across LoRA ranks

  • One estimate determines precision across ranks with λ̂ ∝ r^-1/2, but on RoBERTa-large it diverges from validation-NLL-optimal precision at higher ranks as NLL grows.The optimum falls close to r^-1; temperature fitting removes the loss, matching validation search with temperature at every rank.Source: §7.5 · pp. 8–9, Figure 5 · p. 9

7.4 Predictive performance

  • Without validation precision selection, the method achieves the lowest test NLL at every training-set size, with lower ECE than MAP and marginal likelihood and accuracy matching MAP.Validation search lowers NLL by 0.03–0.09 nats; temperature scaling reduces the gap to at most 0.01 nats from 256 examples onward.Source: §7.4 · p. 8, Figure 4 · p. 9

7.6 Detection of out-of-distribution inputs

  • For sentiment-task OOD detection, AUROC is 0.79–0.81 and highest at every training-set size; on AG News it is 0.83–0.90, while validation search leads by 0.01–0.02.Marginal likelihood is substantially worse in the low-data regime but reaches the highest AUROC at 4,096 examples.Source: §7.6 · p. 9

8 CONCLUSION

The paper’s closed-form estimator matches residual covariances to specify Bayesian LoRA’s prior before fine-tuning, with second-order predictions depending on rank and scale through θ = Kσ^4. On RoBERTa-large, estimated priors improve NLL and calibration over MAP and marginal-likelihood selection, while the conclusion notes limitations of the local approximation and rank estimation.

8 CONCLUSION

  • The closed-form estimator matches covariance between residuals to estimate θ = Kσ^4 and prior precision; under standard LoRA scaling, λ ∝ r^-1/2.A shared LoRA update induces covariance across example residuals, which supplies the moment-matching signal.Source: §8 · p. 9
  • On RoBERTa-large, estimated priors improve NLL and calibration over MAP and marginal-likelihood selection and provide competitive OOD detection; validation search remains slightly better in NLL.Figure 6 reports AUROC for mutual-information-based OOD separation, averaged over three seeds and, in its left panel, three tasks.Source: §8 · p. 10, Figure 6 · p. 10
  • The second-order estimator is limited to local output changes, varies across tasks given one update realization, and identifies combined scale θ rather than rank.Fourth-order information contains rank information, but a stable single-task rank estimator remains open; generative objectives and other settings are future extensions.Source: §8 · p. 10

A.1 Proof of Lemma 1

  • Entries of the update matrix can be correlated when they share a factor row or column, with the shared row of U supplying the surviving correlation in the stated case.The same-column case is the mirror image.Source: §A.1 · p. 14

A.2 From equation 4 to equation 5 and equation 6

  • To leading order, predictive covariance between inputs is θ times their gradient inner product, and predictive variance is θs(x); fourth-order terms bring rank dependence beyond θ.The covariance expansion drops terms of order θ^2.Source: §A.2 · p. 15

A.3 Kurtosis of a linearized output

  • Fourth-order corrections distinguish ranks beyond θ; the kurtosis of a linearized output also depends on the spectrum’s concentration and decreases toward the Gaussian value as rank grows.For the stated linearized output, one factor term has kurtosis 3 + 6ρ, and summing K independent terms divides excess kurtosis by K.Source: §A.2 · p. 15, §A.3 · p. 15

B ESTIMATOR INSTANCES AND DERIVATIONS

This appendix compares estimator instances and explains why simple aggregate matching can fail for classification. It also describes sketch-based computation and the Laplace predictive and evidence procedures used in pretrained-model experiments.

B ESTIMATOR INSTANCES AND DERIVATIONS

  • The naive aggregate estimator fails because single-example binary-label statistics carry only a weak first-moment trace signal, while traces can cancel on balanced tasks.On the balanced task described, the denominator retained about 0.3% of the available signal and the estimate was noise dominated.Source: B ESTIMATOR INSTANCES AND DERIVATIONS · pp. 15–16
  • In regression, independent observation noise drops out of residual pairs; in multiclass classification, pairs sharing an example are excluded to account for dependent class noise.The multiclass estimator recovers the true scale to within a few percent on synthetic data.Source: B ESTIMATOR INSTANCES AND DERIVATIONS · p. 16
  • For multiple adapters, the separated covariance system can recover two scales jointly in simulation, but is noise dominated at real-task sample sizes, making the pooled scale practical.The first-moment analogue fails when adapter trace profiles correlate across examples.Source: B ESTIMATOR INSTANCES AND DERIVATIONS · p. 16
  • A token-level generative instance would use observed-token probability, one backward pass per sequence, and pairs of positions from different sequences.The cross-sequence pairing mirrors the block exclusion used for multiclass outputs.Source: B ESTIMATOR INSTANCES AND DERIVATIONS · p. 17

C SKETCHING AND ITS GUARANTEES

  • With sketch width 48, the estimated θ is within about 10% of exact computation, and its fourth-root scale is within 3%.Independent sketch pairs remove the variance bias in squared inner products, and contractions are computed without forming an N × N matrix.Source: C SKETCHING AND ITS GUARANTEES · p. 17

D POSTERIOR APPROXIMATIONS, PREDICTIVE, AND EVIDENCE

  • The Laplace predictive uses Gaussian logits centered at MAP with covariance J(x)(Λ + λI)^-1J(x)⊤; marginal likelihood selects λ by maximizing approximate evidence.Validation search instead lowers validation NLL, while MAP uses λ = 10^8.Source: D POSTERIOR APPROXIMATIONS,… · p. 17

E TOY EXPERIMENTS

The toy experiments specify controlled network and data settings for testing second-order moments, prior-scale recovery, and rank-related behavior.

E TOY EXPERIMENTS

  • The second-order-moment experiment uses a two-layer tanh classifier with 20 inputs and 12 hidden units, adapting the first layer and estimating moments at six inputs.The moments use 40,000 draws per setting.Source: E TOY EXPERIMENTS · p. 17
  • The prior-scale study uses regression tasks with 4,000 examples and binary logistic tasks with 400,000 examples, across 80 and 96 tasks per scale, respectively.Regression label noise is set to keep the signal-to-noise ratio fixed across prior scales.Source: E TOY EXPERIMENTS · p. 18
  • The rank study uses a 112-input, 112-hidden-unit network, 24 examples averaging two inputs each, and repeats settings across ranks 10–100 and three perturbation strengths.It is repeated 10 times for each rank and perturbation setting.Source: E TOY EXPERIMENTS · p. 18

F RANK TRANSFER IN FINE-TUNING

Across ranks, the precision selected by validation search and marginal likelihood follows a slope close to −1, while the prior-based estimate follows a different scaling. The predictive variance explains why precision values can differ substantially while yielding similar variance levels.

F RANK TRANSFER IN FINE-TUNING

  • The measured slope of log λ∗ against log rank is close to −1 for each task and for validation-search and marginal-likelihood precisions.The rank study used AG News and CR; Table 2 reports λ∗, ˆλ and test NLL at each rank.Source: F RANK TRANSFER IN FINE-TUNING · p. 18, Table 2 · p. 18
  • Precisions chosen by the estimator and validation search can differ by up to 30 times while their products λω agree within 10%.Across ranks, rλω shows no trend; ω is the average predictive logit variance over test examples and classes.Source: F RANK TRANSFER IN FINE-TUNING · p. 18
  • The linearized predictive variance scales as 1/(rλ), accounting for the measured rank slope near −1.The appendix relates this scaling to the variance of the linearized predictive.Source: F RANK TRANSFER IN FINE-TUNING · p. 18

G RANK SCALING OF THE PRIOR AND OF THE LINEARIZED PREDICTIVE

The prior and the trained-factor linearized predictive have different rank scalings because they retain different terms in the low-rank update expansion. The prior’s rank scaling is r^-1/2, whereas the precision that maintains the best linearized-predictive variance falls as 1/r.

G RANK SCALING OF THE PRIOR AND OF THE LINEARIZED PREDICTIVE

  • Under the prior, zero-mean factors leave only the bilinear update term, whose variance yields λ ∝ r^-1/2 when held fixed.The prior’s factors are centered at zero, unlike the trained factors used in the Laplace predictive.Source: G RANK SCALING OF THE PRIOR AND OF THE… · p. 19
  • The trained-factor linearized predictive retains the middle term and omits the bilinear term, producing variance proportional to 1/(rλ).This scaling holds when trained-factor entry sizes do not depend on rank, as long as fine-tuning does not change their scale with rank.Source: G RANK SCALING OF THE PRIOR AND OF THE… · p. 19
  • The precision that maintains the best variance level for the linearized predictive therefore falls as 1/r, rather than following the prior’s r^-1/2 scaling.The estimated prior precision moves away from the best linearized-predictive precision by a factor growing as √r when used directly for that predictive.Source: G RANK SCALING OF THE PRIOR AND OF THE… · p. 19
  • A posterior-sampled predictive retains the bilinear term, so its best precision would scale between r^-1 and r^-1/2 depending on which variance component dominates.The authors did not evaluate this predictive.Source: G RANK SCALING OF THE PRIOR AND OF THE… · p. 19

H ROBERTA-LARGE EXPERIMENTS

The RoBERTa-large experiments compare Bayesian LoRA methods across four classification tasks and examine the effect of training under the estimated prior. Training with that prior improves calibration, while marginal likelihood can refine its precision without validation data when the model was trained under the prior.

H ROBERTA-LARGE EXPERIMENTS

  • The RoBERTa-large evaluation uses SST-2, MR, CR and AG News, with task test sets and training subsets drawn from the stated datasets.The task setup adapts query and value projections across all 24 layers; the supplied result-table captions identify task results but do not provide cell values.Source: H ROBERTA-LARGE EXPERIMENTS · p. 20, Table 3 · p. 21, Table 4 · p. 21

I FINE-TUNING UNDER THE ESTIMATED PRIOR

  • Training under the estimated prior adds the prior’s negative log density to training NLL, making the trained model the posterior mode under that prior.The prior is N(0, ˆλ^-1I) over LoRA parameters.Source: I FINE-TUNING UNDER THE ESTIMATED PRIOR · p. 25
  • Training LoRA under the estimated prior lets marginal likelihood select precisions close to the estimate and match its NLL and calibration without validation data.With LoRA trained without a prior, marginal likelihood instead selects precisions orders of magnitude below the estimate; under-prior training satisfies the evidence method’s mode assumption.Source: I FINE-TUNING UNDER THE ESTIMATED PRIOR · p. 25
  • The estimated prior improves calibration: MAP and marginal-likelihood reliability move toward the diagonal and both methods’ ECE drops sharply.With MAP already well calibrated, methods reach similar NLL; the proposed method maintains its NLL and slightly improves calibration.Source: I FINE-TUNING UNDER THE ESTIMATED PRIOR · p. 25, Figure 10 · p. 26

Related papers

Loading 2610.35457v1…