Multiplying Matrices Without Multiplying

CompML勉強会での論文"Multiplying Matrices Without Multiplying"の紹介資料．

Transcript

1. CompML Multiplying Matrices Without Multiplying - 論⽂紹介 - Ryohei Izawa
2. CompML TL; DR • 計算時に乗算を⾏わない，近似⾏列乗算(Approximate Matrix Multiplication; AMM)の ⼿法であるMADDNESS※を提案． • 直積量⼦化(Product Quantization; PQ)と，その近似探索のハッシュ関数に学習した2 分⽊ハッシュ関数を⽤いたLocality Senstive Hashing(LSH)によって計算を⾼速化． • 素直な⾏列乗算よりも100倍⾼速，最新の近似⼿法よりも10倍⾼速． 1 ※ Multiply-ADDitioN-lESS
3. CompML AMMの概要
4. CompML ⾏列の積と計算量 ⾏列𝐴 ∈ ℝ!×#と⾏列𝐵 ∈ ℝ#×$の乗算（内積）を素朴に計算すると，𝑁×𝑀×𝐷回の乗算が必要 となる． => 𝑁 = 𝐷 = 𝑀のとき，計算量は𝑂(𝑁%) 3 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& = ' '(! $ 𝑎!' 𝑏'! ' '(! $ 𝑎!' 𝑏'" ' '(! $ 𝑎!' 𝑏'# ⋯ ⋯ ' '(! $ 𝑎!' 𝑏'& ' '(! $ 𝑎"' 𝑏'! ⋱ ' '(! $ 𝑎#' 𝑏'! ⋱ ⋮ ⋱ ⋮ ' '(! $ 𝑎%' 𝑏'! ⋯ ⋯ ⋯ ⋯ ' '(! $ 𝑎%' 𝑏'&
5. CompML 近似⾏列乗算(AMM) ⾏列の乗算を近似的に求めるため，⾏列の圧縮やハッシュ関数を⽤いて計算コストを削減す る⽅法． • Linear Approximation ⾏列𝐴，𝐵を低次元空間に圧縮した上で乗算を⾏う⼿法． (e.g. 𝐴の列と𝐵の⾏をランダムにサンプリングして圧縮) • Hashing to Avoid Linear Operations Lookup tableを事前に⽤意しておき，近似探索により計算を⾏う． ⇨ MADDNESSはこちらの⼿法 4 𝐴𝐵 ≈ 𝐴𝑉! 𝑉" #𝐵 , 𝑤ℎ𝑒𝑟𝑒 𝑉! , 𝑉" ∈ ℝ$×&, 𝑑 ≪ 𝐷
6. CompML MADDNESS
7. CompML 問題の定式化 ⾏列𝐴 ∈ ℝ!×#，⾏列𝐵 ∈ ℝ#×$が与えられ，𝑁 ≫ 𝐷 ≥ 𝑀のとき，定数𝛼, 𝛽のもとで誤差𝜀(𝜏)が ⼩さくなるように以下の𝑔 6 , ℎ 6 , 𝑓(6)を求める． 6 𝛼𝑓(𝑔 𝐴 , ℎ 𝐵 ) + 𝛽 − 𝐴𝐵 ' < 𝜀(𝜏) 𝐴𝐵 ' 𝑔 𝐴 : ⼊⼒⾏列𝐴の各⾏の部分空間(サブベクトル)に近似するプロトタイプのインデックスの集まり． ℎ(𝐵): 既知の⾏列𝐵の各列の部分空間とそれに対応するプロトタイプごとの内積が格納されたテーブル． 𝑓(𝑔 𝐴 , ℎ 𝐴 ): 部分空間ごとの内積の和．⼩さいbit幅で量⼦化しているため，定数𝛼, 𝛽で元のbit幅の値に戻す．
8. CompML MADDNESSの⼿法概要 直積量⼦化(Product Quantization; PQ)を⽤いたLookup Tableの作成と，Locality Sensitive Hashing(LSH)によるLookup Tableの⾼速検索が肝． LSHのハッシュ関数としては2分⽊学習モデルが使われる． 7 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ ∘ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& C個の部分空間ごとにベクトル量⼦化し 学習によりK個のプロトタイプを作成
9. CompML 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ ∘ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& MADDNESSの⼿法概要 直積量⼦化(Product Quantization; PQ)を⽤いたLookup Tableの作成と，Locality Sensitive Hashing(LSH)によるLookup Tableの⾼速検索が肝． LSHのハッシュ関数としては2分⽊学習モデルが使われる． 8 プロトタイプごとに内積を予め計算 C個の部分空間ごとにベクトル量⼦化し 学習によりK個のプロトタイプを作成
10. CompML 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ ∘ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& MADDNESSの⼿法概要 直積量⼦化(Product Quantization; PQ)を⽤いたLookup Tableの作成と，Locality Sensitive Hashing(LSH)によるLookup Tableの⾼速検索が肝． LSHのハッシュ関数としては2分⽊学習モデルが使われる． 9 プロトタイプごとに内積を予め計算 C個の部分空間ごとにベクトル量⼦化し 学習によりK個のプロトタイプを作成 ⼊⼒ベクトルaの各部分空間 に対応するプロトタイプの インデックスをLSHで検索
11. CompML 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ ∘ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& MADDNESSの⼿法概要 直積量⼦化(Product Quantization; PQ)を⽤いたLookup Tableの作成と，Locality Sensitive Hashing(LSH)によるLookup Tableの⾼速検索が肝． LSHのハッシュ関数としては2分⽊学習モデルが使われる． 10 プロトタイプごとに内積を予め計算 C個の部分空間ごとにベクトル量⼦化し 学習によりK個のプロトタイプを作成 ⼊⼒ベクトルaの各部分空間 に対応するプロトタイプの インデックスをLSHで検索 部分空間ごとの内積を合計し出⼒
12. CompML MADDNESSの⼿法概要 PQで得られた𝑓(𝑔(𝐴), ℎ(𝐵))は，元の⾏列𝐴，𝐵よりも少ないbit幅で表現しているため，定数 𝛼，𝛽を⽤いて元のbit幅に戻す． 𝛼𝑓 𝑔 𝐴 , ℎ 𝐵 + 𝛽 11 ただし， • 𝑔(𝐴)は⼊⼒⾏列𝐴の各⾏の部分空間(サブベクトル)に近似するプロトタイプのインデックスの集まり • ℎ(𝐵)は既知の⾏列𝐵の各列の部分空間とそれに対応するプロトタイプごとの内積が格納されたテーブル を返す．
13. CompML Product Quantization 直積量⼦化(Product Quantization; PQ)の⼿順; 12 1. プロトタイプの学習 学習⽤の⾏列 : 𝐴の⾏ベクトルを𝐶個の部分ベクトルに 分割し， 部分ベクトルごとに𝐾個のプロトタイプを 学習により獲得 2. ⼊⼒ベクトル𝑎のエンコード：𝑔(𝑎) ⼊⼒ベクトル𝑎の部分空間ごとに最も近いプロトタイ プを選択 3. 内積Lookup Tableの構築：ℎ(𝑏) 学習で獲得したプロトタイプと既知の⾏列𝐵の列ベク トル𝑏の部分空間ごとの内積を予め計算 4. アグリゲーション ℎ(𝑏)から𝑔 𝑎 で得たインデックスで部分空間ごとの 内積を取得し，⾜し合わせて出⼒とする
14. CompML 1. プロトタイプの学習 学習データの⾏列 : 𝐴の各⾏を𝐶個の部分空間に分割し，部分空間ごとに𝐾個のプロトタイプベ クトルをK-means学習※によって獲得する． ※ MADDNESSでは，素朴なK-meansではなく，2分⽊学習モデルを⽤いる． 13 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ C個に分割 K個のプロトタイプ … C個の部分空間 𝑃( ( 𝑃) ( 𝑃* + 𝑃( ) 𝑃, ( 𝑃( * 𝑃( - 𝑃, - … … … … 𝑃) ) 𝑃* ) 𝑃) * 𝑃* * 𝑃, ) 𝑃, * 𝑃. /: 部分空間𝑐の𝑘番⽬のプロトタイプベクトル（⻑さ|𝒥(/)|) 𝑠𝑢𝑏𝑣𝑒𝑐𝑡𝑜𝑟 ! 𝑠𝑢𝑏𝑣𝑒𝑐𝑡𝑜𝑟 " 𝑠𝑢𝑏𝑣𝑒𝑐𝑡𝑜𝑟 # 𝑠𝑢𝑏𝑣𝑒𝑐𝑡𝑜𝑟 $
15. CompML K-meansによるベクトル量⼦化 1 𝐴 ∈ ℝ!×# : 学習⾏列 K: 部分空間ごとのプロトタイプ数 C: 部分空間の数 {𝒥(%)}%'( ) ∶ 各部分空間の要素のインデックス𝑗の集合 𝑃(%) ∈ ℝ*× 𝒥 ! ∶ 部分空間𝑐の学習で得られるプロトタイプ 𝑧(%): 部分空間𝑐において選ばれたプロトタイプのインデックス 14 0 1 1 0 K-meansによりベクトルを2bitに量⼦化する例 以下の最⼩化するプロトタイプの集合をK-meansの学習により獲得する． < ,'( ! < %'( ) < -∈𝒥(!) 1 𝐴,- − 𝑃 /(!),- (%) 1
16. CompML 2. ⼊⼒ベクトル𝑎のエンコード：𝑔(𝑎) ⼊⼒ベクトル𝑎に対し，各部分空間に対応するプロトタイプのインデックスの集合を返す． 各インデックスは，𝑙𝑜𝑔> 𝐾 𝑏𝑖𝑡でエンコードされる（𝐾 = 16のとき4𝑏𝑖𝑡） 15 𝑔 / 𝑎 ≜ 𝑎𝑟𝑔𝑚𝑖𝑛. G 2∈𝒥(*) 𝑎2 − 𝑃.,2 (/) ) ベクトル𝑎の部分空間𝑐に対応するインデックスの取得
17. CompML 3. 内積Lookup Tableの構築：ℎ(𝑎) 16 1の学習で得られたプロトタイプごとに，𝐵の列ベクトル𝑏との内積を計算し，Loopup Tableに格納する． Loopup Table の値は8bitで量⼦化する． ℎ / 𝑏 . ≜ G 2∈𝒥(*) 𝑏2 𝑃 .,2 (/) ベクトル𝑏の部分空間𝑐に対応するサブベクトルと 各プロトタイプ𝑃 . (/)との内積
18. CompML 4.アグリゲーション 17 𝑓(𝑔 𝑎 , ℎ 𝑏 ) ：𝑔(𝑎)で得られたインデックスの集合の各インデックスを⽤いて，ℎ(𝑏)から対 応する内積の値を取得し，⾜し合わせる． 𝑎#𝑏 = G /6( - 𝑎 / #𝑏(/) ≈ G /6( - ℎ / (𝑏). , 𝑘 = 𝑔 / 𝑎 ベクトル𝑎とベクトル𝑏の内積に近似
19. CompML Locality Sensitive Hashing ベクトル𝑎と最も近いプロトタイプの探索を⾏う際，ユークリッド距離を計算すると部分空間 ごとに計算量は𝑂(𝐾)となる． ⇨ LSHにより類似探索を⾏う． 18 … LSHでは，ベクトルをハッシュすると，類似する サブベクトルは同じBucketにハッシュされる プロトタイプは各Bucketにハッシュされた サブベクトルの平均ベクトル 𝐾 部分空間cのサブベクトル
20. CompML 2分⽊ベースのハッシュ関数 プロトタイプ数𝐾 = 16のとき，16個のBucketにハッシュされるように，深さ4の2分⽊をハ ッシュ関数とする．レベル𝑡の各節𝑖 ∈ 2?@Aには， 𝑎(B) の値𝑎 C7 B で分割する閾値𝑣D ?が設定され る．各レベルどのインデックス𝑗の値で分割するかと閾値𝑣D ?は学習により決まる． 19 1 5 2 7 3 9 4 11 6 13 8 14 10 15 12 16 1 5 2 7 3 4 6 8 1 2 3 4 1 2 1 𝑎 2, / ≥ 𝑣( ) 𝑡 = 1, 𝑖 = {1} 𝑡 = 2, 𝑖 = {1,2} 𝑡 = 3, 𝑖 = {1,2,3, 4} 𝑡 = 4, 𝑖 = {1. . 8} 𝑡 = 5, 𝑖 = {1. . 16} ⼊⼒𝑎(/) = 𝑎( / 𝑎) / … 𝑎2 / … 𝑎 |𝒥 * | (/) 𝑎 2- / < 𝑣( ( 𝑘 = 5にハッシュされるケース 𝑎 2. / < 𝑣) * 𝑎 2/ / < 𝑣* 9
21. CompML ハッシュ関数の学習 各レベルで分割する際の値のインデックス𝑗と閾値𝑣D ?を学習により求める． 1. レベル𝑡のノード𝑖にマッピングされたベクトル𝑥を格納するためBucket ℬD ?を導⼊． 2. レベルごとに損失関数ℒ(𝑗, ℬD ?)を最⼩化するようにℬD ?にベクトル𝑥のインデックスを格納し て，ベクトル𝑥C の閾値𝑣D ?を求める． 3. レベルtの損失関数ℒ(𝑗, ℬD ?)の合計ℒ(𝑗, ℬ)を最⼩化するようなインデックス𝑗を，最終的なレ ベルtの分割で使⽤する𝑥C の値のインデックスとする． 20 1 2 3 4 1 2 1 ⼊⼒𝑎(/) … … … … ℬ( ( ℬ( ) ℬ) ) ℬ( * ℬ) * ℬ) * ℬ9 * ℒ 𝑗, ℬ ≜ G :∈ℬ 𝑥2 − 1 ℬ G :0 1∈ℬ 𝑥2 < )
22. CompML ハッシュ関数の学習アルゴリズム xx 21
23. CompML プロトタイプの最適化 再構成した⾏列 : 𝐴の2乗誤差をさらに⼩さくするため，リッジ回帰によりプロトタイプ𝑃を最 適化する． 22 1 𝐴 ∈ ℝ!×# : 学習⾏列 𝐾: 部分空間ごとのプロトタイプ数 𝐶: 部分空間の数 𝑃 ∈ ℝ*)×# ∶ 各部分空間𝑐における𝐾個の学習されたプロトタイプで構成される⾏列 𝐺 ∈ ℝ!×*) : 1 𝐴 ≈ 𝐺𝑃となるように部分空間ごとにプロトタイプを選択するone-hotベクトルが格納された⾏列 𝑃 ≜ (𝐺I𝐺 + 𝜆𝐼)@A𝐺I : 𝐴
24. CompML 𝑓(𝑔(𝐴), ℎ(𝐵))の加算における課題 𝑓(𝑔(𝐴), ℎ(𝐵))の加算において正確な加算を⾏うためには，8bitで値が格納されているℎ(𝐵)か ら検索された各エントリを16bitにアップキャストしてから加算命令を⾏う必要がある． ⇨ オーバーヘッドの発⽣やスループットの半減 23 𝑓(𝑔 𝐴 , ℎ 𝐵 )J,K ≜ Q BLA M 𝑇K,B,N , 𝑘 = 𝑔 B (𝑎J ) 𝑇 ∈ ℝ2×)×* : ⾏列𝐵の全ての列の内積Lookup Table 𝑓(𝑔(𝐴), ℎ(𝐵))の加算
25. CompML ⾼速な8bit加算⼿法 加算命令ではなく，平均化命令を使⽤することで，若⼲の精度損失と引き換えに計算速度を 向上させる． 24 Pairwise Integer Sum and Sum Estimator 𝑠 𝑎, 𝑏 ≜ 𝑎 + 𝑏 ̂ 𝑠 𝑎, 𝑏 ≜ 2𝜇 𝑎, 𝑏 𝑤ℎ𝑒𝑟𝑒 𝜇 𝑎, 𝑏 ≜ Z [ 1 2 (𝑎 + 𝑏 + 1)
26. CompML 計算量 𝑔(𝐴)の計算量はΘ(𝑁𝐶)，ℎ(𝐵)作成時の計算量はΘ(𝑀𝐾𝐶𝐷)，𝑓(𝑔(𝐴), ℎ(𝐵))の計算量はΘ(𝑁𝐶𝑀)． 全体の計算量はΘ(𝑀𝐶(𝐾𝐷 + 𝑁))だが，𝐾 = 16かつ𝑁 ≫ 𝐷よりΘ(𝑀𝐶𝑁)となる． 25 𝑎!! 𝑎!" 𝑎!# ⋯ ⋯ 𝑎!$ 𝑎"! ⋱ 𝑎#! ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎%! ⋯ ⋯ ⋯ ⋯ 𝑎%$ 𝑏!! 𝑏!" 𝑏!# ⋯ ⋯ 𝑏!& 𝑏"! ⋱ 𝑏#! ⋱ ⋮ ⋱ ⋮ 𝑏$! ⋯ ⋯ ⋯ ⋯ 𝑏$& ⾏列𝐴 ⾏列𝐵 𝑁 𝐷 (𝐶個に分割) 𝐷 𝑀 𝑔(𝐴): 𝑁𝐶個のプロトタイプのインデックス を取得するためΘ(𝑁𝐶) ℎ(𝐵): ⻑さ𝐷の𝐾𝐶個のプロトタイプごとの 内積を列ごとに計算するためΘ(𝑀𝑁𝐶𝐷) 𝑓(𝑔 𝐴 ,ℎ 𝐵 ): 各列𝐶個の部分空間プロトタ イプの内積をN⾏M列分計算するため Θ(𝑁𝐶𝑀) 𝑎$$ 𝑎$% 𝑎$& ⋯ ⋯ 𝑎$' 𝑎%$ ⋱ 𝑎&$ ⋱ ⋮ ⋱ ⋮ ⋱ 𝑎($ ⋯ ⋯ ⋯ ⋯ 𝑎(' ∘ 𝑏$$ 𝑏$% 𝑏$& ⋯ ⋯ 𝑏$) 𝑏%$ ⋱ 𝑏&$ ⋱ ⋮ ⋱ ⋮ 𝑏'$ ⋯ ⋯ ⋯ ⋯ 𝑏') ⾏列A𝐵
27. CompML 近似誤差の⼀般化保証 近似誤差に以下の上界を与える． 26 𝔼𝒟 ℒ(𝑎, 𝑏) ≤ 𝔼> ! ℒ 𝑎, 𝑏 + 𝐶?2 𝑏 ) 2 𝜆 1 256 + 8 + 𝜈(𝐶, 𝐷, 𝛿) 2𝑛 𝑤ℎ𝑒𝑟𝑒 ℒ 𝑎, 𝑏 ≜ 𝑎#𝑏 − 𝛼𝑓 𝑔 𝑎 , ℎ 𝑏 − 𝛽 , 𝛼 𝑖𝑠 𝑡ℎ𝑒 𝑠𝑐𝑎𝑙𝑒 𝑢𝑠𝑒𝑑 𝑓𝑜𝑟 𝑞𝑢𝑎𝑛𝑡𝑖𝑧𝑖𝑛𝑔 𝑡ℎ𝑒 𝑙𝑜𝑜𝑘𝑢𝑝 𝑡𝑎𝑏𝑙𝑒𝑠, 𝛽 𝑖𝑠 𝑡ℎ𝑒 𝑐𝑜𝑛𝑠𝑡𝑎𝑛𝑡𝑠 𝑢𝑠𝑒𝑑 𝑖𝑛 𝑞𝑢𝑎𝑛𝑡𝑖𝑧𝑖𝑛𝑔 𝑡ℎ𝑒 𝑙𝑜𝑜𝑘𝑢𝑝 𝑡𝑎𝑏𝑙𝑒𝑠 𝑝𝑙𝑢𝑠 𝑡ℎ𝑒 𝑑𝑒𝑏𝑖𝑎𝑠𝑖𝑛𝑔 𝑐𝑜𝑛𝑠𝑡𝑎𝑛𝑡 𝑓𝑜𝑟 𝑓𝑎𝑠𝑡 8𝑏𝑖𝑡 𝑎𝑔𝑔𝑟𝑒𝑔𝑎𝑡𝑖𝑜𝑛, 𝑎𝑛𝑑 𝜈 𝐶, 𝐷, 𝛿 ≜ 𝐶 4 log) 𝐷 + 256 log 2 − log 𝛿 𝐷: ℝ#に対する確率分布 1 𝐴 ∈ ℝ!×#: 各⾏がDから独⽴に得られた学習⾏列． 𝜎9 で最⼤特異値がバウンドされている． 𝐶: 部分空間の数 𝜆: 正則化パラメータ 𝛿: 0 < 𝛿 < 1の値
28. CompML 実験結果 𝑔(𝑎)によるエンコード時間は，その他⼿法と⽐較して⾼速． CIFAR-10，CIFAR-100での精度と推論時間のトレードオフでも他⼿法を上回る． 27
29. CompML まとめ（再掲） 28 • 計算時に乗算を⾏わない，近似⾏列乗算(Approximate Matrix Multiplication; AMM)の ⼿法であるMADDNESS※を提案． • 直積量⼦化(Product Quantization; PQ)と，その近似探索のハッシュ関数に学習した2 分⽊ハッシュ関数を⽤いたLocality Senstive Hashing(LSH)によって計算を⾼速化． • 素直な⾏列乗算よりも100倍⾼速，最新の近似⼿法よりも10倍⾼速． ※ Multiply-ADDitioN-lESS
30. CompML Reference
31. CompML 30 [1] Blalock, D., & Guttag, J. (2021). Multiplying Matrices Without Multiplying. arXiv preprint arXiv:2106.10860. [2] Manne, S., & Pal, M. (2014). Fast approximate matrix multiplication by solving linear systems. arXiv preprint arXiv:1408.4230.