3Uniform affine quantization: công thức cốt lõi#

3.1Định nghĩa#

Trong xử lý tín hiệu, lượng tử hóa là ánh xạ từ một tập giá trị lớn (thường liên tục) vào một tập hữu hạn nhỏ hơn. Trong deep learning, dạng phổ biến nhất là lượng tử hóa đều (uniform): tập đích gồm các điểm cách đều nhau. Dạng đều được ưa chuộng vì nó cho phép tính toán trực tiếp bằng phép toán số nguyên (Chương 6). Các dạng không đều như NF4 (Mục 11.5) cho sai số nhỏ hơn với cùng số bit, nhưng thường phải giải mã về số thực trước khi tính.

3.2Suy ra công thức#

Ta muốn một ánh xạ tuyến tính từ đoạn số thực $[\alpha, \beta]$ vào đoạn số nguyên $[q_{min}, q_{max}]$, với $\alpha \mapsto q_{min}$ và $\beta \mapsto q_{max}$. Mọi ánh xạ tuyến tính (chính xác hơn là affine) một biến có dạng $q = x/S + Z$. Thế hai điều kiện đầu mút vào:

$$q_{min} = \frac{\alpha}{S} + Z, \qquad q_{max} = \frac{\beta}{S} + Z.$$

Trừ vế theo vế rồi giải ra $S$, sau đó thế lại để tìm $Z$:

$$\boxed{S = \frac{\beta - \alpha}{q_{max} - q_{min}}}, \qquad \boxed{Z = \left\lfloor q_{min} - \frac{\alpha}{S} \right\rceil}$$

$Z$ được làm tròn để là số nguyên (lý do ở Mục 3.4). Thêm phép làm tròn cho $q$ và phép cắt cho các giá trị nằm ngoài dải, ta được cặp công thức cơ bản:

$$\textbf{Quantize:}\quad q = \text{clamp}\left(\left\lfloor \frac{x}{S} \right\rceil + Z,\; q_{min},\; q_{max}\right)$$

$$\textbf{Dequantize:}\quad \hat{x} = S\,(q - Z)$$

Hàm hợp $x \mapsto \hat{x}$ gọi là fake quantization (hay quantize–dequantize). Nó cho giá trị vẫn là số thực nhưng chỉ nhận các giá trị trên lưới lượng tử. Hàm này đóng vai trò trung tâm trong mô phỏng và huấn luyện (Chương 7, 9).

3.3Ví dụ tính tay#

Lượng tử hóa dải $[\alpha, \beta] = [-1, 3]$ sang uint8, tức $[q_{min}, q_{max}] = [0, 255]$:

1. $S = \dfrac{3 - (-1)}{255 - 0} = \dfrac{4}{255} \approx 0{,}015686$.
2. $Z = \left\lfloor 0 - \dfrac{-1}{0{,}015686} \right\rceil = \lfloor 63{,}75 \rceil = 64$.
3. Lấy $x = 1{,}5$: $q = \lfloor 1{,}5 / 0{,}015686 \rceil + 64 = \lfloor 95{,}625 \rceil + 64 = 160$.
4. Khôi phục: $\hat{x} = 0{,}015686 \times (160 - 64) = 1{,}5059$. Sai số $0{,}0059$, nhỏ hơn $S/2 \approx 0{,}0078$.
5. Lấy $x = 0$: $q = 64 = Z$, và $\hat{x} = 0$ chính xác tuyệt đối.

Để nhìn rõ hình dạng của phép lượng tử hóa, Hình 4 dùng cùng dải nhưng chỉ 3 bit ($[0, 7]$), để các bậc thang đủ lớn.

Hình 4 còn cho thấy một chi tiết hay bị bỏ qua. Vì $Z = \lfloor 1{,}75 \rceil = 2$ bị làm tròn, dải thực sự biểu diễn được là $[S(0 - 2),\, S(7 - 2)] = [-1{,}143;\ 2{,}857]$, lệch nhẹ so với dải $[-1, 3]$ ban đầu. Với 8 bit, độ lệch này nhỏ hơn nhiều nhưng vẫn tồn tại.

3.4Vì sao số 0 phải được biểu diễn chính xác#

Số 0 xuất hiện với tần suất rất cao trong mạng nơ-ron: vùng zero-padding của phép tích chập, đầu ra âm bị ReLU đưa về 0, các phần tử bị mask trong attention. Nếu số 0 bị biểu diễn lệch thành một giá trị $\delta \neq 0$, sai số này không ngẫu nhiên mà có hệ thống và cộng dồn qua hàng nghìn phần tử trong một tổng. Bắt $Z$ là số nguyên bảo đảm $x = 0 \Rightarrow q = Z \Rightarrow \hat{x} = 0$. Cũng vì lý do này, dải $[\alpha, \beta]$ luôn được mở rộng để chứa số 0: $\alpha \le 0 \le \beta$.

3.5Lượng tử đối xứng và bất đối xứng#

Bất đối xứng (asymmetric, affine): dùng đúng công thức ở trên, $Z$ tùy ý.

Đối xứng (symmetric): ép $Z = 0$ và dùng dải đối xứng quanh 0, $[-c, c]$ với $c = \max|x|$. Thường dùng miền số nguyên hạn chế $[-(2^{b-1} - 1),\ 2^{b-1} - 1]$, ví dụ $[-127, 127]$ cho int8, bỏ giá trị $-128$ để lưới đối xứng hoàn hảo:

$$S = \frac{c}{2^{b-1} - 1}, \qquad q = \text{clamp}\left(\left\lfloor \frac{x}{S} \right\rceil,\; -(2^{b-1}-1),\; 2^{b-1}-1\right).$$

Đo trên chính dữ liệu của Hình 5 (200 000 giá trị sau ReLU), MSE của lượng tử đối xứng so với bất đối xứng là $2{,}67 \times 10^{-2}$ so với $5{,}79 \times 10^{-3}$ ở 4 bit (gấp 4,6 lần), và $8{,}11 \times 10^{-5}$ so với $2{,}01 \times 10^{-5}$ ở 8 bit. Quy ước phổ biến trong TFLite và PyTorch vì thế là trọng số đối xứng, activation bất đối xứng.

3.6Cài đặt#

import numpy as np

def affine_params(alpha, beta, qmin, qmax):
    alpha, beta = min(alpha, 0.0), max(beta, 0.0)       # dải luôn chứa 0 (Mục 3.4)
    S = (beta - alpha) / (qmax - qmin)
    Z = int(np.clip(np.round(qmin - alpha / S), qmin, qmax))
    return S, Z

def quantize(x, S, Z, qmin, qmax):
    return np.clip(np.round(x / S) + Z, qmin, qmax).astype(np.int32)

def dequantize(q, S, Z):
    return S * (q.astype(np.float32) - Z)

S, Z = affine_params(-1.0, 3.0, 0, 255)
x = np.array([-1.0, 0.0, 1.5, 3.0, 5.0])
q = quantize(x, S, Z, 0, 255)
print(S, Z, q, dequantize(q, S, Z))
# 0.01568627...  64  [  0  64 160 255 255]  [-1.0039  0.  1.5059  2.9961  2.9961]

Lưu ý giá trị $5{,}0$ nằm ngoài dải nên bị cắt về $2{,}9961$. Đó là sai số cắt (clipping error), chủ đề của chương sau.