[論文レビュー] EcoFormer: Energy-Saving Attention with Linear Complexity

Jing Liu, Zizheng Pan+3 more|arXiv (Cornell University)|Sep 19, 2022

Advanced Neural Network Applications被引用数 16

ひとこと要約

EcoFormer は、核化ハッシュを用いてクエリとキーをコン pact なバイナリコードにマッピングすることで、ペアワイズ類似性を保持しながら、高コストな乗加算演算を低エネルギーの累積演算に置き換える線形計算量でエネルギー効率の良いアテンション機構を提案する。ImageNet-1K において標準アテンションと比較して 73% のオンチップエネルギー消費削減を達成し、精度低下はたった 0.33% にとどまる。

ABSTRACT

Transformer is a transformative framework that models sequential data and has achieved remarkable performance on a wide range of tasks, but with high computational and energy cost. To improve its efficiency, a popular choice is to compress the models via binarization which constrains the floating-point values into binary ones to save resource consumption owing to cheap bitwise operations significantly. However, existing binarization methods only aim at minimizing the information loss for the input distribution statistically, while ignoring the pairwise similarity modeling at the core of the attention. To this end, we propose a new binarization paradigm customized to high-dimensional softmax attention via kernelized hashing, called EcoFormer, to map the original queries and keys into low-dimensional binary codes in Hamming space. The kernelized hash functions are learned to match the ground-truth similarity relations extracted from the attention map in a self-supervised way. Based on the equivalence between the inner product of binary codes and the Hamming distance as well as the associative property of matrix multiplication, we can approximate the attention in linear complexity by expressing it as a dot-product of binary codes. Moreover, the compact binary representations of queries and keys enable us to replace most of the expensive multiply-accumulate operations in attention with simple accumulations to save considerable on-chip energy footprint on edge devices. Extensive experiments on both vision and language tasks show that EcoFormer consistently achieves comparable performance with standard attentions while consuming much fewer resources. For example, based on PVTv2-B0 and ImageNet-1K, Ecoformer achieves a 73% on-chip energy footprint reduction with only a 0.33% performance drop compared to the standard attention. Code is available at https://github.com/ziplab/EcoFormer.

研究の動機と目的

• 標準トランスフォーマーの高いエネルギー消費を低減すること、特にエッジデバイスにおいて。
• 従来のバイナリ化手法がアテンション機構におけるペアワイズ類似性を無視するという制限を克服すること。
• ソフトマックスアテンション機構に特化したバイナリ化戦略を開発し、トークンレベルの類似性関係を保持すること。
• 計算複雑度を線形に保ちつつ、性能劣化を最小限に抑えること。
• リソース制約のあるエッジデバイスにおけるスケーラブルなトランスフォーマーモデルの展開を、エネルギー効率の良いハードウェアフレンドリーな演算によって可能にすること。

提案手法

• RBFカーネルを用いた核化ハッシュを用い、高次元のクエリとキーをハミング空間内の低次元バイナリコードにマッピングする。
• 自己教師あり学習によりハッシュ関数を訓練し、ハミング類似度（バイナリコードの内積）を真値のアテンションスコアと一致させる。
• 行列乗算の結合的性質を活用して、アテンションをバイナリコードの内積として線形計算量で表現する。
• 大部分の乗加算演算を単純な累積演算に置き換えることで、オンチップエネルギー消費を削減する。
• 標準アテンションの類似度行列をハッシュ関数学習の監督信号として使用する。
• ハイパーパrameter $ m $ を用いたサポートサンプル戦略を導入し、訓練中のハッシュ関数の一般化性能を向上させる。

実験結果

リサーチクエスチョン

• RQ1 自己アテンションにおけるクエリとキーのバイナリ化を、統計的量子化誤差の最小化ではなく、ペアワイズ類似性関係の保持を目的としてカスタマイズ可能か？
• RQ2 RBFカーネルを用いた核化ハッシュは、競争力のある性能を維持しつつ、線形計算量アテンションを実現可能か？
• RQ3 FP32/FP16 の乗加算演算をバイナリ累積演算に置き換えることで、トランスフォーマーのエネルギー消費をどの程度削減できるか？
• RQ4 サポートサンプル数 $ m $ が、提案手法の性能およびエネルギー効率に与える影響はいかほどか？
• RQ5 EcoFormer は、視覚および自然言語のベンチマークにおいて、標準アテンションと同等の精度を達成しながら、計算コストとエネルギー消費を著しく低減できるか？

主な発見

• PVTv2-B0 において、EcoFormer は標準アテンションと比較してオンチップエネルギー消費を 73% 削減し、ImageNet-1K での精度低下はたった 0.33% にとどまる。
• CIFAR-100 において、$ m=25 $ のサポートサンプルを用いた EcoFormer は、少ないサポートサンプルを用いたモデルを上回る 71.23% の Top-1 精度を達成する。
• 計算複雑度を $ frac{1}{2}N^2 $ から $ frac{1}{2}N $ に削減し、線形時間アテンション計算を実現する。
• 複数のベンチマークで検証された結果、EcoFormer は他の効率的アテンション手法と比較して顕著に低い遅延およびオンチップエネルギー消費を示す。
• ImageNet-1K で訓練をスクラッチから行う実験では、EcoFormer は PVTv2-B0 で乗算回数を 72% 減少させ、エネルギー消費を 72% 低減しながらも高い性能を維持する。
• バイナリ化による性能劣化の主な要因は、非微分可能な符号関数における勾配近似に起因し、高度な最適化手法によりこれを緩和できる。