BitNetと自己教師学習2

前回の自分の記事でSDXLのlinearレイヤーをbitlinearレイヤーに置換して、「学習済み重みの推論」と「量子化重みの推論」とでMSELossを取って正解ラベル(SDXLでは画像)を与えずに学習させようとしたが上手く行かなかった。

https://qiita.com/nishiha/items/2c5ad5d6b0d29eefd343

SDXLの**CrossAttentionのK,V重みの量子化のみ**を考えるのなら事前にTextEncoderの出力$x$と学習済み重み$W$が分かっていればこの二つだけから量子化学習を行えるのではと思った。

ただし、これで量子化出来るのはSDXLの10%ほどの重みに限られるので、前回の試行(推論中に学習可能なレイヤーに置き換える)より目標は後退している。

# 簡単なテストコード
分割数$s$の場合の量子化を下記とする。
```math
\alpha = mean(W)
```
```math
\beta = mean(abs(W-\alpha))
```
```math
quantW=clip(round((W-\alpha)/\beta+\frac{s-1}{2}), 0, s-1)-\frac{s-1}{2}
```
入力$x$が常に固定なら以下の様に書ける。
この時、ループを繰り返すと損失が減って行き、100回程度の計算でもcos類似度が上昇する。ただし$x$のサンプル数が少なすぎる場合、特定のpromptに過学習する事になる。
この時、学習済み重み$W_0$から量子化学習重み$W$は僅かに変化しており、推論結果は量子化済み重みでも元の結果を再現出来ている。

