Hello everyone, this is Goodnote. This article provides a detailed introduction to the SD model fine-tuning method LoRA, including dataset preparation, model fine-tuning process, inference process, advantages and disadvantages, etc.

Paper

Paper URL:LoRA: Low-Rank Adaptation of Large Language Models ^[1]

Concept

LoRA (Low-Rank Adaptation of Large Language Models) is a technique for efficiently fine-tuning large pre-trained models, particularly suitable for models with a huge number of parameters, such as GPT-3, BERT or Stable Diffusion. LoRA provides a solution that allows only a very small number of parameters to be fine-tuned while maintaining performance comparable to full parameter fine-tuning.

Core Principle

LoRA (Low-Rank Adaptation of Large Language Models) is based on the principle of low-rank decomposition to reduce the number of parameters that need to be updated during fine-tuning of large models, achieving more efficient model fine-tuning. It is particularly suitable for pre-trained large models like GPT-3, BERT or Stable Diffusion. During fine-tuning of these models, LoRA introduces two low-rank matrices to replace the direct update of the entire weight matrix, significantly reducing computational costs and storage requirements.

Core Idea: Assume the weight matrix W in the pre-trained model is very large. The basic idea of LoRA is to decompose its update part ΔW into the product of two smaller rank matrices A and B, thereby reducing the number of parameters that need to be updated and only updating these two matrices. Instead of updating the entire W matrix, it maintains the original performance of the pre-trained model.

Core Formula for Weight Matrix Decomposition:

LoRA processes the weight matrix as follows:

Where:

• is the original weight matrix in the pre-trained model, which is not updated.
• is the trainable part introduced by LoRA. To reduce training parameters, LoRA represents as the product of two low-rank matrices:.
• and are two low-rank matrices, where:
• (small-sized matrix, rank far less than ).
• (also a small-sized matrix).

In LoRA, only the matrices and are trained, rather than training the entire weight matrix , thus reducing the number of parameters that need to be updated.

Generally speaking, for matrix A, we use random Gaussian distribution to initialize, and for matrix B we use all zeros initialization, ensuring that the result of multiplying these two matrices at the initial training stage is zero. This guarantees that at the initial stage of training, the weights of the SD model are fully effective.

Small Tasks: Generally, r is taken as 1 to 4.

Medium Tasks: Usually, r is taken as 4 to 8.

Complex Tasks: Generally, r is taken as 8 to 16.

Conclusion Source: A comprehensive analysis of LoRA (Low-Rank Adaptation) model core knowledge ^[2]

Advantages

• Significantly reduce the amount of parameter updates: By performing low-rank decomposition, LoRA only needs to update a very small number of parameters, rather than the entire model parameters. This greatly reduces storage requirements and computational costs.
• Maintain the original performance of the model: LoRA freezes the weights of the SD model, so the fine-tuned model can still utilize the knowledge of the pre-trained model.
• Modular updates and pluggable design: LoRA’s weight adjustments are made through A and B, and can be easily applied during inference without needing to reload the entire model. The A and B matrices trained by LoRA can be saved independently. During inference, you can quickly switch between different tasks without retraining the model.

Training Process

The training process for LoRA fine-tuning can be divided into the following steps:

Loading Pre-trained Model

First, load a pre-trained Stable Diffusion model. This model has been trained on a large-scale text-image dataset and possesses strong generative capabilities. The goal of LoRA is to fine-tune this pre-trained model, requiring only a minimal adjustment of parameters.

Selecting Layers for Fine-Tuning

In the Stable Diffusion model, LoRA typically operates on the linear layers (mainly on the three weight matrices W of the attention module: Q, K, V). Specifically, LoRA will adjust the cross-attention mechanism in U-Net, self-attention mechanism linear layers, etc.

• LoRA adjusts the weights in these linear layers through low-rank decomposition (mainly in the attention module’s Q, K, V weights), introducing two low-rank matrices A and B to adjust weight changes.

Detailed Introduction as Follows:

LoRA mainly acts on the attention mechanism-related linear layers in the U-Net network, particularly for cross-modal interaction in the Cross-Attention layer and Self-Attention layer.

Layers Optimized by LoRA

In the Stable Diffusion U-Net, there are mainly two types of attention mechanisms: Self-Attention and Cross-Attention. LoRA mainly optimizes the linear layers in these attention layers, especially the following parts:

Cross-Attention Layer

