深層学習の量子化に入門してみた〜理論編〜

こんにちは。リサーチャーの勝又です。私はレトリバで自然言語処理、とくに要約や文法誤り訂正に関する研究の最新動向の調査・キャッチアップなどを行っております。

最近、深層学習の量子化について勉強する機会があったので、この記事では量子化の理論的な話をまとめてみようと思います。

深層学習の量子化が流行っている？
量子化の簡単な解説
量子化の種類
まとめ

深層学習の量子化が流行っている？

近年、自然言語処理の分野では事前学習を利用した研究が盛んに行われています。有名な事前学習モデルとしては、BERT（Bidirectional Encoder Representations from Transformers）¹と呼ばれるモデルが盛んに研究されています。事前学習モデルの良いところとして、アノテーションされていないテキストデータを用いて事前学習を行うことで、実際に解きたい課題の精度が向上することが挙げられます。

この事前学習モデル、確かに精度は高いのですが、実運用を考えた際に、いくつかの扱いづらさを抱えています。代表的なものとして、「モデルの大きさ」と「推論速度」が挙げられます。量子化と呼ばれる技術を用いることで、これらの問題を緩和できます。

今回の記事では、この量子化の簡単な説明をしようと思います。

量子化の簡単な解説

量子化に関してはすでに多くの方がまとめてくださっています。今回の記事では、その中でも個人的に一番わかりやすかったWuらの論文²に基づいて、量子化の簡単な説明³をしていきたいと思います。

どうやって値を低精度に変換するのか

ニューラルネットの量子化は、通常では高精度の値で演算を行うところを、工夫して低精度の値で行っています。たとえば、本来fp32だった値を、量子化してint8で計算する、といった具合です。では、どのようにして高精度の値を低精度に変換しているのでしょうか。

たとえば、何かしらの実数値 $x$ を $b$ -bitの値 $x_q$ に変換したいとします。その際、量子化では実数値 $x$ を含む区間 $x \in [\beta, \alpha$ ]を $x_q \in [-2^{b-1}, 2^{b-1}-1$ ]に変換することを考えます。この変換を一様に行う場合は量子化を次の関数 $f(x)$ で表すことができます。

$\begin{eqnarray} f(x) = s \cdot x + z \tag{1} \end{eqnarray}$

では、この $s, z$ の値はどう決めるのでしょうか？ Wuらの論文では次の2種類が解説されています。

Affine Quantization
Symmetric Quantization

それぞれ順に解説します。

1. Affine Quantization

実数値 $x \in \mathbb{R}$ を符号付きintである $b$ -bitの値 $x_q \in \{-2^{b-1}, -2^{b-1}+1, \dots, 2^{b-1}-1 \}$ に変換する際、Affine Quantizationでは以下のようになります。

$\begin{eqnarray} s = \frac{2^b - 1}{\alpha - \beta} \tag{2} \end{eqnarray}$

$\begin{eqnarray} z = -\mathrm{round}(\beta \cdot s) - 2^{b-1} \tag{3} \end{eqnarray}$

ここでの $s$ はscale factor、 $z$ はzero-pointと呼ばれます。たとえば8-bitの場合、 $s = \frac{255}{\alpha - \beta}$ で $z = -\mathrm{round}(\beta \cdot s) - 128$ となります。

このように、Affine Quantizationで量子化を行う場合、scale factorやzero-pointは量子化を行う対象の範囲 $[\beta, \alpha$ ]、および量子化後のbitで決定されます。量子化の際は、実数値の入力 $x$ に対してこのscale factorとzero-pointを用いて変換し、 $\mathrm{round}$ で値を丸め込む操作を行います。具体的には次の通りです。

$\begin{eqnarray} x_q = \mathrm{quantize}(x, b, s, z) = \mathrm{round}(s \cdot x + z) \tag{4} \end{eqnarray}$

この量子化した値 $x_q$ を元に戻す操作（ $\mathrm{dequantize}$ ）も $s$ と $z$ を用いることで実現できます。

$\begin{eqnarray} \hat{x} = \mathrm{dequantize}(x_q, s, z) = \frac{1}{s}(x_q - z) \tag{5} \end{eqnarray}$

当然量子化の際に $\mathrm{round}$ で値を丸めてしまっているので、上記 $\mathrm{dequantize}$ を行っても、元の入力 $x$ とまったく同一の値を得ることは難しく、そのため上式では近似値 $\hat{x} \approx x$ までしか求められていません。

2. Symmetric Quantization

Symmetric QuantizationはAffine Quantizationに以下の制約を加えたものになります。

