[論文レビュー] Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations

Tri Dao, Albert Gu+3 more|arXiv (Cornell University)|Mar 14, 2019

Tensor decomposition and applications参考文献 48被引用数 31

ひとこと要約

本稿では、DFT、DCT、ハダマード変換、畳み込みなどの構造的線形変換—FFT、DCT、畳み込み—を、スパースで構造的なブロック対角行列の積として表現することにより、$O(N\log N)$の高速なアルゴリズムを自動的に学習する微分可能でパラメータ化可能なバタフライ因子分解を提案する。この手法は、エンドツーエンドの機械学習モデルで最先端の性能を達成し、CIFAR-10で全結合層を上回る3.9ポイントの精度向上を達成しながら、パラメータ数を40倍削減し、推論速度を4倍高速化する。

ABSTRACT

Fast linear transforms are ubiquitous in machine learning, including the discrete Fourier transform, discrete cosine transform, and other structured transformations such as convolutions. All of these transforms can be represented by dense matrix-vector multiplication, yet each has a specialized and highly efficient (subquadratic) algorithm. We ask to what extent hand-crafting these algorithms and implementations is necessary, what structural priors they encode, and how much knowledge is required to automatically learn a fast algorithm for a provided structured transform. Motivated by a characterization of fast matrix-vector multiplication as products of sparse matrices, we introduce a parameterization of divide-and-conquer methods that is capable of representing a large class of transforms. This generic formulation can automatically learn an efficient algorithm for many important transforms; for example, it recovers the $O(N \log N)$ Cooley-Tukey FFT algorithm to machine precision, for dimensions $N$ up to $1024$. Furthermore, our method can be incorporated as a lightweight replacement of generic matrices in machine learning pipelines to learn efficient and compressible transformations. On a standard task of compressing a single hidden-layer network, our method exceeds the classification accuracy of unconstrained matrices on CIFAR-10 by 3.9 points -- the first time a structured approach has done so -- with 4X faster inference speed and 40X fewer parameters.

研究の動機と目的

• 手作業による設計なしに、構造的線形変換のための効率的で、四乗より小さい時間計算量のアルゴリズムを自動的に学習すること。
• FFTやcuFFTのようなプラットフォーム依存の手作業最適化ライブラリに依存することを減らし、エンドツーエンドで微分可能な学習によって高速変換を可能にすること。
• 構造的なパラメータ化が、実世界の機械学習ベンチマークで非制約の全結合層を上回る性能を発揮できるかどうかを検証すること。
• スパース行列因子分解を通じて高速アルゴリズムを学習するために必要な最小限のインダクティブバイアスを同定すること。
• 深層学習パイプラインに学習可能な高速変換層を統合し、効率性と性能の両方を向上させること。

提案手法

• 線形変換を$O(\log N)$個のバタフライ行列の積としてパラメータ化する—構造的でスパースなブロック対角行列であり、合計で$O(N)$のパラメータを持つ。
• 目的変換の入出力ペアから、バタフライパラメータを微分可能最適化フレームワークで学習する。
• バタフライ因子分解の再帰的で分割統治型の構造を活用し、$O(N\log N)$の行列-ベクトル乗算を実現する。
• 各ステージが置換と、$2^k$個のサイズ$N/2^k$のブロック対角行列を適用する$O\left(\log N\right)$段階の合成として変換を表現する。
• 全結合層の軽量で圧縮可能な代替手段として、ニューラルネットワークにバタフライ層を統合する。
• 5行のPythonコードで高速乗算アルゴリズムを実装し、変換固有のチューニングなしに、あらゆるハードウェアで効率的な推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 汎用的で微分可能なパラメータ化が、DFT や DCT などの多様な構造的変換に対して、$O(N\log N)$の高速アルゴリズムを学習できるか？
• RQ2 構造的で学習可能な行列因子分解は、精度、パラメータ効率、推論速度の観点で、非制約の全結合層をどの程度上回れるか？
• RQ3 同じバタフライパラメータ化により、Cooley-Tukey FFT などの既知の高速アルゴリズムが、$N=1024$までの実用的サイズで機械精度まで再現できるか？
• RQ4 高度に最適化された手作業チューニング済みカーネル（例：cuFFT や cuDNN）と比較して、学習されたバタフライ変換の性能はいかが？
• RQ5 バタフライ層は、パラメータを大幅に増加させることなく、ResNet18 などの標準的な深層学習アーキテクチャで性能向上をもたらせるか？

主な発見

• 本手法は、$N$が1024までにまで、$O(N\log N)$のCooley-Tukey FFTアルゴリズムを機械精度まで再現し、スパースおよび低ランクベースラインを上回る。
• CIFAR-10では、バタフライパラメータ化を用いた1層ネットワークが93.89%の精度を達成し、非制約の全結合層を3.9ポイント上回るが、パラメータ数は40倍少ない。
• バタフライ層を用いることで、ResNet18のCIFAR-10における精度が0.43ポイント向上し、モデルパラメータの増加はわずか0.07%にとどまる。
• 学習されたバタフライ行列の推論は、密度行列のGEMVより最大2桁の速度向上を達成し、CPU上では専用のFFT/DCTカーネルと3〜5倍の速度差に収まる。
• GPU上で学習する際、バタフライ層は密度行列のGEMMより15%高速であり、FFTに比べて40%以内の性能を示し、競争力のある学習効率を実現する。
• 本手法は、畳み込みやその他の構造的変換を学習する際に一般化可能であり、高速乗算アルゴリズムに変換固有の最適化が不要である。