SDXLの量子化5

# エントロピー視点の量子化
等間隔に分割する量子化について考える。
具体的には$percentile(W, 50)-percentile(W, p1)$を用いて分割間隔を決定する。
この分割方法は下記記事で考案した。

https://qiita.com/nishiha/items/bf06d5dcd7fdf8710934

例えば4分割の場合、エントロピーを最大化するのはp1=25%の時で、この場合は各量子化の割合は0.25ずつでエントロピーは最大2bitになる。

```math
0.25*log2(0.25)+0.25*log2(0.25)+0.25*log2(0.25)+0.25*log2(0.25)=2bit
```

しかしながら、下記6分割の方が同じエントロピーでcos類似度を良く再現できる。(この時のp1=14.05%)
```math
(0.0155*log2(0.0155)+0.1250*log2(0.1250)+0.3595*log2(0.3595))*2=1.998bit
```

$y=W\*x$と$y2=qW\*x$の2個のyのcos類似度の観点から見るとエントロピーを最大化する値とcos類似度を最大化するp1の値にずれがある。
例えば仮にエントロピーを3bit以内に抑える場合、**8分割(緑線)では0.98が限界だが、12分割(赤線)では3bitで0.99に到達できる。**
非等間隔に分割するNF4(NormalFloat)も24分割にして端の点統合して24点=>16点に減らしたと考えるなら今回のように分割数を増やした24等間隔分割の結合と考えられる。
```math
0.01*log2(0.01)+0.01*log2(0.01)=0.133bit
```
```math
0.02*log2(0.02)=0.113bit
```
上記のように頻度の少ない点を統合すればエントロピーをほぼ減らさず分割数を減らせれる。

少なくとも図で見る限り今回の24分割の手法(茶線)は4bit以内のエントロピー相当の分割を行う限りは効率よく見え、その時の引数p1は量子化のエントロピーから決定できる。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/288f9d6d-91cb-e124-5c5b-f1cdcbe98a54.png)

上記グラフは適当な正規分布$W$に対しp1を7%～47%に振った時のN分割における実際の量子化エントロピーと量子化前の$y=W\*x$と量子化後の$y2=qW\*x$のcos類似度$cos similarity(y,y2)$をプロットさせたものである。



```python
def quant4(W, kind='linear', p1=25.0):
    orig_shape = W.shape

    q3 = np.percentile(W.flatten(), p1)
    q4 = np.percentile(W.flatten(), 50.00)
    q5 = np.percentile(W.flatten(), 100-p1)

    qmin = q4 + (q3-q4)*2.0
    qmax = q4 + (q5-q4)*2.0

    W2 = np.clip(W, qmin, qmax).flatten()
    x = [qmin, q3, q4, q5, qmax]
    y = [-0.5, 0.5, 1.5, 2.5, 3.5]
    
    f = scipy.interpolate.interp1d(x, y, kind=kind)
    f_inv = scipy.interpolate.interp1d(y, x, kind=kind)

    W2 = np.clip(np.round(f(W2)), 0, 3)
    qW = f_inv(W2).reshape(orig_shape)
    entropy = 0.0
    for i in range(4):
        rate = np.count_nonzero(W2==i)/len(W2)
        if rate!=0.0:
            entropy += -rate*np.log2(rate)
    return qW, entropy
    
def quant6(W, kind='linear', p1=25.0):
    orig_shape = W.shape

    q3 = np.percentile(W.flatten(), p1)
    q4 = np.percentile(W.flatten(), 50.00)
    q5 = np.percentile(W.flatten(), 100-p1)

    q2 = q4 + (q3-q4)*2.0
    q6 = q4 + (q5-q4)*2.0
    qmin = q4 + (q3-q4)*3.0
    qmax = q4 + (q5-q4)*3.0

    W2 = np.clip(W, qmin, qmax).flatten()
    x = [qmin, q2, q3, q4, q5, q6, qmax]
    y = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5]
    
    f = scipy.interpolate.interp1d(x, y, kind=kind)
    f_inv = scipy.interpolate.interp1d(y, x, kind=kind)

    W2 = np.clip(np.round(f(W2)), 0, 5)
    qW = f_inv(W2).reshape(orig_shape)
    entropy = 0.0
    for i in range(6):
        rate = np.count_nonzero(W2==i)/len(W2)
        if rate!=0.0:
            entropy += -rate*np.log2(rate)
    return qW, entropy

......

def quant24(W, kind='linear', p1=25.0):
    orig_shape = W.shape

    q3 = np.percentile(W.flatten(), p1)
    q4 = np.percentile(W.flatten(), 50.00)
    q5 = np.percentile(W.flatten(), 100-p1)

    q1a = q4 + (q3-q4)*11.0
    q1b = q4 + (q3-q4)*10.0
    q1c = q4 + (q3-q4)*9.0
    q1d = q4 + (q3-q4)*8.0
    q1e = q4 + (q3-q4)*7.0
    q1f = q4 + (q3-q4)*6.0
    q1g = q4 + (q3-q4)*5.0
    q1h = q4 + (q3-q4)*4.0
    q1 = q4 + (q3-q4)*3.0
    q2 = q4 + (q3-q4)*2.0
    q6 = q4 + (q5-q4)*2.0
    q7 = q4 + (q5-q4)*3.0
    q7a = q4 + (q5-q4)*4.0
    q7b = q4 + (q5-q4)*5.0
    q7c = q4 + (q5-q4)*6.0
    q7d = q4 + (q5-q4)*7.0
    q7e = q4 + (q5-q4)*8.0
    q7f = q4 + (q5-q4)*9.0
    q7g = q4 + (q5-q4)*10.0
    q7h = q4 + (q5-q4)*11.0
    qmin = q4 + (q3-q4)*12.0
    qmax = q4 + (q5-q4)*12.0

    W2 = np.clip(W, qmin, qmax).flatten()
    x = [qmin, q1a,q1b,q1c,q1d,q1e,q1f,q1g,q1h, q1, q2, q3, q4, q5, q6, q7,q7a,q7b,q7c,q7d,q7e,q7f,q7g,q7h, qmax]
    y = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5,9.5,10.5,11.5,12.5,13.5,14.5,15.5,16.5,17.5,18.5,19.5,20.5,21.5,22.5,23.5]
    
    f = scipy.interpolate.interp1d(x, y, kind=kind)
    f_inv = scipy.interpolate.interp1d(y, x, kind=kind)

    W2 = np.clip(np.round(f(W2)), 0, 23)
    qW = f_inv(W2).reshape(orig_shape)
    entropy = 0.0
    for i in range(24):
        rate = np.count_nonzero(W2==i)/len(W2)
        if rate!=0.0:
            entropy += -rate*np.log2(rate)
    return qW, entropy
```