zero-point $z$ は0
$\beta = - \alpha$ ; つまり、量子化する実数値の区間を0で対称にする⁴

先ほどのAffine Quantizationと比較すると以下の図のようになります。

f:id:ssskkk420:20220121174929p:plain — 図1: Affine QuantizationとSymmetric Quantizationによるint8への量子化概要図。

この図の通り、Affine Quantizationでは入力 $x$ が0だった際は、量子化後は $z$ だったのに対して、Symmetric Quantizationでは量子化後も0になっています。また、Symmetric Quantizationは量子化を行う実数値区間が0で対称になっていることもわかります。

Symmetric Quantizationの制約をAffine Quantizationの式に反映させると以下のようになります。

$\begin{eqnarray} s = \frac{2^{b-1} - 1}{\alpha} \tag{6} \end{eqnarray}$

$\begin{eqnarray} x_q = \mathrm{quantize}(x, b, s) = \mathrm{round}(s \cdot x) \tag{7} \end{eqnarray}$

$\begin{eqnarray} \hat{x} = \mathrm{dequantize}(x_q, s) = \frac{1}{s}x_q \tag{8} \end{eqnarray}$

制約を入れたことで、Symmetric QuantizationはAffine Quantizationと比べて少しスッキリした形になりました。

演算をIntで行う

それでは、実際にこの変換を適用して高精度な演算を低精度な演算で近似する様子を眺めてみましょう。例として、線形変換 $Y = XW$ を低精度で近似してみたいと思います。ここで、 $X = (x_{ik}) \in \mathbb{R}^{m \times p}$ は入力、 $W = (w_{kj}) \in \mathbb{R}^{p \times n}$ が重み、 $Y = (y_{ij}) \in \mathbb{R}^{m \times n}$ が出力を表しています。

この「線形変換を低精度な演算で近似する」ということは、具体的には入力 $X$ と重み $W$ を量子化することになります。 Symmetric Quantizationを適用し、低精度な値 $X_q = (x_{q, ik}) \in \mathbb{Z}^{m \times p}$ と $W_q = (w_{q, kj}) \in \mathbb{Z}^{p \times n}$ に量子化した後、これらの値を用いて $y_{ij}$ を近似的に求めた式が以下の通りとなります。

$\begin{eqnarray} y_{i j}&=&\sum_{k=1}^{p} x_{i k} \cdot w_{k j} \\ &\approx& \sum_{k=1}^{p} \mathrm{dequantize}\left(x_{q, i k}, s_{x}\right) \cdot \mathrm{dequantize}\left(w_{q, k j}, s_{w}\right) \\ &=&\sum_{k=1}^{p} \frac{1}{s_{x}} x_{q, i k} \cdot \frac{1}{s_{w}} w_{q, k j} \tag{9} \end{eqnarray}$

ここでの $s_x$ や $s_w$ はそれぞれ入力、重みに対するscale factorを表しています。⁵ (9)式を少し整理すると、以下のようになります。

$\begin{eqnarray} \frac{1}{s_{x} \cdot s_{w}} \sum_{k=1}^{p} x_{q, ik} \cdot w_{q, kj} \tag{10} \end{eqnarray}$

この式の通り、内部としては低精度な値 $x_{q, ik}$ と $w_{q, kj}$ で計算を行い、その結果を $\mathrm{dequantize}$ を用いて高精度な値に近似しています。

Affine QuantizationとSymmetric Quantizationの違い

Symmetric Quantizationではなく、Affine Quantizationを用いて線形変換を近似するとどうなるでしょうか。実際にやってみるとzero-pointが増えるため、以下のようになります。

$\begin{eqnarray} y_{i j} &\approx& \sum_{k=1}^{p} \frac{1}{s_{x}}\left(x_{q, i k}-z_{x}\right) \frac{1}{s_{w}}\left(w_{q, k j}-z_{w, j}\right) \\ &=&\frac{1}{s_{x} s_{w}}\left(\sum_{k=1}^{p} x_{q, i k} w_{q, k j}-\sum_{k=1}^{p}\left(w_{q, k j} z_{x}+z_{x} z_{w}\right)-\sum_{k=1}^{p} x_{q, i k} z_{w}\right) \tag{11} \end{eqnarray}$

