Package: bayesian-lora
Bayesian LoRACode for the paperBayesian Low-Rank Adaptation for Large Language Models.See the explanatoryblog postanddocumentation.Installationpipinstallbayesian-loraExampleWe provide a comprehensive example inexamples/example_usage.py, running
through the main methods using Phi-2 on ARC-E.Note that running this requires a local installation with a few extra
dependencies. Run:gitclonehttps://github.com/MaximeRobeyns/bayesian_loracdbayesian_lora
pipinstall-e".[examples]"and thenpython./examples/example_usage.pyThe main functions this library provides are for calculating Kronecker factors,
the marginal likelihood, and the posterior predictive distribution. We show how
to use these in the examples below.Calculating (low-rank) Kronecker factorsFirst, wrap your model call in a function that takes a batch from your data
loader, and returns the relevant logits. For a CausalLM from HuggingFace:deffwd_call(model:nn.Module,batch_prompts:Any)->t.Tensor:inputs=tokenizer(batch_prompts).to(device)outputs=model(**inputs)logits=outputs.logits[:,-1]# Get the last token logitsreturnlogitsYou can now call ourcalculate_kronecker_factorsfunction:frombayesian_loraimportcalculate_kronecker_factorsfactors=calculate_kronecker_factors(model,# Your model (not necessarily PEFT)fwd_call,# Model call wrapper, defined abovetrain_loader,# Your training data loadercfg.n_kfac,# (Optional) rank to usecfg.lr_threshold,# (Optional) threshold for low-rank approximation["lora"],# modules to targetuse_tqdm=True,# (Optional) use tqdm for progress bar)In the above, the["lora"]argument contains a case-insensitive list of
keywords to identify modules to target. Since we're working with a LoRa model,
we choose"lora"to target (e.g.layers.0.q_proj.lora_A, etc).Thefactorsare a dictionary with keys being the full name of the targetted
modules, and a tuple of two tensors as the values: the first being the
(possibly low-rank) Kronecker factor corresponding to the input activations,
and the second being the (possibly low-rank) factor corresponding to the output
gradients.Seethe K-FAC docsfor more detail.Model EvidenceWe provide a function calledmodel_evidencewhich returns the evidence /
marginal likelihood.frombayesian_loraimportmodel_evidenceevidence=model_evidence(model,# Your modellog_likelihood,# A Tensor with model's log likelihood on some eval datasetfactors,# Kronecker factors, as calculated aboven_lora,# rank used in the LoRA adaptersn_kfac,# rank used in the Kronecker factorsprior_var,# prior variance hyperparameter, as a tensor)You can then useevidenceas the loss in a normal training loop, presuming
your parameters (e.g.prior_varhave gradients).Posterior Predictive DistributionTo get the parameters of the Gaussian over the logits, use
thejacobian_meanandvariancefunctions.witht.no_grad():forbatchinvalidation_loaderprompts,classes=batchbatch_inputs=tokenizer(prompts)# Predict the output logit locations# target_ids is a tensor containing the indices of the target tokens# e.g. [354, 355, 356].jacobian,f_mu=jacobian_mean(model,batch_inputs,target_ids)# Predict the output logit variancesf_var=variance(batch_inputs,# inputsjacobian,# the Jacobian dictionary, obtained abovefactors,# Kronecker factors, as calculated aboveprior_var,# prior variance hyperparameter, as a tensorclasses.size(-1),# number of classes to predictn_lora,# rank of the LoRA adaptersn_kfac,# rank of the Kronecker factorsdevice,# device to use)# Now use the parameters to e.g. sample logits from the Gaussian# predictive, parametrised by f_mu, f_varL=t.linalg.cholesky(f_var)samples=100_000f_mu=f_mu.expand(samples,*f_mu.shape)L=L.expand(samples,*L.shape)eps=t.randn_like(f_mu)logits=(f_mu+L@eps).squeeze(-1).softmax(-1).mean(0)The above is a minimal example; seethis
sectionof the documentation for more detail.DevelopmentThis library is intentionally very small and hackable. It has two main files,
and three dependencies (torch,tqdmandjaxtyping.)main.pycontains methods specific tothe paper,kfac.pycontains relatively portable K-FAC methodsFeel free to directly copy the code into your projects and hack on it.