# SDXLでの量子化結果
``emb,proj,proj_in proj_out,to_k,to_v,to_q,to_out,ff,Conv2d``の各レイヤーに対して量子化を行う。[前回](https://qiita.com/nishiha/items/6a21dc3f19b58a74474b)と違うのはprojからproj_inとproj_outを分けている。
この各レイヤーに対する量子化結果は以下の様になった。
基本的にはエントロピーを上げるようにp1を50に近づけると結果の量子化劣化は改善する(例：2列目proj)。ところが``emb,Conv2D``の重みは逆にp1を50に近づけると却って劣化が激しくなる。
この原因を探るため量子化重みでqmax,qminより大きな値(salient)を元に戻すことを行った。また、分布の尖度(kurtosis)を確認した。
この結果``emb,Conv2D,proj_in proj_out``は尖度が非常に高く、重みによっては尖度が10を超える事もままあった(一番酷い尖度は79)。これは重みの分布が正規分布(ガウス分布)と乖離が大きく、この場合``percentile``による分割間隔の決定が良くないことを示している。現にqmax,qminより大きな値(salient)を復元すると``emb,Conv2D,proj_in proj_out``の結果が改善する事からほんの僅かな割合の外れ値によって劣化している。(または分布内に無駄に中央値付近の値が多いと解釈するべきか)

![concat15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/07f31623-11b8-cc94-2fec-c2a46370c871.png)

外れ値の精度を保証するのはBiLLMとか多分以下の論文とかでも見られる。

https://arxiv.org/pdf/2208.07339

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

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

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

entropy = 3.0
case = 0

def quant24(W, entropy=3.0, salient=False):
    entropy_list = [1.8, 2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0]
    p1_list = [10.28, 13.98, 17.55, 21.07, 24.32, 27.28, 30.07, 32.47, 34.62, 36.52, 38.25, 39.84]
    entropy_to_p1 =  scipy.interpolate.interp1d(entropy_list, p1_list, kind='linear')
    p1 = entropy_to_p1(entropy)

    orig_shape = W.shape

    q3 = np.percentile(W.flatten(), p1)
    q4 = np.percentile(W.flatten(), 50.00)
    q5 = np.percentile(W.flatten(), 100-p1)

    q1a = q4 + (q3-q4)*11.0
    q1b = q4 + (q3-q4)*10.0
    q1c = q4 + (q3-q4)*9.0
    q1d = q4 + (q3-q4)*8.0
    q1e = q4 + (q3-q4)*7.0
    q1f = q4 + (q3-q4)*6.0
    q1g = q4 + (q3-q4)*5.0
    q1h = q4 + (q3-q4)*4.0
    q1 = q4 + (q3-q4)*3.0
    q2 = q4 + (q3-q4)*2.0
    q6 = q4 + (q5-q4)*2.0
    q7 = q4 + (q5-q4)*3.0
    q7a = q4 + (q5-q4)*4.0
    q7b = q4 + (q5-q4)*5.0
    q7c = q4 + (q5-q4)*6.0
    q7d = q4 + (q5-q4)*7.0
    q7e = q4 + (q5-q4)*8.0
    q7f = q4 + (q5-q4)*9.0
    q7g = q4 + (q5-q4)*10.0
    q7h = q4 + (q5-q4)*11.0
    qmin = q4 + (q3-q4)*12.0
    qmax = q4 + (q5-q4)*12.0

    W2 = np.clip(W, qmin, qmax).flatten()
    x = [qmin, q1a,q1b,q1c,q1d,q1e,q1f,q1g,q1h, q1, q2, q3, q4, q5, q6, q7,q7a,q7b,q7c,q7d,q7e,q7f,q7g,q7h, qmax]
    y = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5,9.5,10.5,11.5,12.5,13.5,14.5,15.5,16.5,17.5,18.5,19.5,20.5,21.5,22.5,23.5]
    
    f = scipy.interpolate.interp1d(x, y, kind='linear')
    f_inv = scipy.interpolate.interp1d(y, x, kind='linear')

    W2 = np.clip(np.round(f(W2)), 0, 23)
    qW = f_inv(W2).reshape(orig_shape)
    actual_entropy = 0.0
    for i in range(24):
        rate = np.count_nonzero(W2==i)/len(W2)
        if rate!=0.0:
            actual_entropy += -rate*np.log2(rate)

    if salient:
        salient_pos = np.where(W>qmax, 1.0, 0.0)
        salient_pos += np.where(W<qmin, 1.0, 0.0)
        salient_rate = np.sum(salient_pos) /len(W2)
        print('salient_rate=', salient_rate)
        print('qmin,qmax=', qmin, qmax)
        qW = np.where(W>qmax, W, qW)
        qW = np.where(W<qmin, W, qW)
        kur = scipy.stats.kurtosis(W.flatten(), fisher=True)
        print('kurtosis=', kur)

    return qW, actual_entropy