上記(11)式をSymmetric Quantizationを用いた場合の(10)式と比べると、zero-pointに関する項が増えていることがわかります。第1項はSymmetric Quantizationにも存在している部分です。第2項、第3項が新たに含まれる部分で、Affine QuantizationとSymmetric Quantizationの違いとなっています。第2項 $\sum_{k=1}^{p}(w_{q, k j} z_{x}+z_{x} z_{w})$ は、量子化された重み $w_{q, kj}$ と、 $X$ と $W$ の量子化に使用するzero-pointで演算を行っています。これらの値は一度zero-pointさえ決めてしまえば入力 $X$ によらず計算できるので、実際の演算の際には事前に計算した値を使いまわすことができます。一方で、第3項 $\sum_{k=1}^{p} x_{q, i k} z_{w}$ は入力 $X$ を量子化した値 $x_{q, ik}$ に依存しているため、こちらは事前に計算を行うことができず、入力が与えられるたびに計算することになります。

このように、Affine QuantizationはSymmetric Quantizationと比べて、計算上でオーバーヘッドが発生します。そのため、演算の高速化に焦点を当てるとAffine QuantizationよりはSymmetric Quantizationの方が優れていると考えられているようです。実際、今回紹介しているWuらの論文では以下のように表記されています。（ここでのscale quantizationは、本記事でのSymmetric Quantizationを指しています。）

Thus, to maximize inference performance we recommend using scale quantization for weights.

量子化したい値が範囲を超えてしまったら

本記事の冒頭から、量子化はある区間の高精度な値 $[\beta, \alpha$ ]を低精度の値に変換すると説明してきました。では、入力の値がこの範囲を外れてしまったらどうなるのでしょうか。先の図1を見ていただけるとわかるのですが、値が外れた場合、量子化後の値は一番近い端の値になるように $\mathrm{clip}$ を行います。たとえば、Symmetric Quantizationの(7)式は次のようになります。

$\begin{eqnarray} x_q = \mathrm{quantize}(x, b, s) = \mathrm{clip}(\mathrm{round}(s \cdot x), -2^{b-1}+1, 2^{b-1}-1) \tag{12} \end{eqnarray}$

このため、入力の値が想定していた範囲を超えてしまった場合、量子化後は一定の値になってしまい、どのくらい元の値が大きかったのか、といった情報が消えてしまいます。これを防ぐためには、入力の範囲 $[\beta, \alpha$ ]を大きくする方法が考えられると思います。しかしながら、入力の範囲を大きくしてしまうと、今度は範囲内の値の細かい区別がつかないようになっていきます。入力 $X$ には何が来るのかわからないため、適切にこの範囲を決定する必要があります。この範囲の決め方についてはWuらの論文でいくつか実験しているため、気になる方はそちらを確認していただければと思います。

量子化の種類

ここまで、実際に量子化を行って低精度で演算する方法を確認してきました。ここからは、どのタイミングで量子化、とくにSymmetric Quantizationで使用するscale factorを決定するか説明していきます。

タイミングとしては、以下の2通りがよく知られていると思います。

Dynamic Quantization
Static Quantization（Post Training Quantization）

2番のStatic Quantizationに関しては、文献によってはPost Training Quantizationと呼称されていることがあります。⁶

ちなみに、重みの量子化についてはあらかじめ行っておくことが可能です。つまり、重みの値からscale factor $s_w$ をあらかじめ決め、重みを量子化しておくことでモデルのデータサイズを小さくできます。そのため、考える必要があるのは入力の値のscale factor $s_x$ をいつ決定するかです。

それでは、それぞれの量子化タイミング手法について解説します。

1. Dynamic Quantization

Dynamic Quantizationは、その名前の通り、入力 $X$ ごとにscale factor $s_x$ を決める方法です。入力の度にscale factorを決めているため、その分が後述のStatic Quantizationと比べると余計な処理になっています。一方で、その入力に適したscale factorを使用することができるので、Static Quantizationと比べて、モデルの性能の劣化は少ないです。

2. Static Quantization（Post Training Quantization）

Static Quantizationは、Dynamic Quantizationと違って、事前に入力 $X$ に対応したscale factor $s_x$ を計算し、実際の推論時にはその値を使いまわします。このscale factorの計算は、実データを用いて行います。 Dynamic Quantizationは逐次scale factorを計算していましたが、Static Quantizationでは事前計算しているので、その分推論速度が速くなります。

一方で、Dynamic Quantizationと比べたデメリットもあります。 1つ目のデメリットとして、scale factorの計算は、実データを用いて行うので、Dynamic Quantizationと比べてStatic Quantizationはモデルの学習データ（ただし入力側のみで十分です）が必要になります。また、2つ目のデメリットが、実際の推論時の入力がscale factorの事前計算時に使用した入力と全然違う場合、うまく量子化されない可能性があります。

