Light LoRAについて

# 0. Light LoRAについて
[HyperDreamBooth](https://arxiv.org/pdf/2307.06949v1.pdf)のアイデアにLightweight DreamBoothというLoRAよりも更にパラメータ数を削減するアイデアが書かれている。


従来のLoRAは$W_0=(768,320)$の行列に対して、これの差分行列$W=W_0 + \Delta W$を考える時、$\Delta W$を$(768,4)$と$(4,320)$の2層の差分モデルで考える。
出力次元は$320,640,1280$のいずれかでこれはUnetの解像度に依存する。
入力次元は$K,V$のCross-AttentionであればTextEncoderの次元で$768$次元。$Q,Out$のCross-AttentionかSelf-Attentionであれば出力次元と等しい。

Lightweight Dreambooth(Light LoRA)は以下の様に4層のモデルを考える。その内、2層は凍結した重みである。これは例えば$W_0=(768,320)$の行列に対して$\Delta W$を$(768,100)$と$(100,4)$と$(4,50)$と$(50,320)$の4つの行列で計算し、その内、$(768,100)$と$(50,320)$は初期重みで重みを凍結する。学習パラメータは入出力次元に依らず$(100,4)$と$(4,50)$だけでよく、このパラメータはLoRAモデルの1/10以下である。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/3f450a3b-e0bc-b873-ec58-3d58d69e741f.png)

lightLoRAの実装例は

https://github.com/KohakuBlueleaf/HyperKohaku

にあるが、down_aux、up_auxの重みを学習時に毎回代入しているように見える。
また、Hypernetworkで画像入力からdownとupの重みを予測するためか、downとup重みを結合している。
上述の実装手法を参考にしつつ単にLightLoRAを学習するだけの実装してみる。


# 1. up_loraだけzero行列で初期化
以下の様に4つの行列を考えて、そのうちdown_aux、up_auxの重みをorthogonal_重みで固定する。
`C:\ProgramData\Anaconda3\Lib\site-packages\diffusers\models\attention_processor.py`のLoRALinearLayerを改造する。
あとは`diffusers/examples/dreambooth/train_dreambooth_lora.py`で学習させた。

```python
class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, network_alpha=None):
        super().__init__()
...
        self.down_aux = nn.Linear(in_features, 100, bias=False)
        self.down = nn.Linear(100, rank, bias=False)
        self.up = nn.Linear(rank, 50, bias=False)
        self.up_aux = nn.Linear(50, out_features, bias=False)

        self.network_alpha = network_alpha
        self.rank = rank

        seed = 0
        torch.manual_seed(seed)
        nn.init.orthogonal_(self.down_aux.weight, gain=1)
        nn.init.orthogonal_(self.down.weight, gain=1)
        nn.init.zeros_(self.up.weight)
        nn.init.orthogonal_(self.up_aux.weight, gain=1)

        self.down_aux.weight.requires_grad=False
        self.up_aux.weight.requires_grad=False

    def forward(self, hidden_states):
        orig_dtype = hidden_states.dtype
        dtype = self.down.weight.dtype

        down_hidden_states = F.linear(hidden_states.to(dtype), self.down.weight @ self.down_aux.weight)
        up_hidden_states = F.linear(down_hidden_states, self.up_aux.weight @  self.up.weight)

        if self.network_alpha is not None:
            up_hidden_states *= self.network_alpha / self.rank
        return up_hidden_states.to(orig_dtype)
```
**load_lora_weights**は上手く動かないので以下の様にしてLight LoRAを従来LoRAの重みに変換してモデルに読み込む。

保存される重みは`lora.down_aux.weight, lora.up_aux.weight, lora.down.weight, lora.up.weight`の4種なので生成時に行列積を計算して`lora.down.weight, lora.up.weight`を計算する。実際にはdown_auxとup_aux重みは不要だが確認のため保存する。


更に余計な要素を削除するため
`state_dict = {module_name: param for module_name, param in state_dict.items() if 'lora.up.weight' in module_name or 'lora.down.weight' in module_name}`で
state_dict内の**lora.down.weight, lora.up.weight**以外の要素を削除する。

