Hypothesis:

The purpose is to use transfer learning to adapt a pre-trained generative model (e.g., Stable Diffusion) to scientific machine learning tasks, specifically for data assimilation and uncertainty quantification. By leveraging the low-rank adaptation capabilities of LoRA, we can fine-tune the model to generate samples that are consistent with the underlying physics of the system being studied. The transfer learning approach allows us to efficiently adapt the model to new scientific domains without requiring extensive retraining, making it a powerful tool for scientific machine learning applications. We know that LoRA has been used for fine-tuning text-to-image generation, but we want to explore its potential for scientific machine learning tasks, where the goal is to generate samples that are consistent with the underlying physics of the system being studied.

Given a pre-trained generative model (e.g., Stable Diffusion), we can use LoRA to learn a low-rank adaptation that captures the dynamics of a scientific system, enabling the model to generate samples that are consistent with the underlying physics.

#Bayesian LoRA: A generative model for scientific machine learning

$$y=W_o X + \theta_B \theta_A X$$

Where $\theta_A$ and $\theta_B$ are the low-rank adaptation matrices learned by LoRA. Small amount of data from the scientific domain of interest can used to train these parameters making them consistent with the underlying physics of the system being studied.

From a geometrical point of view, standard LoRA constrains the weight update to a low-dimensional subspace of dimension $r \ll k$: the update $θ_B θ_A X$ first projects the input onto the $r-$dimensional row space of $θ_A$, and then lifts the result back to the k-dimensional output space via $θ_B$.

By contruction, this allow for the expressivity of the learned adaptation post fine-tuning to be restricted to the column space of $θ_B$, denoted V = col($θ_B$) ⊂ $R^k$.

We decide to inject noise via $\theta_A$ via the following way:

$$\theta_A \sim \mathcal{N}(\theta_A, \alpha(X) \theta_A^2)$$

where $\alpha(X)$ is a function that controls the variance of noise injected into the low-rank adaptation matrix $\theta_A$ based on the input $X$.

The predictive distribution of the model can then be expressed as:

$$p(y) = \mathcal{N}(y | W_o X + \theta_B \theta_A X, \alpha(X) \theta_B \text{ diag }(\theta_A^ X^2 )\theta_B^T)$$

This when we are optimizing

Literature Review:

  • LoRA: Low-Rank Adaptation of Large Language Models by Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. This paper introduces the concept of low-rank adaptation (LoRA) for fine-tuning large language models, which can be applied to generative models for scientific machine learning tasks.