SpAtten: Attentionの実行に特化したソフトウェアとハードウェア

MITの電気工学・コンピュータサイエンス学科の博士課程の学生であるHanrui Wangらは、計算機の性能に深く依存するAttentionの実行時のDRAMアクセスと計算量の両方を削減するためのテクニックとハードウェアを提案した。

SpAtten: Attentionの実行に特化したソフトウェアとハードウェア

自然言語処理(NLP)は、近年、Attentionメカニズムに牽引されて急速に進歩している。Transformer、BERT、GPT-2などのAttentionモデルは、畳み込みニューラルネットワーク(CNN)やリカレントニューラルネットワーク(RNN)に基づくモデルよりも大幅な性能向上を実現している。BERTは、困難な文分類 タスクでは、人間の性能をも凌駕している。

残念ながら、高い精度は、効率性を犠牲にしている。Attentionは、GPUやCPUなどの汎用プラットフォームでは、データの動きが複雑で演算強度が低いため、動作が非常に遅くなる。例えば、わずか30個のトークンで文章を表現する場合、GPT-2モデルはTITAN Xp GPUで合計370ミリ秒でアテンション推論を実行する。

これは画像を分類するのに6ミリ秒しかかからないMobileNet-V2よりも2桁遅い。リソースが限られたRaspberry Pi ARM CPUでは、Attentionに43秒かかり、対話型のアプリケーションは不可能になる。このような効率の悪さが、Attentionモデルをモバイルデバイスに展開することを阻んでいる。CNNやRNNを高速化するために多くのアクセラレータが提案されているが,それらのアクセラレータは動作が異なるため,アテンションに簡単に適用することはできない。

Attentionの入力は、クエリ(Q)、キー(K)、値(V)を含み、それぞれが複数のヘッドに分割される。Attention確率を V で乗算すると 1 つのヘッドの結果が得られ、すべてのヘッドを連結するとAttentionの出力が得られる。生成段階でのAttention度の演算強度は低く、ベクトル行列乗算(Q × K)では1データあたり 2 回の演算(0.5OPS/バイト)のみだ。生成は、GPT-2 モデルにおける全体的な待ち時間の最大の部分を占めている(32 個のトークンを生成する場合は 97%)。BERT の場合、全体的な性能は計算に依存する。

そこで、MITの電気工学・コンピュータサイエンス学科の博士課程の学生であるHanrui WangらはAttention実行時のDRAMアクセスと計算量の両方を削減するために、図1に示すようなカスケード・トークン・プルーニングを提案する。人間の言語は、前置詞、冠詞、副詞などの構造的で無意味なトークンが多く、冗長性が高いことに着想を得て、結果にほとんど影響を与えずに、重要でないトークンを安全に除去することができる。さらに、アテンションは様々な依存関係を捉えるために多くのヘッドを使用するが、その中には冗長なものもある。

図1. カスケードトークンとヘッドの枝切りは、レイヤー間でグローバルに冗長なトークンとヘッドを削除する。SST-2データセットのBERT-Baseで評価。Source: Wang, et al. (2021)

この研究は今月、IEEE International Symposium on High-Performance Computer Architectureで発表される予定だ。Wangは論文の主著者であり、共著者には、Zhekai Zhangとその指導教員であるSong Han助教授が含まれている。

Wangらが指摘する課題の1つは、Attentionメカニズムを使ったNLPモデルを実行するための専用のハードウェアがないことだ。CPUやGPUのような汎用プロセッサでは、Attentionメカニズムの複雑なデータ移動や演算のシーケンスに問題がある。そして、NLPモデルが複雑になるにつれて、特に長文の場合には問題は悪化するだろう。増え続ける計算需要を処理するためには、アルゴリズムの最適化と専用のハードウェアが必要だ。

研究者たちは、Attentionをより効率的に実行するために、SpAttenと呼ばれるシステムを開発した。彼らの設計には、特殊なソフトウェアとハードウェアの両方が含まれている。SpAttenの重要な進歩の1つは、カスケード・プルーニング(cascade pruning)、つまり計算から不要なデータを排除することだ。アテンション機構が文のキーワード(トークンと呼ばれる)を選ぶと、SpAttenは重要でないトークンを刈り込み、それに対応する計算やデータの動きを排除する。Attentionメカニズムには、複数の計算枝(ヘッドと呼ばれる)も含まれている。トークンと同様に、重要ではないヘッドは識別され、剪定される。一度ディスパッチされると、不要なトークンやヘッドはアルゴリズムの下流の計算には影響しないため、計算負荷とメモリアクセスの両方を削減することができる。

図2. 様々なプラットフォームにおけるエンドツーエンドのGPT-2のレイテンシの内訳と、TITAN Xp GPUにおけるAttentionのレイテンシの内訳。レイテンシ全体の50%以上をAttentionが占めている。Attentionによるレイテンシの73%はデータフローが占めている。Source: Wang, et al. (2021)