• Cross-Attention is the core mechanism in U-Net, responsible for combining the latent features of images (Latent Feature) with text embeddings (Text Embedding). Through the cross-attention mechanism, the model can extract information from text prompts to adjust the content of image generation.
• In Cross-Attention , LoRA mainly targets the linear transformation layers used to generate Query (Q), Key (K) and Value (V) 【】for optimization. These layers process input features and text embeddings and compute attention weights.

Self-Attention Layer

• Self-Attention is the internal image feature interaction mechanism in U-Net, responsible for feature interactions between different regions of the image, capturing the global dependencies of the image.
• In Self-Attention, LoRA also optimizes the linear layers that generate Query (Q), Key (K) and Value (V) .

How LoRA Optimizes These Layers

LoRA reduces the number of parameter updates by performing low-rank decomposition on the linear layers in the attention mechanism (e.g., the linear transformation layers used to generate Q, K, V ). The specific optimization method is as follows:

Operations of Original Linear Layer

In the attention mechanism, input features are typically mapped through linear layers to generate

• input features will be mapped through linear layers to generate Query (Q), Key (K) and Value (V), with the following formula:
• where , , are weight matrices used to generate , , . In standard fine-tuning processes, these weight matrices are updated.

LoRA’s Low-Rank Matrix Decomposition

Compared to directly updating the entire , , matrix, in LoRA , the changes in these matrices are represented by using two smaller matrices A and B:

• Introducing low-rank decomposition to adjust the weight matrix: , , is represented as pre-trained weights plus a low-rank matrix adjustment:

  Similar formulas are used for and .

• and are two low-rank matrices, with rank far less than the input and output dimensions. Thus, LoRA only needs to train and , reducing the number of parameters that need to be trained.
• During the inference phase, the updated weight is the original weight plus .

Key Steps in Optimization

• Keep Original Weights Unchanged: In LoRA fine-tuning, the original weight matrices , , remain unchanged, meaning that the knowledge of the pre-trained model is not destroyed.
• Only Train Low-Rank Matrices and : By training and (and the corresponding , and , ), LoRA fine-tuning can adapt to the needs of new tasks while significantly reducing the amount of parameter updates.

How to Apply Optimization During Inference?

In the inference phase, the optimized part of LoRA is applied to the model as follows:

• Apply Weight Updates: During inference, LoRA uses the fine-tuned low-rank matrices A and B, and adds them to the original weight matrix. For example, calculating Q :

• This updated weight matrix ensures that the model can generate images containing customized content while maintaining its original generative capability.

Why Choose to Optimize These Layers?

LoRA chooses to optimize the linear layers in the Attention layers for the following reasons:

1. Attention Layers are Core: In the U-Net of Stable Diffusion, the Attention layers (especially the Cross-Attention layer) are key to combining image generation with text prompts. Optimizing these layers can directly affect how the model conveys text information to the image generation process.

2. Effectiveness of Low-Rank Decomposition: For fully connected layers in large models, low-rank decomposition can significantly reduce parameter updates without significantly affecting model performance. The weight matrices in attention layers are often high-dimensional, so applying LoRA to them can greatly lower training costs.

3. Reduce Computational Overhead: The linear layers in attention mechanisms usually have a large number of parameters and are computationally intensive. Through LoRA’s low-rank decomposition, it is possible to maintain the model’s generative capability while reducing computational complexity.

LoRA Makes Local Adjustments to Some Cross-Attention and Self-Attention Layers, Which Layers are Primarily Adjusted?

Indeed, U-Net contains multiple Self-Attention and Cross-Attention layers at different resolutions, but LoRA fine-tuning is selective, and does not fine-tune every Attention layer. The design intention of LoRA is to fine-tune only the critical Attention layers through low-rank matrix decomposition.

LoRA typically prioritizes fine-tuning in the following parts, rather than adjusting all layers in the entire U-Net:

1. Attention Layers in the Bottleneck Part

• Reason: The bottleneck part is located between the encoder and decoder and is the layer with the lowest resolution in the entire U-Net. The features in this part represent global high-level information and abstract concepts. Fine-tuning the Cross-Attention layers in this part is crucial because this attention is responsible for combining text embeddings (from CLIP) with latent features of images. Since this part affects the global structure and semantics of the entire generation process, fine-tuning the Cross-Attention layers here maximizes the response to text prompts.

2. High-Resolution Attention Layers in the Decoder

• Reason: The high-resolution layers in the decoder are located in the layers close to the output. These layers progressively map latent features back to high-resolution images, responsible for generating the details of the final image. Fine-tuning these layers’ Self-Attention and Cross-Attention mechanisms helps ensure the details and texture quality of the image. LoRA selectively fine-tunes these layers to maintain the accuracy of generated details.

