外れ値を除外する量子化

SDXLの量子化において外れ値(outlier)の影響は結構大きい。
予め外れ値を除外する処理を入れてから量子化する事を考える。

# ①：特定列の外れ値の除外
例えば外れ値が特定の行や列において存在しているとする。
2回特異値分解をして$W=W_0+U_1\*diag(S_1)\*V_1^T+U_2\*diag(S_2)\*V_2^T$として分割出来る事を期待する。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/5ad2e07f-f0e6-798a-3e4a-59bcc403071c.png)

```python
import numpy as np
import scipy

W0 = np.random.randn(512, 512)

W1 = W0.copy()
W1[:, 100] += np.random.randn(512) * 20
W1[10, :] += np.random.randn(512) * 10
W1[20, :] += np.random.randn(512) * 10

rank = 2
U1, S1, Vh1 = np.linalg.svd(W1)
U1 = U1[:, :rank]
S1 = S1[:rank]
Vh1 = Vh1[:rank, :]

qW1 = U1 @ np.diag(S1) @ Vh1
W2 = W1 - qW1

rank = 2
U2, S2, Vh2 = np.linalg.svd(W2)
U2 = U2[:, :rank]
S2 = S2[:rank]
Vh2 = Vh2[:rank, :]

qW2 = U2 @ np.diag(S2) @ Vh2
W3 = W2 - qW2

alpha0 = np.mean(W0)
beta0 = np.mean(np.abs(W0-alpha0))
kur0 = scipy.stats.kurtosis(W0.flatten(), fisher=True)
alpha1 = np.mean(W1)
beta1 = np.mean(np.abs(W1-alpha1))
kur1 = scipy.stats.kurtosis(W1.flatten(), fisher=True)
alpha2 = np.mean(W2)
beta2 = np.mean(np.abs(W2-alpha2))
kur2 = scipy.stats.kurtosis(W2.flatten(), fisher=True)
alpha3 = np.mean(W3)
beta3 = np.mean(np.abs(W3-alpha3))
kur3 = scipy.stats.kurtosis(W3.flatten(), fisher=True)

print('beta0=', beta0)
print('beta1=', beta1)
print('beta2=', beta2)
print('beta3=', beta3)
print('kurtosis0', kur0)
print('kurtosis1', kur1)
print('kurtosis2', kur2)
print('kurtosis3', kur3)
---------------------------------------
beta0= 0.7976069585433163
beta1= 0.853026542377069
beta2= 0.8079153347662861
beta3= 0.7878064121187903
kurtosis0 -0.010549274212497739
kurtosis1 215.96009479267272
kurtosis2 36.00906251339321
kurtosis3 0.006057399107231465
```
尖度(kurtosis)が元の分布(0)→特定列に外れ値追加(215)→一回目の特異値分解成分を引く(36)→二回目の特異値分解成分を引く(0)。外れ値が僅かでも$mean(abs(W))$に比べ、尖度(kurtosis)は大きく変わる。二回の特異値分解成分を引くことで良くなっている。

# ②：ランダムに存在する外れ値の除外
BiLLMとかであったようにsalient_posを別途保存する。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/0514027d-3c04-3d02-8837-66c20196500b.png)

```python
    qmin = np.percentile(W.flatten(), 0.005*100)
    qmax = np.percentile(W.flatten(), 0.995*100)
    
    salient_pos1 = np.where(W<qmin, 1.0, 0.0)
    salient_pos2 = np.where(W>qmax, 1.0, 0.0)
    qW = salient_pos1 * qmin + salient_pos2 * qmax
    W2 = W - qW
```

これは1bitの重み×2であるが、頻度の低いデータの場合は圧縮可能で少なくて済む。
例えばエントロピーの観点では50%の確率で$0,1$が現れる分布は$0.5\*log2(0.5)+0.5\*log2(0.5)=1bit$だが、0.5%のpercentileから求まる値をより絶対値の大きな値を外れ値としてこの低頻度の重みは$0.005\*log2(0.005)+0.005\*log2(0.005)+0.990\*log2(0.990)=0.091bit$と3値のデータでも頻度が少なければ0.1bitで保持可能である。



