量子化：論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference を読む（２）

「量子化の論文を読む」の2回目です！

Googleが2017年12月にarXivに登録した論文の「Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference」の続きで、今回は2章を読んでいきます。

はじめに
参考サイト
2. Quantized Inference
追記：gemmlowpのソース解析
おわりに

はじめに

量子化の記事一覧です。良かったら参考にしてください。

量子化の記事一覧

・論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference を読む（１）
・論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference を読む（２） ← 今回
・TensorFlow Lite Python で量子化モデルをRaspberry Pi 4で動かす
・TensorFlow Lite C++ で量子化モデルをRaspberry Pi 4で動かす
・Raspberry Pi 4 の TensorFlow Lite C++ を VSCode でリモートデバッグする

参考サイト

arxiv.org

2. Quantized Inference

2.1. Quantization scheme

今回も、すべてChatGPTに丸投げ翻訳です（笑）

このセクションでは、一般的な量子化スキームについて説明します。すなわち、値のビット表現（以下、" $q$ " と表記し、"量子化された値"の意味）と数学的実数値の解釈（以下、" $r$ " と表記し、"実数値"の意味）の対応関係について説明します。私たちの量子化スキームは、推論中に整数のみの算術を使用し、トレーニング中には浮動小数点の算術を使用して実装されます。両方の実装は互いに高い対応性を維持しています。これは、まず我々の量子化スキームの数学的に厳密な定義を提供し、そして整数演算推論と浮動小数点トレーニングの両方でこのスキームを個別に採用することによって達成されます。

重要なことは、学習中はフォワード処理も浮動小数点数を使用しているということですね。一方、推論時は整数のみで演算されます。

私たちの量子化スキームの基本的な要件は、量子化された値に対して整数演算操作のみを使用してすべての算術を効率的に実装できることです（SIMDハードウェア上の純粋な算術と比較してルックアップテーブルを必要とする実装は性能が低い傾向があるため、これらを避けます）。これは、量子化スキームが整数 $q$ から実数 $r$ へのアフィンマッピングであることを要求することと同等であり、すなわち以下の形式であることを要求します。

$r = S(q − Z) \tag1$

これが私たちの量子化スキームであり、定数 $S$ と $Z$ が私たちの量子化パラメータです。私たちの量子化スキームは、各アクティベーション配列内のすべての値と各重み配列内のすべての値に対して1つのセットの量子化パラメータを使用し、別々の配列には別々の量子化パラメータを使用します。

アフィンマッピングとは、線形変換の一種で、元の空間内の点を、別の空間内の点に対応付ける操作です（from ChatGPT）。

量子化操作は、簡単な一次式の変換を使うということですね。

8ビットの量子化では、 $q$ は8ビットの整数として量子化されます（ $B$ ビットの量子化では、 $q$ は $B$ ビットの整数として量子化されます）。いくつかの配列、通常はバイアスベクトルなどは32ビットの整数として量子化されます。詳細はセクション2.4を参照してください。

定数 $S$ （スケール）は任意の正の実数です。これは、通常、実数値 $r$ と同様に、ソフトウェア上で浮動小数点の数量として表されます。セクション2.2では、推論ワークロードでの浮動小数点数量の表現を回避する方法について説明しています。

定数 $Z$ （ゼロポイント）は、量子化された値 $q$ と同じ型であり、実際には実数値 $0$ に対応する量子化された値 $q$ です。これにより、実数値 $r = 0$ が量子化された値で正確に表現できるようになります。この要件の動機は、ニューラルネットワーク演算子の効率的な実装には、しばしば配列の境界周りにゼロパディングが必要であることです。

これまでの議論は、次のような量子化されたバッファデータ構造で要約されます。このようなバッファの1つのインスタンスがニューラルネットワーク内の各アクティベーション配列と重み配列に存在します。型の明確な伝達を可能にするために、C++構文を使用しています。

template<typename QType> // e.g. QType=uint8
struct QuantizedBuffer {
vector<QType> q; // the quantized values
float S; // the scale
QType Z; // the zero-point
};

式 $(1)$ の定義が説明されています。

重要なところは、量子化空間のゼロポイントは、リアル空間の $0$ に対応するというところです。

量子化は、リアル空間のすべての値を量子化空間では表現できません。例えば、リアル空間で非常に近い値の2つの値は、量子化空間では同じ値として表現される可能性が高いです。

