Introduction

Fine-tuning large-scale generative models like Stable Diffusion 1.5 using DreamBooth requires meticulous calibration of hyperparameters, particularly the learning rate (LR) schedule. An optimal LR schedule ensures efficient convergence, minimizes overfitting, and preserves the model’s ability to generalize. Traditional LR schedules, such as constant, linear decay, or cosine annealing, often fall short in addressing the unique challenges posed by DreamBooth training, especially when dealing with limited data and the need for precise subject representation.

Enter RISE—Rise Inverse Stable Evolution—a custom-designed LR scheduler tailored specifically for DreamBooth fine-tuning. RISE introduces a multi-phase approach that combines warmup, plateau, and a smooth decay, optimizing the training dynamics for better performance and stability.

Origins of RISE

The Rationale Behind RISE

DreamBooth fine-tuning involves adapting a pre-trained model to generate images of a specific subject using a limited number of examples. This process is susceptible to overfitting, especially if the learning rate is not carefully managed. High learning rates can lead to rapid convergence but risk overshooting the optimal solution, while low learning rates may result in prolonged training times and underfitting.

RISE addresses these challenges through a structured LR schedule that:

• Warmup Phase (0–20%): Gradually increases the LR from a minimal value to the maximum specified LR. This phase allows the model to adjust to the new data distribution without abrupt weight updates, reducing the risk of early overfitting.
• Plateau Phase (20–40%): Maintains the LR at its peak value, enabling the model to make significant progress in learning the new subject representation.
• Decay Phase (40–100%): Applies a smooth, cosine-based decay to the LR, gradually reducing it to half of the maximum value. This phase ensures fine-tuning of the model weights, promoting stability and preventing overfitting in later stages of training.

Implementation Details

Plot Source Code

# Re-import necessary packages after code execution state reset
import numpy as np
import matplotlib.pyplot as plt
import math

# Parameters
steps = 1000
max_lr = 1.5e-6
min_lr = max_lr / 2

# Build LR schedule with cosine easing decay
lrs = []
for i in range(steps):
    pct = i / steps

    if pct < 0.20:
        # Stage 1: Linear warmup
        lr = min_lr + (max_lr - min_lr) * (pct / 0.20)
    elif pct < 0.40:
        # Stage 2: Constant max_lr
        lr = max_lr
    else:
        # Stage 3: Cosine decay
        t = (pct - 0.40) / (1.0 - 0.40)
        cosine_t = 0.5 * (1 + math.cos(math.pi * t))  # eased in/out
        lr = min_lr + (max_lr - min_lr) * cosine_t

    lrs.append(lr)

# Plot
plt.figure(figsize=(10, 5))
plt.plot(lrs, label="Learning Rate")
plt.xlabel("Training Step")
plt.ylabel("Learning Rate")
plt.title("RISE (Rise Inversion Stable Evolution) algorithm")
plt.grid(True)
plt.legend()
plt.tight_layout()
plt.show()

This implementation ensures a seamless transition between phases, promoting stable and efficient training dynamics.

Advantages of RISE Over Traditional Schedulers

• Tailored for DreamBooth: Unlike generic LR schedules, RISE is specifically designed to address the nuances of DreamBooth fine-tuning, such as limited data and the need for precise subject representation.
• Smooth Transitions: The use of a cosine-based decay ensures gradual transitions between phases, reducing the likelihood of abrupt changes that could destabilize training.
• Flexibility: RISE allows for customization of phase duration’s and LR values, enabling practitioners to adapt the schedule to various datasets and training objectives.
• Enhanced Stability: By preventing sudden LR drops and incorporating a warmup phase, RISE promotes stable convergence and reduces the risk of overfitting.

Empirical Observations

In practice, applying RISE to DreamBooth training with Stable Diffusion 1.5 has demonstrated:

• Improved convergence rates compared to constant or linear decay schedules.
• Enhanced image quality, with better preservation of subject details and reduced artifacts.
• Greater stability during training, with fewer instances of divergence or mode collapse.

These observations align with findings from the broader community, where adaptive LR schedules have been shown to improve fine-tuning outcomes in generative models.

Conclusion

The RISE scheduler offers a robust and flexible approach to managing the learning rate during DreamBooth fine-tuning of Stable Diffusion 1.5. By incorporating a structured progression through warmup, plateau, and decay phases, RISE addresses the unique challenges of this training paradigm, promoting stability, efficiency, and high-quality outputs.

For practitioners seeking to optimize their DreamBooth training workflows, integrating RISE into the training loop can yield significant benefits. As with any training strategy, it’s advisable to experiment with phase duration’s and LR values to best suit the specific characteristics of your dataset and objectives.

Implementation & Source Code

The RISE scheduler is available in the following GitHub repository, maintained as a fork of the DreamBooth extension for Automatic1111’s Web UI:

🔗 GitHub: Anime4000/sd_dreambooth_extension

To use the RISE scheduler:

🛠️ Installation

git clone https://github.com/Anime4000/sd_dreambooth_extension --branch rise

Then, inside your environment:

pip install -r extensions\sd_dreambooth_extension\requirements_strict.txt

Once installed, you can select RISE from the learning rate scheduler dropdown in the UI or define it in your training config JSON:

    "lr_scheduler": "rise_inverse_stable_evolution",

Training Results

Base Checkpoint: Anything-V3.0-pruned.ckpt

Checkpoint Download: https://civitai.com/models/18729?modelVersionId=1685718

Attempt #1

Attempt #2

Hentai

Futanari (Hermaphrodite)