```python
pip install pyzpaq
-----------------------------------------------------------
import numpy as np
import zpaq
import bz2
import gzip
import lzma

text = ''
for i in range(1024*1024):
    random_num = np.random.rand()
    if random_num < 0.005:
        text = text + 'a'
    elif random_num > 0.995:
        text = text + 'b'
    else:
        text = text + 'c'

text = text.encode("utf-8")
text_compress_zpaq = zpaq.compress(text)
text_compress_bz2  = bz2.compress(text)
text_compress_gzip = gzip.compress(text)
text_compress_lzma = lzma.compress(text)

print('text_compress_bit_zpaq=', len(text_compress_zpaq)/len(text)*8)
print('text_compress_bit_bz2=',  len(text_compress_bz2)/len(text)*8)
print('text_compress_bit_gzip=', len(text_compress_gzip)/len(text)*8)
print('text_compress_bit_lzma=', len(text_compress_lzma)/len(text)*8)
-----------------------------------------------------------
text_compress_bit_zpaq= 0.09407806396484375
text_compress_bit_bz2= 0.10107421875
text_compress_bit_gzip= 0.1347198486328125
text_compress_bit_lzma= 0.175323486328125
```

# 特定のレイヤーの量子化
これら二通りの外れ値を除外した分布を量子化する事で結果を改善したい。
SDXLのUnetはfp16で4.78GB、LoRAはrank=4で27.5MB、0.1bit相当のsalient_posを持つには4.78GB\*(0.1/16)bit=30MBで二つ合わせて1%程度の重みである。
要するにrank=4のLoRAはUnet重みの0.1bit相当である。
Conv2Dは4次元の重み行列を2次元に変換して特異値分解をやった。

絶対値平均$\beta=mean(abs(W))$を使い、これを基準に正規分布の分割後のエントロピーが指定した値になるような間隔p1を使い等間隔に分割する。
分割間隔を非等間隔にした方がcos類似度を向上することは可能だが、自分がやった限りではガウス分布では量子化する重みに対して間隔の調整が必要な割に激的には向上しないのでやらない(0.95が0.952になる程度)。
case=1の場合、(10240,1280)の行列がありこの特異値分解を求めるのが処理時間が非常に遅い。
case=7の場合、(1280,2560,3,3)の行列があり(3840,7680)に変形し特異値分解を求めるのも非常に遅い。
また、(1280,2048)の特異値分解も対象の数が多い事も相まってまあまあ遅い。