3. Some Intermediate Layers

• Reason: Intermediate layers appear in the middle positions of both the encoder and decoder parts, that is, the transition layers when the encoder gradually reduces resolution or when the decoder gradually upsamples. In certain tasks, particularly those requiring global consistency, LoRA may adjust some intermediate layers. This typically targets only certain key intermediate layers, rather than covering all intermediate Attention layers.

Other layers will not be changed. DreamBooth will modify all Attention layers.

Low-Rank Decomposition

For each layer that needs fine-tuning, LoRA performs low-rank decomposition on the weight matrix W of that layer:

Where:

• is the weight matrix of the pre-trained model, which remains unchanged.
• is the part adjusted by LoRA, only training the low-rank matrices and .

Training Low-Rank Matrices of LoRA

During the training process of LoRA, we only train the two low-rank matrices A and B, while keeping the original weights in the model unchanged. This method significantly reduces the number of parameters that need to be updated, lowering training costs.

The specific process is as follows:

1. Image and Text Input: Train the model using data pairs containing images and text descriptions, similar to the pre-training tasks of Stable Diffusion.

2. Loss Function: The loss function of LoRA is similar to that of the original Stable Diffusion, primarily training through predicting noise residuals. The loss function is usually Mean Squared Error (MSE), with the following formula:

• where is the original image, is the noisy image after diffusion steps, is the text condition, is the noise.
• Minimize the difference between the predicted noise residuals by the model and the real noise through the loss function.

Only Update Low-Rank Matrices: During backpropagation, LoRA only updates A and B, while keeping the rest of the pre-trained model unchanged. Through this method, the model can learn new tasks with minimal parameter updates.

More detailed processes are as follows:

Forward Pass

Input Data Preparation:

• Input includes images and text descriptions, which will be feature-extracted through the text encoder (like CLIP) and Stable Diffusion’s U-Net.

Weight Replacement:

• In LoRA, the original pre-trained weight matrix will not be directly updated but will be expanded as:
• where is the original weight matrix, which remains unchanged, and and are the introduced low-rank matrices, respectively and dimensions, is the low-rank dimension.
• Thus, the linear layer weights in U-Net (especially in Cross-Attention and Self-Attention for Q, K, V weights) will be replaced by in the forward pass to complete the linear transformation.

Generating Q, K, and V Matrices:

• For the input image features and text embeddings, the model generates Q, K, and V matrices through linear layers:

• Here 、 and represent the linear layer weights of the Attention layers. In LoRA fine-tuning, these weights are expanded as:

• During the forward pass, these new Q, K, V weights are used to compute the Attention output.

Calculating Attention Scores:

• After obtaining the Q and K matrices, compute the Attention scores through dot product and use Softmax for normalization to determine the relevance between image features and text features:
• where is 、is the dimension for scaling to stabilize gradients. This step assigns higher weights to important features.

Applying Attention Scores to Value Matrices:

• Apply the Attention scores to matrices to generate the output features after attention aggregation:

• This output contains important information extracted from the input features, with the feature details of the new task.

Attention Output Used for Noise Prediction【Same as DDPM Diffusion Process】

DDPM principle reference: Stable Diffusion notes collection in<span> "Diffusion Model Principle" </span>.

Loss Calculation:

• Based on the output of the forward pass, calculate the difference between the predicted noise residuals and the target using the loss function (usually mean squared error MSE):

• where is the original image, is the noisy image after diffusion steps, is the text condition, is the noise.
• Minimize the difference between the predicted noise residuals by the model and the real noise through the loss function.

Backward Pass

Calculating the Gradient of Loss:

• Calculate the gradient of the loss with respect to A and B based on the loss function L:

and

Lock Original Weight Matrices :

• During backpropagation, LoRA only updates and , while keeping the original weight matrix unchanged. This ensures that the pre-trained knowledge is not destroyed, while adapting to new task features through updates to A and B.

Updating Low-Rank Matrices A and B:

• Based on the gradients, update the parameters of A and B using an optimization algorithm (like Adam):

• where is the learning rate, and the updated and allow the model to adapt to new task features without altering the pre-trained weights.

Preventing Overfitting

To avoid the model overfitting to the new dataset during fine-tuning, LoRA introduces regularization techniques and typically limits the number of parameters in low-rank matrices by selecting a smaller rank . A smaller rank r can ensure that the model while learning new tasks does not lose its original generative ability, while avoiding overfitting.