メモリ使用量をさらに削減するために、研究者は「プログレッシブ量子化」と呼ばれる技術も開発した。この方法では、アルゴリズムはより小さなビット幅のチャンクでデータを処理し、メモリから可能な限り少ないデータをフェッチすることができる。ビット幅が小さくなるほどデータの精度が低くなり、単純な文章には精度が高くなり、複雑な文章には精度が高くなるようになっている。

さらにWangらはDRAMアクセスを削減するために、Attention入力に対してプログレッシブ量子化を提案する。少数のトークンが分布を支配している場合、量子化誤差は小さく、MSBのみが必要となり、フラットな分布の場合はLSBとMSBの両方が必要となる。この現象については、セクションIII-Dで理論的な証明を行う。この観察に基づいて、Wangらは、Attention確率が支配的なAttentionに対してはより積極的に量子化を行い、その他のAttentionに対してはより保守的に量子化を行ったという。メモリアクセスが少なくて済むように計算を行うことができるので、メモリに依存するモデルには有利だ。プログレッシブ量子化により、さらに5.1倍のメモリアクセスを節約することができる。

要約すると、SpAttenは、精度を維持しつつ、スパースおよび量子化されたAttention計算のためのアルゴリズム・アーキテクチャの共同設計を行う。SpAttenは4つの貢献をしている。

  • カスケード・トークン・プルーニングは、累積トークン重要度スコアに応じて重要でないトークンを削除し、DRAMアクセスと計算量を最大3.8倍削減する。
  • カスケード・ヘッド・プルーニングは、重要でないヘッドを削除し、DRAMアクセスと計算量を1.1倍削減する。
  • プログレッシブ量子化では、メモリアクセスが少なくなる代わりに、計算量が少し増える。Attention度の高いヘッドやレイヤーのビット幅をAttention度分布に応じて変更することで、DRAMアクセスを5.1×削減。
  • 時間的に複雑なSpecialized High Parallelism top-k Engineを搭載し、オンザフライでのトークンとヘッドの選択を効率的にサポートする。

並列性を重視したハードウェア

これらのソフトウェアの進歩と並行して、研究者らは、メモリアクセスを最小限に抑えながらSpAttenとAttentionメカニズムを実行することに特化したハードウェアアーキテクチャも開発した。このアーキテクチャ設計では、複数の処理要素で複数の演算を同時に処理する「並列性」を重視しており、Attentionメカニズムは一文一文を一度に解析するために有効である。この設計により、SpAttenは、トークンとヘッドの重要度を(潜在的な剪定のために)少ないクロックサイクル数でランク付けすることが可能になった。全体的に、SpAttenのソフトウェアとハードウェアのコンポーネントは、不必要で非効率なデータ操作を排除し、ユーザーの目的を達成するために必要なタスクだけに焦点を当てることができる。

SpAttenの概要を図8の通りだ。トークン/ヘッドの刈り込みをサポートするために、トークン/ヘッドの重要度をランク付けする。刈り込みは計算とメモリトラフィックを削減するが、ランダムアクセスが発生する。このため、クロスバーを適用してアドレスを処理し、各メモリチャネルをビジー状態に保ち、帯域幅利用率を向上させる。プログレッシブ量子化をサポートするために、フェッチされたビットの分割や、MSBとLSBの連結を処理するために、オンチップのビット幅変換器を実装している。

図 8. SpAtten アーキテクチャの概要 クリティカルパス上のモジュールは、スループットを最大化するために完全にパイプライン化されている。 Source: Wang, et al. (2021)

SpAttenは、Attentionをヘッドごとに処理し、クエリごとにクエリを処理することで、プルーニングの粒度とパーセンテージのバランスを良くしている。

1つのヘッドのクエリが一度にパイプラインに供給され、ヘッド単位とレイヤ単位の粒度の両方でトークンの刈り込みが可能になる。インナーヘッド並列化により、チップ上のすべての計算リソースをビジー状態に保つことができるため、ヘッド間並列化は必要ない。要約段階では、カスケードトークン刈り込みに耐えたKとVがオンチップSRAMにフェッチされ、複数のクエリに渡って再利用される。生成段階では、Qは単一のベクトルであるため、KとVの再利用はなく、オンチップSRAMに格納する必要はない。

参考文献

  1. Hanrui Wang, et al. SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning. arXiv:2012.09852 [cs.AR]. [v2] Mon, 4 Jan 2021 03:49:57 UTC (2,039 KB).
  2. Daniel Ackerman. A language learning system that pays Attention — more efficiently than ever before. MIT News. Feb 10, 2021.

Photo by Joel Naren on Unsplash

Special thanks to supporters !