```python
import torch

def quant(W, split_num=4):
    split_num = float(split_num)
    alpha = torch.mean(W)
    beta = torch.mean(torch.abs(W-alpha))
    qW = torch.clamp(torch.round((split_num-1.0)/2.0 + (W-alpha)/(beta)), 0.0, (split_num-1.0)) - (split_num-1.0)/2.0
    if split_num==2.0:
        beta *= 2.0
    qW = qW * beta
    qW = (qW - W).detach() + W
    return qW

x = torch.randn((10,768)) * 3.2 + 1.5
x = torch.nn.LayerNorm(768)(x)
x = x.detach()
W = torch.randn((360,768), requires_grad=True)
W0 = W.clone()

optimizer = torch.optim.Adam([W], lr=0.0025)
for i in range(100):
    y_true = (torch.nn.functional.linear(x, W0)).detach()
    y_pred = torch.nn.functional.linear(x, quant(W, 2))
    loss = torch.nn.MSELoss()(y_pred, y_true)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    cos_sim = torch.mean(torch.nn.functional.cosine_similarity(y_pred, y_true)).item()
    if i % 10==0:
        print('cos_sim_%03d=' % (i), cos_sim)
        print('loss=', loss.item())

print('W0=', W0)
print('W=', W)
print('qW=', quant(W, 2))
print('y_true=', torch.nn.functional.linear(x, W0))
print('y_pred=', torch.nn.functional.linear(x, quant(W, 2)))
print('y_pred_init=', torch.nn.functional.linear(x, quant(W0, 2)))
---------------------------------------------
cos_sim_000= 0.8010725975036621
loss= 272.1870422363281
cos_sim_010= 0.8735825419425964
loss= 182.60972595214844
cos_sim_020= 0.9228243827819824
loss= 118.36371612548828
cos_sim_030= 0.9530874490737915
loss= 74.97147369384766
cos_sim_040= 0.9713737368583679
loss= 46.717227935791016
cos_sim_050= 0.9826198816299438
loss= 28.60102653503418
cos_sim_060= 0.9883385896682739
loss= 18.659488677978516
cos_sim_070= 0.9919894933700562
loss= 12.495198249816895
cos_sim_080= 0.9941555857658386
loss= 9.006094932556152
cos_sim_090= 0.9955546259880066
loss= 6.822784900665283
W0= tensor([[ 0.8177, -1.6385,  0.8304,  ..., -1.5038, -1.6542, -0.9612],
        [-0.2075,  0.6370,  0.6033,  ...,  0.9107, -1.2797,  0.6104],
        [ 0.2519,  0.1167, -0.2490,  ..., -0.3666,  0.2556, -1.5370],
        ...,
        [ 1.0984,  0.9009, -0.5185,  ..., -1.3383,  1.0501, -1.5089],
        [ 0.7476, -0.6304, -0.8311,  ...,  0.0246, -0.5052, -0.5490],
        [ 0.7840, -0.2057,  0.4212,  ...,  0.1208,  1.3372, -1.2740]],
       grad_fn=<CloneBackward0>)
W= tensor([[ 0.8530, -1.5036,  0.6316,  ..., -1.4779, -1.8706, -1.0332],
        [-0.3586,  0.4883,  0.7505,  ...,  1.0871, -1.1571,  0.4658],
        [ 0.4234,  0.1454, -0.2180,  ..., -0.2327,  0.3872, -1.4187],
        ...,
        [ 1.2002,  1.0089, -0.6320,  ..., -1.4489,  0.9234, -1.3995],
        [ 0.6657, -0.7892, -0.8473,  ..., -0.0664, -0.7419, -0.3784],
        [ 0.8509, -0.0903,  0.5842,  ...,  0.2528,  1.2808, -1.4366]],
       requires_grad=True)
qW= tensor([[ 0.8028, -0.8028,  0.8028,  ..., -0.8028, -0.8028, -0.8028],
        [-0.8028,  0.8028,  0.8028,  ...,  0.8028, -0.8028,  0.8028],
        [ 0.8028,  0.8028, -0.8028,  ..., -0.8028,  0.8028, -0.8028],
        ...,
        [ 0.8028,  0.8028, -0.8028,  ..., -0.8028,  0.8028, -0.8028],
        [ 0.8028, -0.8028, -0.8028,  ..., -0.8028, -0.8028, -0.8028],
        [ 0.8028, -0.8028,  0.8028,  ...,  0.8028,  0.8028, -0.8028]],
       grad_fn=<AddBackward0>)
y_true= tensor([[-23.4319, -24.8198, -48.9392,  ...,  14.6610,  11.1746, -40.8871],
        [ 91.0460, -21.2551,  16.5874,  ...,  41.1129,   6.3804,  43.5760],
        [ -2.7527, -14.4890, -11.1396,  ...,  13.4700, -78.2040,   5.7792],
        ...,
        [ 24.7213,  34.6690, -37.5822,  ..., -36.4096,  28.8839,   4.5973],
        [ 22.3427,   2.9519,  71.4005,  ..., -10.5914,  -8.2352,  -7.0887],
        [  1.1459, -45.0301, -12.6591,  ...,  27.6968, -54.3258,  23.3390]],
       grad_fn=<MmBackward0>)
y_pred= tensor([[-23.2912, -24.5324, -47.9054,  ...,  13.2326,  14.3372, -41.3761],
        [ 94.8195, -18.7135,  16.8677,  ...,  39.6075,   2.7964,  42.5792],
        [ -8.1964, -13.9528,  -8.5078,  ...,  15.7938, -72.3384,   5.3755],
        ...,
        [ 28.3064,  31.1511, -37.3318,  ..., -36.4068,  32.4308,   7.2599],
        [ 23.7614,   6.6244,  70.8209,  ..., -10.2627,  -8.4588,  -7.9016],
        [ -1.7436, -48.7466, -10.7339,  ...,  26.6857, -51.1799,  22.7436]],
       grad_fn=<MmBackward0>)
y_pred_init= tensor([[ -8.3595,  -2.1513, -22.5430,  ...,  11.5740, -14.3811, -19.2176],
        [ 53.4883,   5.3388,   5.6832,  ...,  24.3135,  20.2834,  41.7076],
        [  1.4805, -20.4796,  14.9275,  ...,  25.6816, -48.3107,  14.7088],
        ...,
        [ 13.1133,   6.6648,  -9.9837,  ..., -30.3539,  19.1926, -11.4375],
        [ 15.3122,   9.9555,  25.2670,  ...,  -3.7394,  -0.7025,   5.6137],
        [-18.2769, -23.8019, -14.6380,  ...,  13.2467, -37.0301,  -6.3800]],
       grad_fn=<MmBackward0>)
```