Quantization Aware Training

ここまでの説明した通り、量子化はそもそも高精度な演算を低精度に近似して行っているため、その分性能が劣化してしまいます。この量子化由来の性能劣化を軽減する手法の1つにQuantization Aware Trainingがあります。

この手法は、量子化での近似に伴う数値のズレをモデルに考慮してもらう、というものです。たとえば、以下の式のように $\mathrm{quantize}$ の出力に $\mathrm{dequantize}$ を施した際、元の値である $x$ と処理後の値 $\hat{x}$ は異なる値になってしまいます。

$\begin{eqnarray} \hat{x} = \mathrm{dequantize}(\mathrm{quantize}(x, b, s), b, s) \tag{13} \end{eqnarray}$

この $x$ と $\hat{x}$ のズレをモデルに考慮してもらう必要があります。どのように考慮させるかというと、上記の(13)式のような操作（fake quantization）をモデルに入れ、実際に量子化での近似に伴う数値のズレを発生させて、モデルの学習⁷を行います。ちなみに、この学習時に量子化は行うのですが、データタイプはfloatのまま扱います。

実運用でこのQuantization Aware Trainingを使うことを考えてみます。先ずはscale factorを求める必要があるので、Static Quantizationと同様に学習データの入力側データを使用して、適したscale factorを決定します。その後Quantization Aware Trainingを行うのですが、この学習は学習データの入出力を使用します。このように、通常のStatic Quantizationと比較するとQuantization Aware Trainingは学習データの入力側だけでなく、出力側も必要になります。

量子化のタイミングについてまとめると......

ここまで、Dynamic Quantization、Static Quantizationの違いと、量子化での性能劣化に対応したQuantization Aware Trainingを確認しました。改めて、これらのscale factorを決定するタイミングやその他の特徴を次の表にまとめました。

f:id:ssskkk420:20220127164528p:plain — 表1: 量子化タイミングなどに関してまとめた表。

この表の通り、Dynamic Quantizationは量子化にとくにデータを必要としない分、一番お手軽に試すことができます。モデルのデータサイズを小さくすること自体はどの手法でも可能なので、それが目的ならばDynamic Quantizationがオススメです。一方で、速度を追い求めるならStatic QuantizationやQuantization Aware Trainingといった、事前にscale factorを計算する方式が良さそうです。性能を保ちたい、というところであればQuantization Aware Trainingが一番確実なのかなとは思います。⁸

まとめ

今回の記事ではWuらの論文を中心に、ニューラルネットの量子化について簡単ではありますがご紹介させていただきました。今回はscale factorの具体的な決め方や、そもそも量子化でどのくらいモデルの性能が劣化するのかなどはとくに解説していません。これらの内容をもっと知りたい方は是非、Wuらの論文をご覧ください。他にも個人的な感想ではありますが、こちらの英語ブログ記事も参考になりました。興味がある方は是非ご覧ください。

今後の記事では、実際にモデルを量子化してみてどのくらい推論速度が向上するのかなどについて検証したいと思います。

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩
Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation. Hao Wu, Patrick Judd, Xiaojie Zhang, Mikhail Isaev, Paulius Micikevicius. [paper]↩
この論文では一様な量子化について解説しています。同様に、本記事でも一様な量子化、つまり高精度な値を低精度な値に一様変換する方法について解説します。↩
0で対称にするため、Affine Quantizationと比べて量子化後の値が1つ使えなくなることに注意してください。↩
Wuらの論文でも線形変換に量子化を適用した例が記載されているのですが、この記事とscale factorの添え字が少し異なっています。Wuらはscale factorとしてtensor単位なのか、それとももっと細かい粒度なのかに関しても論じているためです。気になる方は是非Wuらの論文をご確認ください。↩
実際、Wuらの論文ではPost Training Quantizationとして記述されています。しかし、PyTorchのドキュメントなどではStatic Quantizationの呼称がされていることもあり、本記事ではそちらを呼称しています。↩
量子化を実際に行って学習を行うのですが、そもそも量子化は内部で $\mathrm{round}$ を使っているため、微分ができません。そのため、学習時のbackwordの処理の際にはStraing-through Estimator（STE）と呼ばれる近似を用いて対応します。詳細はWuらの論文をご確認ください。↩
Dynamic Quantizationも性能劣化が少ないとしていますが、これは量子化に伴う近似計算のズレの話ではなくて、入力に対して逐次scale factorを計算しているため、他の2手法と比べてその点で性能劣化しずらい、ということを意味しています。↩