Shogo Otani, 林祐輔, 鈴木卓也, Mayumi Nakamura, Kinoco, Masatoshi Yokota, Yohei Onishi, Tomochika Hara, 秋元 善次, Satoshi Takeda, Ken Manabe, Yasuhiro Hatabe, 4383, lostworld, ogawaa1218, txpyr12, shimon8470, tokyo_h, kkawakami, nakamatchy, wslash, TS, ikebukurou, 太郎, bantou.

700円/月の支援

Axionは吉田が2年無給で、1年が高校生アルバイトの賃金で進めている「慈善活動」です。有料購読型アプリへと成長するプランがあります。コーヒー代のご支援をお願いします。個人で投資を検討の方はTwitter(@taxiyoshida)までご連絡ください。

デジタル経済メディアAxionを支援しよう
Axionはテクノロジー×経済の最先端情報を提供する次世代メディアです。経験豊富なプロによる徹底的な調査と分析によって信頼度の高い情報を提供しています。投資家、金融業界人、スタートアップ関係者、テクノロジー企業にお勤めの方、政策立案者が主要読者。運営の持続可能性を担保するため支援を募っています。
Takushi Yoshida is creating writing/journalism | Patreon
Patreon is a membership platform that makes it easy for artists and creators to get paid. Join over 200,000 creators earning salaries from over 6 million monthly patrons.

投げ銭

投げ銭はこちらから。金額を入力してお好きな額をサポートしてください。

https://www.paypal.com/paypalme/axionyoshi?locale.x=ja_JP

Read more

アドビ、日本語バリアブルフォント「百千鳥」発表  往年のタイポグラフィー技法をデジタルで再現

アドビ、日本語バリアブルフォント「百千鳥」発表 往年のタイポグラフィー技法をデジタルで再現

アドビは4月10日、日本語のバリアブルフォント「百千鳥」を発表した。レトロ調の手書き風フォントで、太さ(ウェイト)の軸に加えて、字幅(ワイズ)の軸を組み込んだ初の日本語バリアブルフォント。近年のレトロブームを汲み、デザイン現場の様々な要望に応えることが期待されている。

By 吉田拓史
新たなスエズ危機に直面する米海軍[英エコノミスト]

新たなスエズ危機に直面する米海軍[英エコノミスト]

世界が繁栄するためには、船が港に到着しなければならない。マラッカ海峡やパナマ運河のような狭い航路を通過するとき、船舶は最も脆弱になる。そのため、スエズ運河への唯一の南側航路である紅海で最近急増している船舶への攻撃は、世界貿易にとって重大な脅威となっている。イランに支援されたイエメンの過激派フーシ派は、表向きはパレスチナ人を支援するために、35カ国以上につながる船舶に向けて100機以上の無人機やミサイルを発射した。彼らのキャンペーンは、黒海から南シナ海まですでに危険にさらされている航行の自由の原則に対する冒涜である。アメリカとその同盟国は、中東での紛争をエスカレートさせることなく、この問題にしっかりと対処しなければならない。 世界のコンテナ輸送量の20%、海上貿易の10%、海上ガスと石油の8~10%が紅海とスエズルートを通過している。数週間の騒乱の後、世界の5大コンテナ船会社のうち4社が紅海とスエズ航路の航海を停止し、BPは石油の出荷を一時停止した。十分な供給があるため、エネルギー価格への影響は軽微である。しかし、コンテナ会社の株価は、投資家が輸送能力の縮小を予想している

By エコノミスト(英国)
新型ジェットエンジンが超音速飛行を復活させる可能性[英エコノミスト]

新型ジェットエンジンが超音速飛行を復活させる可能性[英エコノミスト]

1960年代以来、世界中のエンジニアが回転デトネーションエンジン(RDE)と呼ばれる新しいタイプのジェット機を研究してきたが、実験段階を超えることはなかった。世界最大のジェットエンジン製造会社のひとつであるジー・エアロスペースは最近、実用版を開発中であると発表した。今年初め、米国の国防高等研究計画局は、同じく大手航空宇宙グループであるRTX傘下のレイセオンに対し、ガンビットと呼ばれるRDEを開発するために2900万ドルの契約を結んだ。 両エンジンはミサイルの推進に使用され、ロケットや既存のジェットエンジンなど、現在の推進システムの航続距離や速度の限界を克服する。しかし、もし両社が実用化に成功すれば、超音速飛行を復活させる可能性も含め、RDEは航空分野でより幅広い役割を果たすことになるかもしれない。 中央フロリダ大学の先端航空宇宙エンジンの専門家であるカリーム・アーメッドは、RDEとは「火を制御された爆発に置き換える」ものだと説明する。専門用語で言えば、ジェットエンジンは酸素と燃料の燃焼に依存しており、これは科学者が消炎と呼ぶ亜音速の反応だからだ。それに比べてデトネーシ

By エコノミスト(英国)