論文解説|I-ViT:Integer-only Quantization for Efficient Vision Transformer…

I-ViT: Integer-only Quantization for Efficient Vision Transformer Inference

はじめに

Vision Transformers (ViTs)は、コンピュータビジョンで最先端のパフォーマンスを達成しましたが、ストレージと計算のオーバーヘッドが大きく、スマホなどの小型デバイスでの効率的な推論は難しいです。従来手法のFasterTransfomer[1]は入力と出力だけを量子化(整数値)に変換して計算量を削減していました。この論文では、モデル全体で量子化を行うことで、高速な推論を実現しています。整数値でどのようにLayer Normや活性化関数を表現するか、などがこの提案手法の見どころです。

従来手法との比較

事前知識

量子化の方法はFasterTransfomer[1]やCNN[2]などで提案されており本論文でもこれらを用いて整数化を行っています。浮動小数点を$R$、整数値を$I$、$R$がとり得る最小-最大値を$m$^[1]、整数型のバイト数を$k$(int32だったら32、int8だったら8)とすると変換式は以下のようになります。

$$I = \lfloor \frac{clip(R,-m,m)}{S}\rceil \\ S = \frac{2m}{2^k-1}$$

ただし、$\lfloor \cdot \rceil$はround(四捨五入)を示しています。
試しに$R=3.141592$、$m=10$、整数型をint8として変換してみます。

$$S = \frac{2*10}{2^8-1} = 0.0784313...\\ I = \lfloor \frac{3.141592}{S}\rceil I = \lfloor 40.055298...\rceil = 40$$

となります。このようにして浮動小数点から整数に変換できます。今回の論文のViT内では$S$と$I$の値を使い、処理を行っています。

提案手法

Transformerの構造は変わらずに浮動小数点を整数値にしたようなモデル構造になっています。本論文ではすべての処理を整数型で行うために4つの手法を提案しています。

1. 従来手法のダイアディック算術(Dyadic arithmetic)による線形演算
2. 整数専用SoftmaxのShiftmax
3. 整数専用GELUのShiftGELU
4. 整数専用LayerNormのI-LayerNorm

手法の概要

ダイアディック算術による線形演算

ダイアディック算術[2]は、整数のビットシフトを使用して浮動小数点演算のスケーリング係数を効率的に実現する手法です。これにより、線形演算を整数のみの算術で行うことができます。もともとはCNN(畳み込みニューラルネットワーク)用に設計されていましたが、ViTs(ビジョントランスフォーマー)の線形演算、特に埋め込み層のConv(畳み込み演算)やトランスフォーマー層のMatMul(行列乗算)やDense(密結合層)にも適用可能なようです。論文ではMatMul(行列乗算)を計算例として挙げています。
int8の$Q=(S_Q,I_Q)$と$K=(S_K,I_K)$のMatMulをした結果を$A'$とすると、$A'$は以下のように求められます。

$$A' = S_Q \cdot S_K \cdot (I_Q \times I_K^T)$$

このとき$S_{A'} = S_Q \cdot S_K$、$I_{A'} = I_Q \times I_K^T$とします。すると、出力$I_{A'}$はint32になってしまい、MatMulをするたびに計算に必要なbit数が増えてしまいます。そこでint32の$I_{A'}$をint8の$I_{A}$に変換します。

$$I_{A} = \lfloor \frac{S_{A'}\cdot I_{A'}}{S_A} \rceil = \lfloor \frac{S_{Q}\cdot S_{K}}{S_A} \cdot (I_Q \times I_K^T) \rceil$$

$S_A$は出力$A$のスケールファクターです。$\frac{S_{Q}\cdot S_{K}}{S_A}$は浮動小数点のままですが2進数に変換すれば整数値に変換できます。二進数変換関数を$DN$とすると

$$DN(\frac{S_{Q}\cdot S_{K}}{S_A}) = \frac{b}{2^c}$$

となります。$b$、$c$は正の整数。こうすることによって求めたかったint8の$I_A$は以下のように簡単に求められます。

$$I_{A} = (b\cdot(I_Q \times I_K^T)) \gg c$$

$\gg$はビットシフトを表しています。

整数専用SoftmaxのShiftmax

SoftmaxはViTsにおいてアテンションスコアを確率に変換しますが、非線形性のためダイアディック算術は使用できません。そのため、筆者はShiftmaxという近似手法を提案しています。Shiftmaxの導出までの流れを説明します。非常に長いです…。
最初にsoftmax関数を振り返ると、softmaxはよくオーバーフローを防ぐために変数の最大値を引く操作をしています。

$$Softmax(x_i) = \frac{e^{x_i - x_{max}}}{\Sigma_j e^{x_j - x_{max}}}$$

この式を$(S,I)$で示すとこのようになります。

$$Softmax(x_i) = \frac{e^{S_i\cdot(I_i - I_{max})}}{\Sigma_j e^{S_j\cdot(I_j - I_{max})}}$$

今後の式を簡単にするために$I_{\Delta_i}=I_i - I_{max} < 0$、$S_{\Delta_i}=S_i$を導入しておきます。

$$Softmax(x_i) = \frac{e^{S_{\Delta_i}\cdot I_{\Delta_i}}}{\Sigma_j e^{S_{\Delta_j}\cdot I_{\Delta_j}}}$$

このままでは当然整数型で計算できないので様々な近似をします。ここで筆者は$\log_2e\approx (1.0111)_2$と近似して計算を行いました。

