Quantization

馒头and花卷發表於2024-12-04

目錄
  • 引
  • 線性量化 (Linear Quantization)
    • 對稱量化
    • 非對稱量化
  • 非線性量化
    • Power-of-X
  • Rounding
    • Deterministic rounding
    • Stochastic rounding

[1] 進擊的程式猿-模型壓縮-神經網路量化基礎.

[2] Przewlocka-Rus D., Sarwar S. S., Sumbul H. E., Li Y. and De Salvo B. Power-of-two quantization for low bitwidth and hardware compliant neural networks. tinyML Research Symposium, 2022.

[3] Li Y., Dong X. and Wang W. Additive powers-of-two quantization: an efficient non-uniform discretization for neural networks. ICLR, 2020.

[4] Xia L., Anthonissen M., Hochstenbach M. and Koren B. Improved stochastic rounding. 2020.

引

整理了一下量化中的基礎概念和方法. 其實, 在計算機中, 我們所見之浮點數也只不過是真正數字的一個高精度量化 (e.g, FP32), 但是隨著模型的逐步增加, 我們需要將模型進一步'精簡'. 當然了, '精簡'的法子有很多: 裁剪 (pruning), 蒸餾 (distillation), 包括這裡講到的量化 (quantization). 個人感覺, 相較於裁剪和蒸餾, 量化是一種即插即用的方法 (雖然會有一些追求極限精度的量化方法需要一些校準). 本文討論如下的例子:

\[\underbrace{[-3.1, -0.03, 0.1, 1.2]}_{x_f} \mathop{\longrightarrow} \limits^{Q} [...] \subset \underbrace{[-128, 127]}_{\text{Int8}}/\underbrace{[0, 255]}_{\text{UInt8}}. \]

線性量化 (Linear Quantization)

對稱量化

  • \(x_q = Q(x_f)\) 的對稱量化過程如下:

    1. 確定 \(x_f\) 的(絕對值)範圍: \(\Delta_f = \max_i(|x_f[i]|)\);
    2. 確定 \(x_q\) 的範圍: \(\Delta_q = 2^{N-1} - 1\), \(N\) 表示量化後的精度, 比如 Int8 時 \(\Delta_q = 127\).
    3. 量化:

      \[x_q = Q(x_f) := \text{round}(\frac{x_f}{\Delta_f} \cdot \Delta_q). \]

  • 如上圖所示, \(\Delta_f\) 首先將 \(x_f\) 壓縮到 \([-1, 1]\) 之中, 然後再對映回 \([-128, 127]\) 的量化後的精度中.


def quant(f: torch.Tensor, N: int = 8):
    delta_f = f.abs().max()
    delta_q = 2 ** (N - 1) - 1
    return (f * delta_q / delta_f).round().to(torch.int), delta_f
  • 利用上述程式碼可得:

\[[-127, -1, 4, 49]. \]

  • 反解的過程:

    \[\hat{x}_f = Q^{-1}(x_q) := \frac{x_q}{\Delta_q} \cdot \Delta_f. \]


def dequant(q: torch.Tensor, delta_f: torch.Tensor, N: int = 8):
    delta_q = 2 ** (N - 1) - 1
    return (q * delta_f / detal_q).float()
  • 得到:

    \[[-3.1000, -0.0244, 0.0976, 1.1961]. \]

非對稱量化

  • \(x_q = Q(x_f)\) 的非對稱量化過程如下:

    1. 確定 \(x_f\) 的(絕對值)範圍: \(\Delta_f = x_f^{\max} - x_f^{\min}\);
    2. 確定 \(x_q\) 的範圍: \(\Delta_q = 2^{N} - 1\), \(N\) 表示量化後的精度, 比如 UInt8 時 \(\Delta_q = 255\).
    3. 確定零點偏移量: \(x_{offset} = -x_f^{\min}\)
    4. 量化:

      \[x_q = Q(x_f) := \text{round}(\frac{x_f + x_{offset}}{\Delta_f} \cdot \Delta_q). \]

  • 注意到:

    \[x_f + x_{offset} = x_f - x_f^{\min} \in [0, \Delta_f]. \]


def quant(f: torch.Tensor, N: int = 8):
    offset = -f.min()
    delta_f = f.max() + offset
    delta_q = 2 ** N - 1
    return ((f + offset) * delta_q / delta_f).round().to(torch.int), delta_f, offset

  • 透過如上程式碼可得量化結果:

    \[[0, 182, 190, 255] \]

  • 透過如下方式進行反解:

    \[\hat{x}_f = Q^{-1}(x_q) := \frac{x_q}{\Delta_q} \cdot \Delta_f - x_{offset}. \]


def dequant(q: torch.Tensor, delta_f: torch.Tensor, offset: torch.Tensor, N: int = 8):
    delta_q = 2 ** N - 1
    return (q * delta_f / delta_q).add(-offset).float()
  • 得到:

    \[[-3.1000, -0.0310, 0.1039, 1.2000] \]