最後に`pipe.unet.load_attn_procs(state_dict)`でLoRA重みは適用されるはず。

```python
model_id = "./cat_1_TI/"
lora_path = './cat_1_TI_Li/pytorch_lora_weights.bin'
ddim = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id,torch_dtype=torch.float16,scheduler=ddim).to("cuda")

state_dict = torch.load(lora_path, map_location="cpu")

for name in list(state_dict.keys()):
    if 'unet' in name:
        if 'down_blocks.0' in name:
            hidden_dim = 320
        if 'down_blocks.1' in name:
            hidden_dim = 640
        if 'down_blocks.2' in name:
            hidden_dim = 1280
        if 'up_blocks.1' in name:
            hidden_dim = 1280
        if 'up_blocks.2' in name:
            hidden_dim = 640
        if 'up_blocks.3' in name:
            hidden_dim = 320
        if 'mid_block' in name:
            hidden_dim = 1280

        if 'attn2.processor' in name:
            # Cross-Attention
            if 'to_k' in name or 'to_v' in name:
                input_dim = 768
            else:
                input_dim = hidden_dim
        else:
            # Self-Attention
            input_dim = hidden_dim

        output_dim = hidden_dim
        
        embedding1 = nn.Embedding(input_dim, 100)
        embedding2 = nn.Embedding(50, output_dim)
        
        seed = 0
        torch.manual_seed(seed)
        nn.init.orthogonal_(embedding1.weight.T, gain=1)
        nn.init.orthogonal_(embedding2.weight.T, gain=1)
        
        if 'lora.down_aux.weight' in name:
            print(embedding1.weight.T)
            print(state_dict[name])
        if 'lora.up_aux.weight' in name:
            print(embedding2.weight.T)
            print(state_dict[name])

        if 'lora.down.weight' in name:
            state_dict[name] = state_dict[name] @ embedding1.weight.T
            print(name, state_dict[name].shape)
        if 'lora.up.weight' in name:
            state_dict[name] = embedding2.weight.T @ state_dict[name]
            print(name, state_dict[name].shape)

state_dict = {module_name: param for module_name, param in state_dict.items() if 'lora.up.weight' in module_name or 'lora.down.weight' in module_name}

pipe.unet.load_attn_procs(state_dict)
```


# 2. down_loraだけzero行列で初期化
down_loraだけzero行列で初期化する。

```python
class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, network_alpha=None):
        super().__init__()
...

        nn.init.orthogonal_(self.down_aux.weight, gain=1)
        nn.init.zeros_(self.down.weight)
        nn.init.orthogonal_(self.up.weight, gain=1)
        nn.init.orthogonal_(self.up_aux.weight, gain=1)

```
# 3. 両方zero行列で初期化
up_lora, down_lora共にzero行列で初期化する。

```python
class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, network_alpha=None):
        super().__init__()
...

        nn.init.orthogonal_(self.down_aux.weight, gain=1)
        nn.init.zeros_(self.down.weight)
        nn.init.zeros_(self.up.weight)
        nn.init.orthogonal_(self.up_aux.weight, gain=1)

```


# 4. 学習行列を結合して学習、特異値分解で分割

down_loraとup_loraを一個の行列$(100,50)$として考えaux行列は固定する。これを画像生成時に特異値分解して$U\Gamma V^T=(100,4), (4,4), (4,50)$から$(100,4), (4,50)$にわけ、aux行列を掛けてdown_loraとup_loraに戻す。


```python
class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, network_alpha=None):
        super().__init__()
...
        self.down_aux = nn.Linear(in_features, 100, bias=False)
        self.down_up = nn.Linear(100, 50, bias=False)
        self.up_aux = nn.Linear(50, out_features, bias=False)
...
        nn.init.orthogonal_(self.down_aux.weight, gain=1)
        nn.init.zeros_(self.down_up.weight)
        nn.init.orthogonal_(self.up_aux.weight, gain=1)

        self.down_aux.weight.requires_grad=False
        self.up_aux.weight.requires_grad=False
...
    def forward(self, hidden_states):
        orig_dtype = hidden_states.dtype
        dtype = self.down_up.weight.dtype

        x = self.down_aux(hidden_states.to(dtype))
        x = self.down_up(x)
        up_hidden_states = self.up_aux(x)
...
```