$$e^{S_{\Delta_i}\cdot I_{\Delta_i}} = 2^{S_{\Delta_i}\cdot I_{\Delta_i}\cdot \log_2e} \approx 2^{S_{\Delta_i}\cdot I_{\Delta_i}\cdot (1.0111)_2}$$

このとき

$$(1.0111)_2 = (1)_2 + (1)_2\gg1 - (1)_2\gg4$$

であるため

$$e^{S_{\Delta_i}\cdot I_{\Delta_i}} \approx 2^{S_{\Delta_i}\cdot (I_{\Delta_i} + I_{\Delta_i}\gg1 - I_{\Delta_i}\gg4)}$$

となります。$I_{\Delta_i} + I_{\Delta_i}\gg1 - I_{\Delta_i}\gg4$は浮動小数点になってしまうため、もっと簡単にする必要があります。$S_{\Delta_i}(I_{\Delta_i} + I_{\Delta_i}\gg1 - I_{\Delta_i}\gg4)$の整数部分を$q$、特殊な定義になりますが小数部分を$S_{\Delta_i}r$とすると先ほどの式は以下のように変換できます。($I_{\Delta_i} < 0$であるため-がつく)

$$2^{S_{\Delta_i}\cdot (I_{\Delta_i} + I_{\Delta_i}\gg1 - I_{\Delta_i}\gg4)} = 2^{-q - S_{\Delta_i}r} = 2^{-S_{\Delta_i}r} \gg q$$

論文ではさらに$2^{-S_{\Delta_i}r} \approx S_{\Delta_i}[(-r) \gg 1 + \frac{1}{S_{\Delta_i}}]$と近似しています。ここまで来てやっと

$$e^{S_{\Delta_i}\cdot I_{\Delta_i}} \approx S_{\Delta_i}I_{exp}$$

が得られてsoftmaxが計算できます。ただし$I_{exp} = [(-r) \gg 1 + \frac{1}{S_{\Delta_i}}] \gg q$です。
非常に長かったですが論文ではもう少し処理をしているようです。詳しくは論文を見てください…。

整数専用GELUのShiftGELU

ShiftGELUは整数型で計算可能にしたGELUです。GELUは以下のように表されています。

$$GELU(x) = x\cdot \frac{1}{\sqrt{2\pi}}\int_{-\infin}^{\infin}e^{-t^2/2}dt \approx x\cdot \sigma(1.702x)$$

ここで$\sigma$はsigmoid関数です。このときsimgoidはsoftmax関数の形に変換可能なためsfiftmaxと同様の計算を行うことができます。

整数専用LayerNormのI-LayerNorm

I-LayerNormは、Vision Transformers(ViTs)での入力の正規化を行います。具体的には、入力の平均値と分散を使用して、特徴次元を正規化します​​。しかし、整数演算では標準偏差の平方根計算をサポートしていません​。I-LayerNormは、ビットシフトを用いた軽量な整数反復アプローチ[3]により改善されています。具体的には、

$$I_{i+1} = (I_i + ⌊Var(Ix)/I_i⌋)/2 = (I_i + ⌊Var(x)/I_i⌋) ≫ 1$$

という式を用いています​​。初期値$I_0$は2⌊bit(Var(Ix))/2⌋で計算されます。10回の反復でほとんどの収束を達成できることが実験的に見出され、ハードウェア実装を容易にするために反復数に基づく停止基準に変更されています​​。

実験結果

この論文では学習時には浮動小数点で学習して、その学習パラメータや計算過程を提案手法によって整数型に置き換えています。
I-ViTは、浮動小数点(FP)モデルに比べて、3.72～4.11倍の加速を実現しています。また、わずかに高い精度を達成しています​​。
DeiT-TとDeiT-Sモデルの量子化において、FasterTransformerは3.45倍と3.53倍の加速しか達成できませんでした。これは、浮動小数点算術の非効率さに起因しています​​。

I-ViTは、INT8精度で評価され、Top-1精度は81.74%(I-BERTと比較して-0.11%)、遅延は7.93ms(3.88倍の加速)でした。これに対し、FP32ベースラインのTop-1精度は81.35%、遅延は16.8msでした​​。
Swin-Sモデルにおいても、I-ViTは83.01%のTop-1精度(FP32ベースラインに対して-0.19%)と6.92msの遅延(4.02倍の加速)を実現しました​​。

まとめと感想

ディープラーニングでは浮動小数で計算を行うことが当たり前の中、すべての処理を整数型で計算できる手法を提案しています。非常に面白いが部分的に整数型にしているだけのFasterTransformerとあまり大差ない推論スピードなのはなぜなのでしょうか。憶測ですが、今のGPUは浮動小数点の計算に特化しているからかなとも思いました。整数型の処理が早いハードウェアとかがあればもっと有意差が出て面白いかもしれません。

6. 論文のリクエスト

解説してほしい論文のリクエストを受け付けています！
リクエストから2週間程度で記事を作成したします。
どなたでもお気軽にリクエストしてください！

私はAI、コンピュータビジョンを専門としています。それ以外の分野の論文は投稿に時間が掛るか、間違った情報を流さないためにお断りするケースがございます。

https://forms.gle/kDRYWM7K9k1pYvZ16

参考文献

[1]

NVIDIA. FasterTransformer, https://github.com/nvidia/fastertransformer.git, 2022.

[2]

Jacob, Benoit, et al. "Quantization and training of neural networks for efficient integer-arithmetic-only inference." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[3]

Richard Crandall and Carl Pomerance. Prime numbers.
Springer, 2001.

脚注

1. 実装では$m$は1つではなく、最小値と最大値は違う値になっている。また、学習や推論の度にモーメンタムなどで更新している。