量子化により、情報を少なくしている以上、仕方ないことですが、リアル空間の $0$ は、パディングなどで多用されるため特別重要だから、量子化空間で正確に表現できるようにしているということです。

もし、この考慮をしない場合、リアル空間の $0$ は、量子化空間からリアル空間に戻したときに、微妙に $0$ ではない値になってしまいます。

以上が2.1章でした。

2.2. Integerarithmeticonly matrix multiplication

さて、整数演算のみを使用して推論を行う方法について考えてみましょう。つまり、実数値の計算を量子化された値の計算に変換するために方程式 $(1)$ をどのように使用し、後者をどのように設計して、尺度の値 $S$ が整数でなくても整数演算のみを含めるかについてです。

実数 $r1$ と $r2$ （その積は $r3 = r1r2$ で表される）の $N \times N$ 正方行列の乗算を考えます。これらの行列のそれぞれのエントリは $r_\alpha (\alpha=1, 2 \hspace{0.2em} \rm{or} \hspace{0.2em} 3)$ で、 $r_\alpha^{(i, j)}$ （ただし、 $1 \leqslant i, j \leqslant N$ ）と示します。それらが量子化される量子化パラメータは、 $(S_{\alpha}, Z_{\alpha})$ と示します。量子化されたエントリは $q_\alpha^{(i, j)}$ とします。

すると、方程式 $(1)$ は次のようになります：

$r_\alpha^{(i, j)} = S_\alpha(q_\alpha^{(i, j)} - Z_\alpha) \tag2$

行列の乗算の定義から、次のようになります：

$S_3(q_3^{(i, k)} − Z_3) = \sum\limits_{j=1}^{N}S_1(q_1^{(i, j)} - Z_1)S_2(q_2^{(j, k)} - Z_2) \tag3$

これは次のように書き換えられます：

$q_3^{(i, k)} = Z_3 + M\sum\limits_{j=1}^{N}(q_1^{(i, j)} - Z_1)(q_2^{(j, k)} - Z_2) \tag4$

ここで、乗数 $M$ は次のように定義されます：

$M := \dfrac{S_1S_2}{S_3} \tag5$

方程式 $(4)$ では、唯一の非整数は乗数 $M$ です。これは、量子化尺度 $S1, S2, S3$ にのみ依存する定数であり、オフラインで計算できます。我々は実験的に、常に区間 $(0, 1)$ にあることがわかっており、したがって正規化された形式で表すことができます：

$M = 2^{-n}M_0 \tag6$

ここで、 $M_0$ は区間 $[0.5, 1)$ にあり、 $n$ は非負の整数です。正規化された乗数 $M_0$ は、固定小数点乗数（ハードウェアの能力に応じてint16またはint32など）としてうまく表現できます。たとえば、int32が使用される場合、 $M_0$ を表す整数は、 $2^{31}M_0$ に最も近いint32値です。 $M_0 \geqslant 0.5$ であるため、この値は常に少なくとも $2^{30}$ であり、したがって少なくとも30ビットの相対精度を持ちます。したがって、 $M_0$ による乗算は固定小数点乗算として実装できます。一方、 $2^{-n}$ による乗算は効率的なビットシフトで実装できますが、これは正確な最も近い丸め動作を持つ必要があります。この問題については、付録Bで取り上げます。

前半は、 $N \times N$ 行列の $r1$ （の1行目）と $r2$ （の1列目）を積和して、 $r3$ の左上の1要素分の計算過程を示しています。

量子化係数 $S_1, S_2, S_3$ を $M$ にまとめて、後半は $M$ の説明です。

ここは論文を読むだけでは理解できなかったので、gemmlowp の実装を確認しているところです。理解ができたら、ここに追記します。

2.3. Efficient handling of zeropoints

方程式 $(4)$ の評価を効率的に実装するために、 $2N^3$ 回の減算を行う必要がなく、乗算のオペランドを16ビット整数に展開する必要もなくするためには、まず方程式 $(4)$ の乗算を分配することによって、それを次のように書き直すことができることに気付きます：

$q_3^{(i, k)} = Z_3 + M \left( NZ_1Z_2 - Z_1a_2^{(k)} - Z_2a_1^{(i)} + \sum\limits_{j=1}^{N}q_1^{(i, j)}q_2^{(j, k)} \right) \tag7$

ここで、