• Small Tasks (simple image generation, basic text tasks): Generally, r is taken as 1 to 4. For simple tasks, this rank can usually capture enough feature changes and significantly reduce fine-tuning parameters.

• Medium Tasks (e.g., generating images in specific styles or scenes): Usually, r is taken as 4 to 8. This rank setting can adapt to slightly more complex tasks while maintaining fine-tuning efficiency.

• Complex Tasks (high-resolution image generation, large-scale customization): Generally, r is taken as 8 to 16. For very complex generation tasks, a higher rank can capture more detailed features but will increase the number of parameters and computational costs.

Inference Process

The inference process is similar to the standard Stable Diffusion inference process, but will add the adjustments of LoRA on top of the original model.

The inference flow is as follows:

Loading the Fine-Tuned Base Model and LoRA Weights

• In the inference phase, we need to load the pre-trained Stable Diffusion model and the fine-tuned LoRA weights (i.e., values of A and B). These weights are typically small storage files, easy to load and apply.
• The model’s inference process will utilize the original weights W0 and the fine-tuned ΔW=AB.

Input Text Conditions

• Users can input a text description as a prompt for the Stable Diffusion model to generate images. This text description will be encoded into text embeddings through the CLIP model and used as a condition input for the generation process.

Image Generation

• During inference, LoRA generates the final image by adding the trained ΔW to the original weights W0:

• where is the input latent vector or image feature, is the generated result after LoRA fine-tuning.

The model will go through a series of reverse diffusion steps to gradually remove noise and generate images that meet the input text description.

Control and Adjustment

• By adjusting the input text prompts, users can flexibly control the style, content, and other details of the generated images.
• LoRA provides a modular design, allowing for easy replacement of different LoRA fine-tuning weights to adapt to different tasks and datasets. For example, multiple different LoRA fine-tuning weights can be applied to the same model to generate images of different styles or content.

Recommended Reading

Efficient Parameter Fine-Tuning Using LoRA for Stable Diffusion^[3]

LoRA Training Scripts from Around the World, Unite!^[4]

LoRA Using Scripts to Train Unet and Text Encoder Throughout^[5]

Using set_adapters to Load LoRAS for Inference, Customize Adapter Strength^[6]

Load adapters, load Dreambooth, Textual inversion, LoRA, IP-Adapter:

• https://huggingface.co/docs/diffusers/using-diffusers/loading_adapters#LoRA^[7]

• https://huggingface.co/docs/diffusers/using-diffusers/merge_loras^[8]

Note:

• LoRA can be used alongside other methods, such as textual inversion and dreambooth.
• LoRA primarily modifies the unet when changing the model, but can also fine-tune the text encoder.
• LoRA checkpoints in the diffusion community are almost always obtained through DreamBooth^[9].

References

[1]

LoRA: Low-Rank Adaptation of Large Language Models: https://arxiv.org/abs/2106.09685

[2]

A Comprehensive Analysis of LoRA (Low-Rank Adaptation) Model Core Knowledge: https://zhuanlan.zhihu.com/p/639229126

[3]

Efficient Parameter Fine-Tuning Using LoRA for Stable Diffusion: https://huggingface.co/blog/lora

[4]

LoRA Training Scripts from Around the World, Unite!: https://huggingface.co/blog/sdxl_lora_advanced_script

[5]

LoRA Using Scripts to Train Unet and Text Encoder Throughout: https://huggingface.co/docs/diffusers/training/lora

[6]

Using set_adapters to Load LoRAS for Inference, Customize Adapter Strength: https://huggingface.co/docs/diffusers/main/en/tutorials/using_peft_for_inference

[7]

https://huggingface.co/docs/diffusers/using-diffusers/loading_adapters#LoRA: https://huggingface.co/docs/diffusers/using-diffusers/loading_adapters#LoRA

[8]

https://huggingface.co/docs/diffusers/using-diffusers/merge_loras: https://huggingface.co/docs/diffusers/using-diffusers/merge_loras

[9]

DreamBooth: https://huggingface.co/docs/diffusers/main/en/training/dreambooth

Popular Articles Review

Welcome to subscribe to our collection of AI foundational knowledge, the collections of Machine Learning Notes and Deep Learning Notes have been basically updated, welcome click the image below to enter the collection directory.

Currently, the collection of Multimodal Paper Notes related to image generation has been basically updated (subsequent updates will include video generation-related papers), welcome click the image below to enter the collection directory.