def quant48(W, entropy=3.0, salient=False):
    entropy_list = [1.8, 2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.2, 4.4, 4.6, 4.8, 5.0]
    p1_list = [10.28, 13.98, 17.55, 21.07, 24.32, 27.28, 30.07, 32.47, 34.62, 36.52, 38.25, 39.84, 41.16, 42.28, 43.26, 44.14, 44.90]
    entropy_to_p1 =  scipy.interpolate.interp1d(entropy_list, p1_list, kind='linear')
    p1 = entropy_to_p1(entropy)

    orig_shape = W.shape

    q3 = np.percentile(W.flatten(), p1)
    q4 = np.percentile(W.flatten(), 50.00)
    q5 = np.percentile(W.flatten(), 100-p1)

    q1a = q4 + (q3-q4)*23.0
    q1b = q4 + (q3-q4)*22.0
    q1c = q4 + (q3-q4)*21.0
    q1d = q4 + (q3-q4)*20.0
    q1e = q4 + (q3-q4)*19.0
    q1f = q4 + (q3-q4)*18.0
    q1g = q4 + (q3-q4)*17.0
    q1h = q4 + (q3-q4)*16.0
    q1i = q4 + (q3-q4)*15.0
    q1j = q4 + (q3-q4)*14.0
    q1k = q4 + (q3-q4)*13.0
    q1l = q4 + (q3-q4)*12.0
    q1m = q4 + (q3-q4)*11.0
    q1n = q4 + (q3-q4)*10.0
    q1o = q4 + (q3-q4)*9.0
    q1p = q4 + (q3-q4)*8.0
    q1q = q4 + (q3-q4)*7.0
    q1r = q4 + (q3-q4)*6.0
    q1s = q4 + (q3-q4)*5.0
    q1t = q4 + (q3-q4)*4.0
    q1 = q4 + (q3-q4)*3.0
    q2 = q4 + (q3-q4)*2.0
    q6 = q4 + (q5-q4)*2.0
    q7 = q4 + (q5-q4)*3.0
    q7a = q4 + (q5-q4)*4.0
    q7b = q4 + (q5-q4)*5.0
    q7c = q4 + (q5-q4)*6.0
    q7d = q4 + (q5-q4)*7.0
    q7e = q4 + (q5-q4)*8.0
    q7f = q4 + (q5-q4)*9.0
    q7g = q4 + (q5-q4)*10.0
    q7h = q4 + (q5-q4)*11.0
    q7i = q4 + (q5-q4)*12.0
    q7j = q4 + (q5-q4)*13.0
    q7k = q4 + (q5-q4)*14.0
    q7l = q4 + (q5-q4)*15.0
    q7m = q4 + (q5-q4)*16.0
    q7n = q4 + (q5-q4)*17.0
    q7o = q4 + (q5-q4)*18.0
    q7p = q4 + (q5-q4)*19.0
    q7q = q4 + (q5-q4)*20.0
    q7r = q4 + (q5-q4)*21.0
    q7s = q4 + (q5-q4)*22.0
    q7t = q4 + (q5-q4)*23.0
    qmin = q4 + (q3-q4)*24.0
    qmax = q4 + (q5-q4)*24.0

    W2 = np.clip(W, qmin, qmax).flatten()
    x = [qmin, q1a,q1b,q1c,q1d,q1e,q1f,q1g,q1h,q1i,q1j,q1k,q1l,q1m,q1n,q1o,q1p,q1q,q1r,q1s,q1t, q1, q2, q3, q4, q5, q6, q7,q7a,q7b,q7c,q7d,q7e,q7f,q7g,q7h,q7i,q7j,q7k,q7l,q7m,q7n,q7o,q7p,q7q,q7r,q7s,q7t, qmax]
    y = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5,10.5,11.5,12.5,13.5,14.5,15.5,16.5,17.5,18.5,19.5,20.5,21.5,22.5,23.5,
              24.5,25.5,26.5,27.5,28.5,29.5,30.5,31.5,32.5,33.5,34.5,35.5,36.5,37.5,38.5,39.5,40.5,41.5,42.5,43.5,44.5,45.5,46.5,47.5]
    
    f = scipy.interpolate.interp1d(x, y, kind='linear')
    f_inv = scipy.interpolate.interp1d(y, x, kind='linear')

    W2 = np.clip(np.round(f(W2)), 0, 47)
    qW = f_inv(W2).reshape(orig_shape)
    actual_entropy = 0.0
    for i in range(48):
        rate = np.count_nonzero(W2==i)/len(W2)
        if rate!=0.0:
            actual_entropy += -rate*np.log2(rate)

    if salient:
        salient_pos = np.where(W>qmax, 1.0, 0.0)
        salient_pos += np.where(W<qmin, 1.0, 0.0)
        salient_rate = np.sum(salient_pos) /len(W2)
        print('salient_rate=', salient_rate)
        print('qmin,qmax=', qmin, qmax)
        qW = np.where(W>qmax, W, qW)
        qW = np.where(W<qmin, W, qW)
        kur = scipy.stats.kurtosis(W.flatten(), fisher=True)
        print('kurtosis=', kur)
        
    return qW, actual_entropy


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()
        if entropy > 4.1:
            M2, actual_entropy = quant48(M, entropy=entropy)
        else:
            M2, actual_entropy = quant24(M, entropy=entropy)
        M2 = torch.from_numpy(M2)
        param.data = M2.to(orig_data_type)
        del M, M2
        print('actual_entropy=', actual_entropy)

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))
```

# まとめ
$percentile(W, p1)$を用いてSDXLにて分割する方法を示した。
しかし、embやConv2DはBlockwise-NF4などの$max(abs(W))$を用いて分割する方が、BitNetなどの$mean(abs(W))$や今回の$percentile(W, p1)$を用いて分割する方法よりも良い。
これは正規分布からの乖離が大きいためである。尖度によってこれら($\beta$や$p1$)に補正を加えるか、qmax,qminより大きな値(salient)を復元すると良くなる。また、任意分布に何らかの関数を掛けて正規分布へ変換できる変換関数を考えてもよい。（これは自明ではないが正規分布の累積分布関数と任意分布の累積分布関数をfitする）


今回示した手法ならば外れ値を除いた最大値基準の等間隔の分割(Blockwise)より0.5bit程度良い。とはいえ、これはおそらく非等間隔分割のBlockwise-NF4と僅かな差しかない。
今回示したのは理論エントロピーで実際の圧縮比率ではないがzpaqを使うとbz2よりも良く、ほぼ理論エントロピーに近い圧縮が可能なのでエントロピー基準の方が表記的には効率的なのではと考えた。