# SDXLのCross-AttentionのK,V量子化
SDXLのpipelineから$x,W$を取り出し、$y_{true}=W_0\*x$、$y_{pred}=quant(W)\*x$としてこの二個のMSElossを取って$W$を学習させる。この量子化の学習(cuda)をし、事前にTextEncoderの推論を実行するのでSDXLの``pipe.enable_model_cpu_offload()``のせいでcudaとcpuが混在して動かなくなったりする。これのせいで``torch.Generator(device="cpu")``とせざるを得なかった。この為、前回までと出力が若干変わる。

変換中のlogを示すと学習が進むにつれ損失が減り、cos類似度が上昇しているのが確認できる。

```python
down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight torch.Size([640, 2048]) cpu
cos_sim_000= 0.8646824955940247
loss= 1.2962104082107544
cos_sim_025= 0.9022889137268066
loss= 0.6338478922843933
cos_sim_050= 0.9089803695678711
loss= 0.5473829507827759
cos_sim_075= 0.8997079133987427
loss= 0.5937912464141846
cos_sim_last= 0.9140064120292664
loss_last= 0.5820772051811218
down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v.weight torch.Size([640, 2048]) cpu
cos_sim_000= 0.7531372904777527
loss= 0.18519467115402222
cos_sim_025= 0.7568215131759644
loss= 0.24309834837913513
cos_sim_050= 0.8272229433059692
loss= 0.17504878342151642
cos_sim_075= 0.8463781476020813
loss= 0.1371772289276123
cos_sim_last= 0.8142029643058777
loss_last= 0.17443372309207916
down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_k.weight torch.Size([640, 2048]) cpu
cos_sim_000= 0.8964177966117859
loss= 1.5087316036224365
cos_sim_025= 0.9318107962608337
loss= 0.48295220732688904
cos_sim_050= 0.9347710609436035
loss= 0.3807266652584076
cos_sim_075= 0.9358013272285461
loss= 0.40587469935417175
cos_sim_last= 0.9258428812026978
loss_last= 0.44096773862838745
down_blocks.1.attentions.0.transformer_blocks.1.attn2.to_v.weight torch.Size([640, 2048]) cpu
...
```

量子化学習の結果を示すと以下の様になった。劣化の激しいbit1、bit1.58量子化の結果が量子化学習によって結果が改善している。

注目する点は自己教師あり学習なので**学習には正解ラベル(画像)が必要ではない点である。**

![concat13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/af0c42fd-b8f3-649a-2a3f-d8433d5c4d3c.png)

ただし、特定のprompt(入力$x$)に過学習している可能性は否定できず、その場合は学習promptの多様性が必要であろう。アボガドの形の椅子とかオレンジ色のフードのカンガルーとか適当な10個の学習prompt+1個の空のnegative_promptの場合、学習promptに宇宙飛行士や馬が含まれていないにもかかわらずtrain_num=50,split_num=2で馬に乗る宇宙飛行士が描ける。しかし、反面この場合学習回数の調整がシビアである。学習が安定するにはある程度negative_promptの比率を上げる必要があるのかもしれない。
![concat14.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/5cfbb81c-6974-9b1c-fb85-bb720b24c89b.png)