```python
state_dict = torch.load(lora_path, map_location="cpu")
for name in list(state_dict.keys()):
...
        if 'lora.down_up.weight' in name:
            M = state_dict[name]
            U, S, Vh = torch.linalg.svd(M.to(torch.float32))
            rank = 4

            U = U[:, :rank]
            S = S[:rank]
            Vh = Vh[:rank, :]

            U = U @ torch.diag(S)
            Vh = Vh

            state_dict[name.replace('lora.down_up.weight', 'lora.down.weight')] = Vh @ embedding1.weight.T
            state_dict[name.replace('lora.down_up.weight', 'lora.up.weight')] = embedding2.weight.T @ U

```

# 5. normalで初期化
normalをorthogonalの代わりに使った時、lossは下がらず学習は上手く行かなかった。

```python
class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, network_alpha=None):
        super().__init__()
...

        nn.init.normal_(self.down_aux.weight, std=1)
        nn.init.normal_(self.down.weight, std=1 / rank)
        nn.init.zeros_(self.up.weight)
        nn.init.normal_(self.up_aux.weight, std=1)
```


# 結果
TI学習の結果から1～4のLight LoRAの学習した結果を示す。
学習率は2.0e-03、学習stepは500。
学習率が大きい場合、fp16の学習に失敗する場合があって以降--mixed_precision="no"とした。

1.up_loraのみ、2.down_loraのみゼロ行列で初期化した結果はLight LoRAでも学習出来ているように見える。一方、3.両方ゼロ行列で初期化した結果は学習出来ていなかった。
4.学習後に特異値分解で次元を落としたデータでも上手く学習出来ているようには見える。

## 0. TI学習のみ(学習初期)
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/76d17439-3118-6cba-82ec-492d7d6af4ad.png)
## 0. TI+従来LoRA
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/01c6fbe0-925a-aac1-3386-8db0de70948d.png)
## 1. TI+LiLoRA(up_loraだけzero行列で初期化)
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/dfc2b4e5-0879-9274-2127-8a711373b02c.png)
## 2. TI+LiLoRA(down_loraだけzero行列で初期化)
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/25dfbcd8-e71d-2f08-b9ef-23cd49f89163.png)
## 3. TI+LiLoRA(両方zero行列で初期化)
0.のTI学習のみの結果と変わらない。
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/bc1656f8-82ae-0591-edf2-b741778eb1af.png)
## 4. TI+LiLoRA(特異値分解で分割)
![cat_cat_1_TI_1_sample_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/f78e48bb-f4d6-bff0-56ec-d9a06a9c1312.png)

再現度は従来LoRAが最も高そうには見える。個人的には特異値分解で分割もその次に良く見える。

# LiLoRAの大きさ
rank=4の時のLoRA出力ファイルの大きさを示す。
モデル保存時にdown_aux, up_auxを保存するとLoRAサイズは59MBで却って大きくなる。
また4.の$(100,50)$でdown_upを保存した場合は61MBであった。
down_aux, up_auxを削除すると0.4MBまで小さくなる。サイズ感的にはXTIのモデルサイズと張り合えるようになる。

| モデル重み | サイズ |
|:-----------|:------------|
|TI(1token)  |4KB          |
|TI(8token)  |25KB         |
|XTI(8\*16token)|  400KB？|
| 従来LoRA    | 3,211 KB     |
| **LiLoRA(down+up)**     | 379 KB |
| LiLoRA(down+up+down_aux+up_aux)| 58,844 KB |
| LiLoRA(down_up+down_aux+up_aux)| 60,995 KB |



# まとめ
Light LoRA(Lightweight Dreambooth)について確認してみた。
自分でやった限りでは両方ゼロ行列で初期化した場合、何故か上手く行かなかった。また、orthogonalの初期化は意外に重要なようだった。
Light LoRAはHyperDreamBoothの構成の一要素に過ぎないため、その他は自分はよく分かってない。そのほかHypernetworkはStableDiffusionなら別の構造を指すので名称的に良くないと思った。