$a_2^{(k)} := \sum\limits_{j=1}^{N}q_2^{(j, k)}, \hspace{0.5em} a_1^{(i)} := \sum\limits_{j=1}^{N}q_1^{(i, j)} \tag8$

各々の $a_2^{(k)}$ または $a_1^{(i)}$ を計算するには $N$ 回の加算が必要であり、したがって合計で $2N^2$ 回の加算しか必要ありません。方程式 $(7)$ の評価の残りのコストは、ほとんどが次の整数行列の乗算蓄積に集中しています：

$\sum\limits_{j=1}^{N}q_1^{(i, j)}q_2^{(j, k)} \tag9$

この部分は $2N^3$ の算術演算を必要とします。実際、 $(7)$ に関与するすべての要素は $O(N^2)$ で、 $O$ 内の小さな定数が含まれています。したがって、式 $(7)$ への展開と、 $a_2^{(k)}$ と $a_1^{(i)}$ の計算を因数分解することにより、 $N$ が最小値以外の任意のゼロ点を扱うためのオーバーヘッドが少なくなり、問題を同じ核となる整数行列乗算蓄積 $(9)$ に削減します。この核となる乗算蓄積は、他のゼロ点フリーの量子化スキームで計算する必要があるものと同じです。

式 $(4)$ を素直に計算すると、乗算以外に、 $2N^3$ 回の減算が必要になってしまうが、展開することで、 $2N^2$ 回の加算に減らすことができると言っています。

言い回しが分かりにくいですが、細かい計算があるけど、式 $(9)$ の計算が支配的なので、対称量子化とほとんど同じような計算量だと言ってると思います。

2.4. Implementation of a typical fused layer

セクション2.3の議論を続けますが、今回は明示的にすべての関連する数量のデータ型を定義し、量子化された行列の乗算 $(7)$ を修正して、バイアスの加算と活性化関数の評価を直接統合します。この層全体を単一の操作に統合することは、単なる最適化だけではありません。トレーニングで使用される算術と同じ算術を推論コードで再現する必要があるため、推論コード内の統合された演算子の粒度（8ビット量子化入力を取り、8ビット量子化出力を生成する）は、トレーニンググラフの「フェイク量子化」演算子の配置と一致する必要があります（セクション3）。ARMおよびx86 CPUアーキテクチャでの実装では、gemmlowpライブラリ[18]を使用します。そのGemmWithOutputPipelineエントリーポイントは、我々がここで説明する統合された操作をサポートします。

$q_1$ 行列を重みとし、 $q_2$ 行列を活性化とします。重みと活性化の両方はuint8型です（int8を選択することもできますが、適切に修正されたゼロ点を持つ）。uint8値の積を蓄積するには32ビットのアキュムレータが必要であり、アキュムレータにはすぐに明らかになる理由で符号付きの型を選択します。式 $(9)$ の合計は次の形式です：

$\rm{int32} \mathrel{+}= \rm{uint8} × \rm{uint8} \tag{10}$

量子化されたバイアスの加算を、int32バイアスをこのint32アキュムレータに加算することにします。バイアスベクトルは次のように量子化されます：int32をその量子化されたデータ型として使用します。量子化ゼロ点 $Z_{bias}$ として $0$ を使用します。そして、その量子化スケール $S_{bias}$ は、重みと入力活性化のスケールの積であるアキュムレータのスケールと同じです。セクション2.3の表記では、

$S_{bias} = S_1S_2, \hspace{0.5em} Z_{bias} = 0 \tag{11}$

バイアスベクトルは32ビットの値として量子化されますが、ニューラルネットワークのパラメーターのごく一部に過ぎません。さらに、バイアスベクトルの高い精度の使用は実際の必要性に応えます：各バイアスベクトルのエントリが多数の出力活性化に加算されるため、バイアスベクトルの量子化誤差は全体のバイアス（つまり、非ゼロ平均のエラー項）として作用し、エンドツーエンドのニューラルネットワークの精度を保持するために回避する必要があります。

最終的なint32アキュムレータの値を持って、行うべきことが3つ残っています。まず、最終的なスケールにスケールダウンし、次にuint8にダウンキャストし、最終的な8ビット出力活性化を生成するために活性化関数を適用します。

スケーリングダウンは、式 $(7)$ の乗数 $M$ によるものです。セクション2.2で説明したように、これは正規化された乗数 $M_0$ による固定小数点乗算と、ラウンディングビットシフトによって実装されます。その後、uint8に飽和キャストを実行し、範囲 $[0, 255$ ] に飽和します。

