[論文レビュー] Hardware Acceleration of Explainable Machine Learning using Tensor Processing Units

Zhixin Pan, Prabhat Mishra|arXiv (Cornell University)|Mar 22, 2021

Advanced Neural Network Applications参考文献 10被引用数 4

ひとこと要約

本論文は、行列畳み込みとフーリエ変換を活用して複雑な解釈タスクを効率的な行列演算に変換することで、解釈可能機械学習における最初のTPUベースのハードウェア加速フレームワークを提案する。GPUベースの手法と比較して分類時間で最大38.6倍、解釈時間で12.8倍の高速化を達成し、多様なMLワークロードにおいてリアルタイムの結果解釈を可能にする。

ABSTRACT

Machine learning (ML) is successful in achieving human-level performance in various fields. However, it lacks the ability to explain an outcome due to its black-box nature. While existing explainable ML is promising, almost all of these methods focus on formatting interpretability as an optimization problem. Such a mapping leads to numerous iterations of time-consuming complex computations, which limits their applicability in real-time applications. In this paper, we propose a novel framework for accelerating explainable ML using Tensor Processing Units (TPUs). The proposed framework exploits the synergy between matrix convolution and Fourier transform, and takes full advantage of TPU's natural ability in accelerating matrix computations. Specifically, this paper makes three important contributions. (1) To the best of our knowledge, our proposed work is the first attempt in enabling hardware acceleration of explainable ML using TPUs. (2) Our proposed approach is applicable across a wide variety of ML algorithms, and effective utilization of TPU-based acceleration can lead to real-time outcome interpretation. (3) Extensive experimental results demonstrate that our proposed approach can provide an order-of-magnitude speedup in both classification time (25x on average) and interpretation time (13x on average) compared to state-of-the-art techniques.

研究の動機と目的

• リアルタイム展開を制限する解釈可能機械学習（XML）における遅い解釈処理という重要なボトル neck を解消すること。
• 反復的で時間がかかる最適化問題に依存する既存のXML手法の非効率性を克服すること。
• TPUが行列計算を高速化する特性を活用して、リアルタイムの結果解釈を実現すること。
• 画像分類やマルウェア検出などの多様なMLアルゴリズムと分野に適用可能な汎用的なフレームワークを開発すること。
• ハードウェアに配慮したアルゴリズム設計により、分類および解釈フェーズの両方で顕著な高速化を達成すること。

提案手法

• 解釈可能機械学習におけるモデル distillation プロセスを線形代数問題にマッピングし、一連の行列演算に変換する。
• 行列畳み込みとフーリエ変換の相乗効果を活用して、解釈計算を簡素化・高速化する。
• TPUのシステリカルアレイアーキテクチャ——特に256×256行列乗算ユニット（MXU）——を活用し、8ビット整数の超高速行列乗算を実現する。
• データ分解を適用して大きな行列を小さなサブ行列に分割し、複数のMXUに効率的に分散することで帯域幅の圧力を軽減する。
• モデル distillation をTPU加速型フォワードパスと統合し、入力特徴量の寄与要因を生成することで結果の解釈を可能にする。
• 分類と解釈を1つの最適化されたパイプラインとしてTPUハードウェア上で実行可能なフレームワークを設計する。

実験結果

リサーチクエスチョン

• RQ1 TPUを用いたハードウェア加速により、解釈可能機械学習における結果解釈に要する時間が著しく短縮可能か？
• RQ2 XMLにおけるモデル distillation プロセスを、TPUで効率的に実行可能な行列計算に再定式化できるか？
• RQ3 提案されたTPUベースのフレームワークは、行列サイズやモデルの複雑さが増大するにつれてどのようにスケーリングするか？
• RQ4 提案手法は、分類および解釈速度の両面でGPUベースの最先端技術をどの程度上回るか？
• RQ5 高いパフォーマンスを維持しながら、マルウェアの攻撃モードなど、局所的で解釈可能な説明（例：重要な特徴の同定）を提供できるか？

主な発見

• 提案されたTPUベースのフレームワークは、最先端のGPUベース手法と比較して分類時間で平均25倍の高速化を達成した。
• 解釈時間は平均13倍の高速化を達成し、CIFAR-100ベンチマークではピークで39.5倍の高速化を記録した。
• 1024×1024行列では、ベースラインと比較して30倍以上も高速であり、優れたスケーラビリティを示した。
• フレームワークはMIRAIマルウェアトレースにおける「ATTACK_VECTOR」変数といった重要な特徴を効果的に同定し、実行可能な解釈可能なインサイトを提供した。
• 大規模ニューラルネットワークワークロードにおいて、計算時間を数時間から数秒に短縮し、解釈可能AIのリアルタイム展開を可能にした。
• 本手法は画像分類（CIFAR-100）やマルウェア検出（ResNet50）を含む多様な分野に一般化可能で、効果的であることが確認された。