特定のレイヤーの量子化を行い、[前回の結果](https://qiita.com/nishiha/items/0e72cb1202af8c939f6e)(下段)と比較すると**embやConv2Dの量子化結果は外れ値を除くことで大幅に改善している。** 外れ値を保持することで+0.2bit必要だが、それを差し引いてもそれ以外でも2～3bitの結果で前回より割と良く見える。
![concat15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/e9bd645b-0db4-ffe8-6cf4-bf1253c597fe.png)
尖度(kurtosis)は傾向として外れ値を除外することでゼロに近づき、特異値分解が有効かどうかは重み分布による。特異値分解によって尖度(kurtosis)が大幅に減少する事もあれば、特異値分解ではあまり変化しない場合もある。


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

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

prompt = "a photo of an astronaut riding a horse on mars"
seed = 42
generator = torch.Generator(device="cuda")
generator = generator.manual_seed(seed)

entropy = 2.0
case = 3

def quant(W, entropy=3.0):
    split_num = 64
    entropy_list = [2.0,   2.5,   3.0,   3.5,   4.0,   4.5,   5.0]
    p1_list      = [1.358, 0.937, 0.655, 0.461, 0.325, 0.230, 0.162]
    entropy_to_p1 =  scipy.interpolate.interp1d(entropy_list, p1_list, kind='linear')
    p1 = entropy_to_p1(entropy)

    orig_shape = W.shape
    if len(orig_shape)==4:
        print('orig_shape=', orig_shape)
        W = W.transpose(0, 2, 1, 3).reshape(orig_shape[0]*orig_shape[2], orig_shape[1]*orig_shape[3])
    
    rank = 2
    U1, S1, Vh1 = np.linalg.svd(W)
    U1 = U1[:, :rank]
    S1 = S1[:rank]
    Vh1 = Vh1[:rank, :]
    print(U1.shape, S1.shape, Vh1.shape)
    qW1 = U1 @ np.diag(S1) @ Vh1
    
    W2 = W - qW1
    U2, S2, Vh2 = np.linalg.svd(W2)
    U2 = U2[:, :rank]
    S2 = S2[:rank]
    Vh2 = Vh2[:rank, :]
    print(U2.shape, S2.shape, Vh2.shape)
    qW2 = U2 @ np.diag(S2) @ Vh2
    
    W3 = W2 - qW2
    
    qmin = np.percentile(W3.flatten(), 0.005*100)
    qmax = np.percentile(W3.flatten(), 0.995*100)
    
    salient_pos1 = np.where(W3<qmin, 1.0, 0.0)
    salient_pos2 = np.where(W3>qmax, 1.0, 0.0)
    qW3 = salient_pos1 * qmin + salient_pos2 * qmax
    W4 = W3 - qW3
    
    alpha = np.mean(W4)
    beta = np.mean(np.abs(W4-alpha))
    gamma = beta * p1
    kur = scipy.stats.kurtosis(W.flatten(), fisher=True)
    kur2 = scipy.stats.kurtosis(W2.flatten(), fisher=True)
    kur3 = scipy.stats.kurtosis(W3.flatten(), fisher=True)
    kur4 = scipy.stats.kurtosis(W4.flatten(), fisher=True)
    print('kurtosis=', kur)
    print('kurtosis2=', kur2)
    print('kurtosis3=', kur3)
    print('kurtosis4=', kur4)
    
    W5 = np.clip(np.round((split_num-1.0)/2.0 + (W4-alpha)/gamma), 0.0, (split_num-1.0))
    entropy2 = 0.0
    for i in range(split_num):
        rate = np.count_nonzero(W5==i)/len(W5.flatten())
        if rate!=0.0:
            entropy2 += -rate*np.log2(rate)
    print('actual_entropy=', entropy2)
    
    qW4 = W5 - (split_num-1.0)/2.0
    qW4 = qW4 * gamma + alpha
    
    qW = qW1 + qW2 + qW3 + qW4
    if len(orig_shape)==4:
        qW = qW.reshape(orig_shape[0], orig_shape[2], orig_shape[1], orig_shape[3]).transpose(0, 2, 1, 3)
    return qW, entropy2


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

total_param = 0
quant_param = 0
for name, param in pipe.unet.named_parameters():
    total_param += param.data.numel()
    if case==0:
        condition = (len(param.size())==2 and 'emb' in name) # 0.010
    if case==1:
        condition = (len(param.size())==2 and not('emb' in name) and not('proj_' in name) and ('proj' in name)) # 0.319
    if case==2:
        condition = (len(param.size())==2 and ('to_k' in name)) # 0.106
    if case==3:
        condition = (len(param.size())==2 and ('to_v' in name)) # 0.106
    if case==4:
        condition = (len(param.size())==2 and ('to_q' in name)) # 0.080
    if case==5:
        condition = (len(param.size())==2 and ('to_out' in name)) # 0.080
    if case==6:
        condition = (len(param.size())==2 and ('ff' in name) and not('proj' in name)) # 0.160
    if case==7:
        condition = len(param.size())==4 # Conv2D 0.130
    if case==8:
        condition = (len(param.size())==2 and not('emb' in name) and ('proj_' in name)) # 0.009
    if case==9:
        condition = (len(param.size())==2 and ('to_' in name)) # 0.372

    if condition:
        print(name, param.size(), param.data.dtype)
        quant_param += param.data.numel()
        orig_data_type = param.data.dtype

        M = param.data.clone().to(torch.float32).numpy()
        M2, actual_entropy = quant(M, entropy=entropy)
        M2 = torch.from_numpy(M2)
        param.data = M2.to(orig_data_type)
        del M, M2

torch.cuda.empty_cache()
print('entropy=', entropy, ', case=', case)
print('quant_param/total_param=', quant_param/total_param)
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_case%d_entropy(%3.1f).png" % (case, entropy))
```


# 特異値分解のrank
特異値分解の外れ値を除外を二回、np.whereの外れ値を除外を二回やる。

外れ値の除外が十分かどうかは重みの分布によって異なる。
例えば以下の様に分布の最大値と最小値の差を$\beta$で割った値が一定値以下になるのを探せばいい。
今回のケースではほとんどの分布では外れ値の除外を行って$\epsilon$は6～8程度だが、$conv,emb$では未だ$\epsilon$は16～50程度あり、十分ではないかもしれない。

```math
W_{n}=W_{n-1}-W_{n,outlier}
```
```math
\epsilon_{n}=\frac{max(W_{n})-min(W_{n})}{mean(abs(W_{n}))}<C
```
特異値分解の外れ値を二回、1%のnp.whereのランダム外れ値を二回除いた場合の$\epsilon_{average}$及び平均尖度(kurtosis)は以下の通りである。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/253959dd-7a12-ee88-44af-da30cd9b7cc4.png)


rank=2の特異値分解の2回からrank=4の特異値分解2回にした場合は以下。
特異値分解の値を自身を8bit量子化すればrankを倍にしても良いが、rankを増やしてもあまり差はない。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/a7122822-40b0-c25a-06fc-5981ff930ffd.png)

特異値分解の処理をskipした場合は以下の様である。
$emb,to\\_k,to\\_out$は特異値分解による除外が有効であり、$to\\_q,conv$は特異値分解による除外は有効ではない。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/29fb4717-93a1-22f3-cb52-9fd4fd072510.png)

# 全レイヤーの量子化

Conv2Dを3bit、embとproj_in,proj_outを4bit、それ以外を共通エントロピーbit数として量子化した時、以下の様になった。参考程度。

outlier(0.2bit)は特異値分解の外れ値0.1bit、1%の外れ値0.1bitを保持する。
outlier(0.3bit)は特異値分解の外れ値0.1bit、1%の外れ値0.1bit×2を保持する。

3.0～3.5bitの量子化を見たが中々に量子化は難しい。
最後の2個の図は過去の記事からBlockwise-NF4などの最も良かった量子化結果。
![concat15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/bbc40e6f-9f66-34f6-4a07-e2fcad30d6ca.png)


# 外れ値除外の繰り返し
特異値分解の除外を2回(0.1bit分)やった後、前述の$\epsilon$の値を確認しながらこの値が8以下になるまで、1回あたり0.1bit使用する外れ値除外を繰り返す。$\beta=0.798\sigma$とすればおおよそ$-3.2σ～3.2σ$となるまで外れ値を除外を行う事とする。

次に、この外れ値除外の繰り返しを行う$\epsilon$の域値で変えてみた。
$\epsilon=8～6.5$は結果と使用bitにそれほど変化はない。
$\epsilon=6.0$だと使用bitが0.1bitも増えてしまうが見合った改善効果はない。
![concat18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/d0966092-2ae9-3f99-df21-13fd72e17901.png)

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

Conv2D_entropy = 3.0
emb_entropy = 4.0
base_entropy = 2.5

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

prompt = "a photo of an astronaut riding a horse on mars"
seed = 42
generator = torch.Generator(device="cuda")
generator = generator.manual_seed(seed)

def Int8_quant(W):
    c = 127.0 / np.max(np.abs(W))
    qW = np.round(W * c)
    qW = qW / c
    return qW

def quant(W, entropy=3.0, max_epsilon=8.0):
    split_num = 128
    entropy_list = [2.0,   2.5,   3.0,   3.5,   4.0,   4.5,   5.0]
    p1_list      = [1.358, 0.937, 0.655, 0.461, 0.325, 0.230, 0.162]
    entropy_to_p1 =  scipy.interpolate.interp1d(entropy_list, p1_list, kind='linear')
    p1 = entropy_to_p1(entropy)

    orig_shape = W.shape
    if len(orig_shape)==4:
        print('orig_shape=', orig_shape)
        W = W.transpose(0, 2, 1, 3).reshape(orig_shape[0]*orig_shape[2], orig_shape[1]*orig_shape[3])
    
    rank = 2
    if rank==0:
        qW1 = np.zeros_like(W)
    else:
        U1, S1, Vh1 = np.linalg.svd(W)
        U1 = U1[:, :rank]
        S1 = S1[:rank]
        Vh1 = Vh1[:rank, :]
        print(U1.shape, S1.shape, Vh1.shape)
        qW1 = U1 @ np.diag(S1) @ Vh1
    W2 = W - qW1

    if rank==0:
        qW2 = np.zeros_like(W)
    else:
        U2, S2, Vh2 = np.linalg.svd(W2)
        U2 = U2[:, :rank]
        S2 = S2[:rank]
        Vh2 = Vh2[:rank, :]
        print(U2.shape, S2.shape, Vh2.shape)
        qW2 = U2 @ np.diag(S2) @ Vh2
    W3 = W2 - qW2
    
    W4 = W3.copy()
    qW3 = np.zeros_like(W4)
    iter = 0
    for i in range(20):
        iter += 1
        qmin = np.percentile(W4.flatten(), 0.005*100)
        qmax = np.percentile(W4.flatten(), 0.995*100)
    
        salient_pos1 = np.where(W4<qmin, 1.0, 0.0)
        salient_pos2 = np.where(W4>qmax, 1.0, 0.0)
        qW3 += salient_pos1 * qmin + salient_pos2 * qmax
        W4 = W3 - qW3

        alpha = np.mean(W4)
        beta = np.mean(np.abs(W4-alpha))
        gamma = beta * p1
        delta = np.max(W4) - np.min(W4)
        epsilon = delta/beta
        print('epsilon%d=%f' % (iter, epsilon))
        if epsilon < max_epsilon:
            break
    
    if epsilon > 20:
        print('actual_entropy=', 8.0)
        if len(orig_shape)==4:
            W = W.reshape(orig_shape[0], orig_shape[2], orig_shape[1], orig_shape[3]).transpose(0, 2, 1, 3)
        qW = Int8_quant(W)
        return qW, 8.0

    W5 = np.clip(np.round((split_num-1.0)/2.0 + (W4-alpha)/gamma), 0.0, (split_num-1.0))
    entropy2 = 0.0
    for i in range(split_num):
        rate = np.count_nonzero(W5==i)/len(W5.flatten())
        if rate!=0.0:
            entropy2 += -rate*np.log2(rate)
    entropy2 += (0.1+iter*0.1)
    print('actual_entropy=', entropy2)
    
    qW4 = W5 - (split_num-1.0)/2.0
    qW4 = qW4 * gamma + alpha
    
    qW = qW1 + qW2 + qW3 + qW4
    if len(orig_shape)==4:
        qW = qW.reshape(orig_shape[0], orig_shape[2], orig_shape[1], orig_shape[3]).transpose(0, 2, 1, 3)
    return qW, entropy2


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

total_param = 0
quant_param = 0
for name, param in pipe.unet.named_parameters():
    total_param += param.data.numel()
    print(name, param.size(), param.data.dtype)
    M = param.data.clone().to(torch.float32).numpy()
    
    orig_data_type = param.data.dtype
    if len(param.size())==2 and 'emb' in name: # 0.010
        M2, actual_entropy = quant(M, entropy=emb_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and not('emb' in name) and not('proj_' in name) and ('proj' in name): # 0.319
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and ('to_k' in name): # 0.106
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and ('to_v' in name): # 0.106
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and ('to_q' in name): # 0.080
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and ('to_out' in name): # 0.080
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and ('ff' in name) and not('proj' in name): # 0.160
        M2, actual_entropy = quant(M, entropy=base_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==4: # Conv2D 0.130
        M2, actual_entropy = quant(M, entropy=Conv2D_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    elif len(param.size())==2 and not('emb' in name) and ('proj_' in name): # 0.009
        M2, actual_entropy = quant(M, entropy=emb_entropy)
        quant_param += param.data.numel() * actual_entropy/16.0
    else:
        quant_param += param.data.numel() * 8.0/16.0
        M2 = Int8_quant(M)

    M2 = torch.from_numpy(M2)
    param.data = M2.to(orig_data_type)
    del M, M2

total_entropy = quant_param/total_param*16.0

torch.cuda.empty_cache()
print('Conv2D_entropy=', Conv2D_entropy)
print('emb_entropy=', emb_entropy)
print('base_entropy=', base_entropy)
print('total_entropy=', total_entropy)

generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_Conv2D(%3.1f)_emb(%3.1f)_base(%3.1f)_total(%3.1f).png" % (Conv2D_entropy, emb_entropy, base_entropy, total_entropy))
```

# projレイヤー3.0bit固定
前回projレイヤーの方が使用bitが多めだったのを鑑みて、projの使用bitを3.0bitに固定して他の共通bit(base,max_epsilon)を変えると以下の様になる。左上から2番目のtotal(3.02bit)がバランス良く見える。
![concat19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/82228ebc-2da0-e89f-485d-6c18847a1fbb.png)


# まとめ
BiLLMのように外れ値を除外する量子化を考えた。
SDXLの全ての重みを2回特異値分解させるのに1時間20分くらいかかる。その割には一部の重み以外では特異値分解をやっても尖度では**あまり改善効果がない。**
外れ値を除けばembやconv2Dの量子化bitを4bit以下に出来る事は可能だが、SDXL内でそれぞれ0.010と0.130とそもそもの割合が低いため効果は薄い。conv2Dの量子化bitを5→3で2bit(totalだと0.26bit)減らしても、外れ値の除外に0.2bit～0.3bit費やすだけならあまり意味はない。
レイヤーごとに異なる量子化bitで量子化すればもっと良くなるだろうが、特異値分解に時間が掛かりすぎるため条件見極めは難しい。