我々は、単なるクランプである活性化関数に焦点を当てています。例えば、ReLU、ReLU6などです。数学的な関数については、付録A.1で議論されており、現時点ではそのような層に統合されていません。したがって、我々の統合された活性化関数が行う唯一のことは、最終的なuint8出力活性化を格納する前にuint8値を $[0, 255$ ] のいくつかのサブ区間にさらにクランプすることです。実際には、量子化されたトレーニングプロセス（セクション3）は、出力uint8 $[0, 255$ ] 間全体を利用するよう学習する傾向がありますので、活性化関数はもはや何も行わず、その効果はuint8への飽和キャストに含まれる $[0, 255$ ] へのクランプに取り込まれます。

バイアスはint32で量子化すればよくて（入力と重みの行列演算の結果が32bitなので）、ゼロポイントは $0$ に固定する（32bitあるのでゼロポイントは自由に決めていいということだと思います）と言っています。

十分理解できなかったので、こちらも実装を確認してから、ここに追記しようと思います。

追記：gemmlowpのソース解析

gemmlowp は、（TensorFlow Lite のベースになった？今も使われてる？）TensorFlow Lite のリファレンスコードから呼び出されてる Google の GitHub の行列演算ライブラリです。

論文中の注釈に TensorFlow Lite のリファレンスコードへのリンクが貼られていて、そのソースコードを見ると、gemmlowp が使われていました。

それは、全結合層（FullyConnected）の実装でした。gemmlowp が使われているのは、行列演算が完了して、バイアスが加算された後に呼び出されています。

以下は、抜粋です。（MultiplyByQuantizedMultiplierSmallerThanOne(acc, output_multiplier, output_shift); のところです）

  for (int b = 0; b < batches; ++b) {
    for (int out_c = 0; out_c < output_depth; ++out_c) {
      int32 acc = 0;
      for (int d = 0; d < accum_depth; ++d) {
        int32 input_val = input_data[b * accum_depth + d];
        int32 filter_val = filter_data[out_c * accum_depth + d];
        acc += (filter_val + filter_offset) * (input_val + input_offset);
      }
      if (bias_data) {
        acc += bias_data[Offset(bias_dims, out_c, 0, 0, 0)];
      }
      acc = MultiplyByQuantizedMultiplierSmallerThanOne(acc, output_multiplier,
                                                        output_shift);
      acc += output_offset;
      acc = std::max(acc, output_activation_min);
      acc = std::min(acc, output_activation_max);
      output_data[out_c + output_depth * b] = static_cast<uint8>(acc);
    }
  }

MultiplyByQuantizedMultiplierSmallerThanOne() の定義されている実装を見ます。

inline int32 MultiplyByQuantizedMultiplierSmallerThanOne(
    int32 x, int32 quantized_multiplier, int right_shift) {
  using gemmlowp::RoundingDivideByPOT;
  using gemmlowp::SaturatingRoundingDoublingHighMul;
  return RoundingDivideByPOT(
      SaturatingRoundingDoublingHighMul(x, quantized_multiplier), right_shift);
}

SaturatingRoundingDoublingHighMul() と RoundingDivideByPOT() を見ていきます。

gemmlowp のリポジトリは、以下です。

github.com

まず、SaturatingRoundingDoublingHighMul() です。引数 a は全結合層の計算結果で、引数 b は乗数 $M_0$ です。

乗数 $M_0$ （引数 b）は正と論文に書かれていたので、引数 a の正負で nudge は変わります。引数 a が正の場合は、nudge は 0x40000000 で、負の場合は、nudge は 1-0x40000000

// This function implements the same computation as the ARMv7 NEON VQRDMULH
// instruction.
template <>
inline std::int32_t SaturatingRoundingDoublingHighMul(std::int32_t a,
                                                      std::int32_t b) {
  bool overflow = a == b && a == std::numeric_limits<std::int32_t>::min();
  std::int64_t a_64(a);
  std::int64_t b_64(b);
  std::int64_t ab_64 = a_64 * b_64;
  std::int32_t nudge = ab_64 >= 0 ? (1 << 30) : (1 - (1 << 30));
  std::int32_t ab_x2_high32 =
      static_cast<std::int32_t>((ab_64 + nudge) / (1ll << 31));
  return overflow ? std::numeric_limits<std::int32_t>::max() : ab_x2_high32;
}