[論文レビュー] Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps

Tri Dao, Nimit S. Sohoni+6 more|arXiv (Cornell University)|Dec 29, 2020

Neural Networks and Applications参考文献 79被引用数 9

ひとこと要約

K-matrices（K行列）は、構造的線形写像を学習可能でパラメータ効率の良い表現として提供するもので、バタフライ行列の積として構成され、低ランク、スパース、フーリエ変換などの多様な構造を近似的に最適な空間的・時間的複雑性で捉える。エンドツーエンドの微分可能学習を可能にし、ImageNetでは最大5%の精度向上、Transformerでは36%の推論高速化を実現する。

ABSTRACT

Modern neural network architectures use structured linear transformations, such as low-rank matrices, sparse matrices, permutations, and the Fourier transform, to improve inference speed and reduce memory usage compared to general linear maps. However, choosing which of the myriad structured transformations to use (and its associated parameterization) is a laborious task that requires trading off speed, space, and accuracy. We consider a different approach: we introduce a family of matrices called kaleidoscope matrices (K-matrices) that provably capture any structured matrix with near-optimal space (parameter) and time (arithmetic operation) complexity. We empirically validate that K-matrices can be automatically learned within end-to-end pipelines to replace hand-crafted procedures, in order to improve model quality. For example, replacing channel shuffles in ShuffleNet improves classification accuracy on ImageNet by up to 5%. K-matrices can also simplify hand-engineered pipelines -- we replace filter bank feature computation in speech data preprocessing with a learnable kaleidoscope layer, resulting in only 0.4% loss in accuracy on the TIMIT speech recognition task. In addition, K-matrices can capture latent structure in models: for a challenging permuted image classification task, a K-matrix based representation of permutations is able to learn the right latent structure and improves accuracy of a downstream convolutional model by over 9%. We provide a practically efficient implementation of our approach, and use K-matrices in a Transformer network to attain 36% faster end-to-end inference speed on a language translation task.

研究の動機と目的

• 深層学習における効率性と精度を高めるために、手動で選択・設計された構造的線形写像（例：低ランク、スパース、フーリエ変換）の課題に対処すること。
• 既存の構造的行列クラスの限界（微分不能、非効率、多様な構造を十分に表現できない）を克服すること。
• すべての構造的線形写像を近似的に最適なパラメータ数と算術的複雑性で一元的に表現できる、微分可能で効率的な統合的表現を開発すること。
• ニューラルネットワーク内での構造的コンポonentのエンドツーエンド学習を可能にし、チャネルシャッフルやフィルターバンクといった手作業による手続きを置き換えること。
• K-matricesの実用的有用性を、画像分類、音声認識、Transformer推論といった実世界のタスクで示すこと。

提案手法

• バタフライ行列の積として構成される、K-matrices（K行列）を提案。バタフライ行列は、行列-ベクトル乗算が効率的であることで知られている。
• 任意の構造的線形写像が $ s \ll n^2 $ の演算で実現可能な場合、その算術回路複雑性に一致する総パラメータ数を持つスパース行列の積として表現可能であるという理論的基盤を活用する。
• スパース因子分解におけるスパースパターン探索の非微分性を回避するため、事前に定義された固定されたスパース構造を持つバタフライ行列を採用する。
• K行列全体のパラメータ化を微分可能にすることで、バックプロパゲーションによるエンドツーエンドの最適化を可能にする。
• バタフライ構造を活用した効率的な行列-ベクトル乗算アルゴリズムを実装し、$ O(n \log n) $ の複雑性を達成する。
• Transformerを含むニューラルネットワークにK-matricesを学習可能なレイヤーとして統合し、チャネルシャッフルやFFTレイヤーといった固定された構造的コンポーネントを置き換える。

実験結果

リサーチクエスチョン

• RQ1 一様で微分可能なパラメータ化によって、近似的に最適な効率性を達成する形で、すべての主要な構造的線形写像クラスを統合的に表現できるか？
• RQ2 K-matricesがエンドツーエンドの深層学習パイプライン内で有効に学習可能であり、手作業で設計された構造的コンポーネントと比較してモデル性能を向上させられるか？
• RQ3 音声認識における従来のフィルターバンク計算といった従来の前処理手順を、K-matricesで置き換えることで、精度損失を最小限に抑えられるか？
• RQ4 K-matricesは、置換された画像分類といった複雑なタスクにおいて、潜在的な構造を発見し、下流のモデル精度を向上させられるか？
• RQ5 実世界のモデル（例：Transformer）において、標準的な線形レイヤーをK-matricesに置き換えた場合、推論速度とメモリ効率にどの程度の向上が得られるか？

主な発見

• K-matricesは、行列の算術回路複雑性の理論的下界と一致する近似的に最適なパラメータ数と算術的複雑性で、任意の構造的線形写像を表現可能である。
• ShuffleNetのチャネルシャッフルを学習可能なK-matrixレイヤーに置き換えることで、推論コストを増加させずにImageNetのトップ-1精度が最大5%向上した。
• TIMIT音声認識タスクにおいて、手作業で設計されたフィルターバンクを学習可能なK-matrixレイヤーに置き換えたところ、精度低下はわずか0.4%にとどまり、パイプラインが簡素化された。
• 挑戦的な置換された画像分類タスクにおいて、置換のK-matrixベースの表現により、下流の畳み込みモデルの精度が9%以上向上した。
• TransformerモデルにK-matricesを統合したところ、言語翻訳タスクでエンドツーエンドの推論速度が36%高速化され、精度に顕著な損失は生じなかった。
• K-matrix表現により、ヒューリスティック的または高コストなスパースパターン探索を伴わない、効率的で微分可能な構造的コンポーネントの学習が可能になった。