非線性量化

  • 其實, 一般的量化都可以用如下的過程統一表示:

    1. 將 \(x_f\) normalize 到 \(\phi(x) \in [0/-1, 1]\) 區間內;
    2. 在該區間內設定 \(2^N\) (或者 \(2^N - 1\)) 個點, 記為:

      \[\mathcal{Q} = \{q_i\}_{i \in \mathcal{I}}. \]

    3. 尋找最近鄰的作為量化結果:

      \[Q(x) = \mathop{\text{argmin}} \limits_{i \in \mathcal{I}} |\phi(x) - q_i|. \]

  • 之前的線性量化實際上就是:

    \[\mathcal{Q} = \{i / (2^{N - 1} - 1)\}_{i \in \mathcal{I}}, \\ \mathcal{I} = \{-(2^{N-1} - 1), \ldots, 0, 1, \ldots, 2^{N-1} - 1\}, \]

    和

    \[\mathcal{Q} = \{i / (2^{N} - 1)\}_{i \in \mathcal{I}}, \\ \mathcal{I} = \{0, 1, \ldots, 2^{N} - 1\}. \]

  • 而 round 就是一種自動的 nearest 的搜尋方式.

  • 當然了, \(\mathcal{Q}\) 的分佈不必像線性量化一樣那樣的均勻, 實際上這種分佈往往不是最優的. 容易感覺到, 最優的分佈應當在 \(x_f\) 中比較稠密的範圍分佈較多的量化點, 否則會造成大量的浪費. 當然了, 這種非線性量化的方式相較於線性量化有它的不足之處, 就是得根據資料'找'一個合適的分佈, 這個'找'的過程往往是比較耗時的.

Power-of-X

  • Power-of-X 的特點就是 \(\mathcal{Q}\) 的分佈是服從指數分佈的:

    \[\mathcal{Q} = \{\alpha^{i-1} \}_{i \in \mathcal{I}} \cup \{\mathcal{Q}_0 = 0\}, \\ \mathcal{I} = \{1, \ldots, 2^{N} - 1\}. \]

    其中 \(\alpha \in (0, 1)\).

  • 其量化過程為:

    1. 確定 \(x_f\) 的(絕對值)範圍: \(\Delta_f = \max_i(|x_f[i]|)\);
    2. 確定 \(x_q\) 的範圍: \(\Delta_q = 2^{N-1} - 1\), \(N\) 表示量化後的精度, 比如 Int8 時 \(\Delta_q = 127\).
    3. 量化:

      \[x_q = Q(x_f) := \text{sign}(x_f) \cdot \text{clip} \Bigg( \text{round}\bigg(\log_{\alpha} \bigg(\frac{|x_f|}{\Delta_f}\bigg) + 1 \bigg), 0, \Delta_q \Bigg). \]

    4. 反解的過程為:

      \[\hat{x}_f = \text{sign}(x_q) \cdot \alpha^{|x_q| - 1} \cdot \Delta_f. \]

注: 這裡 \(+1, -1\) 是為了將 \(0\) 保留給真正的 0 值.


def quant(f: torch.Tensor, alpha: float, N: int = 8):
    sign, f = f.sign(), f.abs()
    delta_f = f.max()
    delta_q = 2 ** (N - 1) - 1
    logalpha = math.log2(alpha)
    return (f / delta_f).log2().div(logalpha).add(1).round().clip(0, delta_q).mul(sign).to(torch.int), delta_f

def dequant(q: torch.Tensor, alpha: float, delta_f: torch.Tensor, N: int = 8):
    sign, q = q.sign(), q.abs()
    return (alpha ** q.add(-1)).mul(delta_f).mul(sign)

  • 如上圖所示, \(\alpha \rightarrow 0\), 不均勻性大大增加. 這個性質在有些時候會很有用, 因為很多資料的分佈是形容高斯分佈的, 此時採用不均勻的會更加高效.

  • 需要宣告的一點是, \(\alpha \rightarrow 1\) 的時候, 整體會顯得更加均勻一點, 但是並不能等價線性量化, 注意到:

    \[\alpha^{2^N} \mathop{\longrightarrow} \limits^{\alpha \rightarrow 1} 1. \]

    這意味著, 除了 \(0\) 之外, \(\mathcal{Q}\) 所能表示的最小值會隨著 \(\alpha\) 增加逐步增大. 這其實是另一種極為嚴重的 '不均勻' 性質.

  • 此外 \(\alpha = 0.5\) 是比較特殊的一種情況, 可以透過位操作加速.

Rounding

有些時候 Rounding 的方式也很重要, 畢竟它是誤差的來源.

Deterministic rounding

Stochastic rounding

  • 隨機 rounding 形式如下:

    \[\text{round}(x) = \lfloor x \rfloor + \mathbb{I}\big[x - \lfloor x \rfloor \xi \big], \quad \xi \sim \mathcal{U}([0, 1]). \]

    其中 \(\lfloor x \rfloor\) 表示 round down 操作.

  • 舉個例子, 對於 \(0.4\), 它有 0.4 的機率為 \(1\) 有 \(0.6\) 的機率為 0. 對於 \(-1.6\), 它有 \(0.6\) 的機率為 \(-2\), 有 \(0.4\) 的機率為 \(-1\).

  • 一個比較好的性質是:

    \[\begin{array}{ll} \mathbb{E}[\text{round}(x)] &= \lfloor x \rfloor + \mathbb{E}\bigg[\mathbb{I}\big[x - \lfloor x \rfloor \xi \big] \bigg] \\ &= \lfloor x \rfloor + 1 \cdot (x - \lfloor x \rfloor) \\ &= x \end{array}. \]

相關文章