```python
from diffusers import DiffusionPipeline
import numpy as np
import torch

def quant(W, split_num=4):
    split_num = float(split_num)
    alpha = torch.mean(W)
    beta = torch.mean(torch.abs(W-alpha))
    qW = torch.clamp(torch.round((split_num-1.0)/2.0 + (W-alpha)/(beta)), 0.0, (split_num-1.0)) - (split_num-1.0)/2.0
    if split_num==2.0:
        beta *= 2.0
    qW = qW * beta + alpha
    qW = (qW - W).detach() + W
    return qW

def train_ssl(x, W, train_num=100, lr=0.0025, split_num=2):
    orig_data_type = W.dtype
    W = W.to(torch.float32)
    x = x.detach()
    W = W.detach()
    W0 = W.clone()
    W = W.requires_grad_(True)
    y_true = (torch.nn.functional.linear(x, W0)).detach()
    optimizer = torch.optim.Adam([W], lr=lr)
    if train_num!=0:
        for i in range(train_num):
            #y_true = (torch.nn.functional.linear(x, W0)).detach()
            y_pred = torch.nn.functional.linear(x, quant(W, split_num))
            loss = torch.nn.MSELoss()(y_pred, y_true)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            cos_sim = torch.mean(torch.nn.functional.cosine_similarity(y_pred, y_true)).item()
            if i % 25==0:
                print('cos_sim_%03d=' % (i), cos_sim)
                print('loss=', loss.item())

        
        print('cos_sim_last=', cos_sim)
        print('loss_last=', loss.item())

    del W0, x, y_true, optimizer,
    return quant(W, split_num).to(orig_data_type)

model_id = './stable-diffusion-xl-base-1.0/'
output_path = './sdxl_ssl_CA_kv/'

prompt = "a photo of an astronaut riding a horse on mars"
prompt2 = ['a stained glass window of a panda eating bamboo',
           'a corgi wearing a red bowtie and a purple party hat',
           'a hedgehog using a calculator',
           'an armchair in the shape of an avocado',
           'A green sign that says "Very Deep Learning" and is at the edge of the Grand Canyon. Puffy white clouds are in the sky.',
           'A portrait photo of a kangaroo wearing an orange hoodie and blue sunglasses standing on the grass in front of the Sydney Opera House holding a sign on the chest that says Welcome Friends!',
           'New York Skyline with "Hello World" written with fireworks on the sky.',
           'Greek statue of a man tripping over a cat.',
           'Hovering cow abducting aliens.',
           'An origami of a monkey dressed as a monk riding a bike on a mountain.']

torch.cuda.empty_cache()
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16")
pipe.enable_model_cpu_offload()
torch.cuda.empty_cache()

prompt_embeds, negative_prompt_embeds, _, _ = pipe.encode_prompt(prompt=prompt, negative_prompt='')
x = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0).clone().to(torch.float32).to("cuda")
print(x.shape, x.device)

train_num = 100
split_num = 4

for name, param in pipe.unet.named_parameters():
    if len(param.shape)==2 and param.shape[1]==2048:
        print(name, param.shape, param.device)
        W = param.clone().to("cuda")
        param.data = train_ssl(x, W, train_num=train_num, split_num=split_num).to("cpu")
        del W

del x, prompt_embeds, negative_prompt_embeds, _
torch.cuda.empty_cache()
pipe.to("cpu")
torch.cuda.empty_cache()

seed = 42
generator = torch.Generator(device="cpu")
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_ssl%02d_%04d.png" % (split_num, train_num))
```

# まとめ
SDXLで自己教師あり学習でCross-AttentionのK,V重みの量子化学習をした。
Cross-AttentionのK,V重みは入力$x$がTextEncoderの出力で固定なのでそれほど難しくない。

SDXLのCross-Attention以外の重みで学習しようとすると入力$x$自体がprompt、step、seed値によって変化するので多分学習の収束は更に遅い。また、推論中のUnetの中間層の入力を取り出すのは結構難しそうで、それなら前回記事のように推論中に量子化学習可能なレイヤーを考える方が易しい。