量子化について

ideinのクリスチャンです。機械学習周りの研究開発を担当しています。

edge-computing で機械学習と言えば、量子化を利用して処理を加速する人は多いと思われます。（n.b.: 現時点の Actcast で量子化は利用されてませんが）

この記事では量子化アプローチとこの開発ノウハウをざっと紹介します。 

量子化とは

機械学習で行われている量子化の実際は activations 又は weights で通常用いられる32bit 数値をbitwidth の少ない数値で代替することです。 メリットとしては、モデルサイズが小さくなるだけではなく、行列乗算を（なので畳込みも）速く処理できます。デメリットは勿論、モデルの精度が低下してしまう恐れがあることです。

どうやってビット数の少ない計算を早くするのは hardware-dependent なので、この周りの説明をするつもりはありません。

量子化をする方法の一つとして、単純に学習済みモデルの weights を取って、量子化関数を作用することです。推論の時、 activation に同じ関数を作用するだけのも可能ですが、このアプローチを使うと、モデルの精度は随分下がってしまいます。 この方法よりも器用なやり方は、学習済みモデルで学習データセットを推論することです。こうすると、レイヤーのactivations 毎に分布の平均と分散を確認できて、この情報で相応しい量子化 parameter を設定できます。

一方で、最も広く研究されている方法であり、量子化を意識している学習を使えます。実用的には、学習の時にも量子化を行って、back-propagation の時に上手く勾配を推定することです。（量子化の関数は微分不可能または微分は殆どゼロ可能性ありますので、ここで注意しないと行けません）

アプローチの紹介

良い参考になると思いますので、有名な３つアプローチを紹介します。

DoReFa-Net

この論文で、activation とweight だけでなく、勾配も量子化する手法が提案されています。こうしたら、学習も加速できるので、edge-device とかでも学習は可能となります。 それぞれのbitwidth をセットして、速度と精度の釣り合いを細かく設定出来ます。量子化数値の処理は以下のようになります：

x  \\ float32 行列
k  \\ bitwidth 
q_x = round((2^k - 1) * x) / (2^k - 1)

このアプローチで bit-width は量子化の唯一hyper parameter です。（実はweightとactivationとgradientそれぞれのパラメターあるので、３つありますが）

Integer-Arithmetic-Only Inference (gemmlowp)

tensorflow lite の組み込みquantization です。 DoreFa-Netと違って、bit-width の他にも量子化パラメターがあって、activation の分布の間の極値です。学習の時にこのparameter の調整はback-propagation でなくて、移動平均で行います。 このアプローチで整数処理で計算を加速することは目的なので、対象デバイスで使える整数型（uint8等）を設定します。（推論の時、デバイスの整数blasカーネルを呼び出したいので）bit-widthを選考するよりも、整数型で設定されてます。 量子化数値の処理は以下のようになります：

x  \\ float32 行列
k  \\ bitwidth
a  \\ float32 数値 (学習せれてる下限)
b  \\ float32 数値 (学習されてる上限)

clamp(r) -> min(max(r, a), b)
s = (b - a) / (2^k - 1)

q_x = round((clamp(x) - a) / s )  * s + a

LQ-Nets

LQはLearning Quantizationの略です。上のアプローチで、k と a と bを設定したら、可能な量子化数値は {a + i * (b-a)/(2**k - 1) | i∈ [1, 2^k] }（即ち、実数から整数への変換はただのアフィン写像） のセットに限られています。 LQ-Netsではこの可能量子化数値を可変として、学習しようと提案されてます。2つparameter で設定されてなくて（aとb）、vector（size: k）で設定されてます。このベクトルを学習するため、 新たなアルゴリズム(Quantization Error Minimization)を導入されています。 量子化数値の処理は以下のようになります：

x  \\ float32 行列
k  \\ bitwidth
l_i, i=1..2^k  \\ float32 
v  \\ float32 vector (size: k)

t_i = (l_(i-1) + l_i) / 2, i=1..2^k
j = argmin_i(x <= t_(i+1))

q_x = dot(v, bit_repr(j))

実装アドバイス

実感として、論文に導入されたアプローチを実装しようとしたり、論文では評価されていないモデルの量子化をしてみたりすると、結果はイマイチなことが多いです。 ML枠組みで新しい量子化レイヤーの実装とモデルの適応に関して、ざっとノウハウを共有したいと思います。 推論処理の最適化はハードウェア依存なので、以下に記載してません。

  • 移動平均で学習されてる parameter を warm-up

量子化特定のparameter は back-propagation でなく、移動平均で学習されてることはあります。 この場合、初期設定が理想的な数値と幅広くずれたら、モデルが精度の良くない状態へ収束する恐れがあります。このため、最初の数epoch は量子化無しで学習するのはオススメです。

  • 量子化直前にbatch norm

様々な論文に書いてあると思いますが、activation の分布を狭い間にあるために、量子化レイヤーの前に batch normalization を取り入れると誤差は減少します。 一般的に、レイヤー順番をBatch Normalization->(k-bits) Activation-> Convolution -> Poolingにするのはオススメです。

  • CPUでもいい

機械学習と言えば、GPUで計算を速くするのは当たり前です。ただし、量子化特定の処理のある部分はもしかすると軽いため、周波数が高いCPUを利用した方が学習速くなる場合があります。 それを判断するには、見比べるしかなさそうです。

  • 全レイヤーを量子化？

軽さと速度のため、全てのレイヤーを量子化したいものの、最初レイヤーと最後レイヤーの量子化は精度に大幅に影響あります。 trade-offだと思うのですが、精度を上がるように、最初と最後のlayer をfp32のままにしても良いかもしれません。 それに、モデル内で最初の layer の channel 数は少ないので、量子化しても加速はそんなにないと思われます。

  • 入力を正規化

最初のレイヤーは量子化されたら、このレイヤーの入力はモデルの入力自体です。画像系のモデルで、数値は広く使われているuint8 の RGB (チャンネル３つ、数値 0~255)だとすると、 分布の間は広くて、量子化したら情報を失うことは多い場合があります。（量子化アプローチによって） 精度に結構に影響あるので、入力を正規化した方がいいかと思います。（もしくは、上に記載の通り、最初のレイヤーを量子化しない）

  • weightsよりactivations

論文にも言及されてて、weight の量子化よりも activation のが精度に影響あります。それぞれのbit-width を設定出来ますので、 act_bitwidth >= weight_bitwidth を薦めします。

最後に

上に紹介したapproachチだけでなく、様々な面白い量子化系の論文がちょいちょい投稿されてます。 使っているMLフレームワークで再現実装する場合が良くあるので、ご参考として読んでいただければと思います。

ビット数の少ない数値で速いモデルを作りましょう！

Jan. 23 2019