深層学習に特化したプロセッサ、Cloud TPU の設計

この投稿は米国時間 8 月 30 日に投稿されたもの（投稿はこちら）の抄訳です。

Posted by Google Cloud デベロッパーアドボケイト 佐藤一憲

Tensor Processing Unit (TPU) は、機械学習の計算処理のために Google がゼロから開発したカスタム LSI （ASIC）です。Google 翻訳、フォト、検索、そして Gmail 等の主要な Google サービスの運用を支える基盤技術です。Google Cloud では、機械学習を利用するエンジニアやデータサイエンティストに向けて、この TPU の能力をスケーラブルで使いやすいクラウドサービスとして提供する  Cloud TPU を提供しています。今年 6 月にサンフランシスコで開催された Google Next ‘18 では、無償トライアルを含むすべてのユーザーに向けて Cloud TPU v2 の一般提供（GA）が開始され、また新たに Cloud TPU v3 のアルファ版がアナウンスされました。

この TPU について、多くの方から「CPU や GPU とはどう違うのか？」という質問をいただきます。そこで Google では、この疑問に応えるプレゼンテーションとアニメーションを掲載したデモコンテンツ（日本語字幕付き）を公開しました。

この記事では、このコンテンツからハイライト部分を抜粋して紹介します。

ニューラル ネットワークの動き

例として、手書き数字画像を認識する簡単な一層のニューラル ネットワークを考えてみます。

入力データとして 28 x 28 のモノクロ画像を想定すると、合計 784 個の値でこの画像を表せます。ここで、数字の「8」を認識するニューロンは、個々の値に対してパラメータ値（上図の赤い線）を掛け算します。

このパラメータは、入力データの特徴を抽出し、画像が「8」の字にどの程度似ているかを測るフィルターとして働きます。以下のような動きです。

ニューラル ネットワークによるデータ分類は、おおざっぱには上図の仕組みで動いています。つまり、入力データに対して個々のパラメータを掛け算（上図で円に色を付ける部分）し、その結果を足し算（右側で円を並べる部分）しています。その結果、もし一番大きな数値が得られれば、入力データとパラメータが最もよく似ていることを表し、それが認識結果となります。

このように、ニューラル ネットワークでは入力データとパラメータの間で大量の掛け算（乗算）と足し算（加算）が必要となります。この乗算と加算は通常、高校や大学で習う行列演算として実行されます。つまり、この乗算と加算をいかに速く、少ない電力で実行できるかが、ニューラル ネットワークさらにはディープラーニング開発の効率性を大きく左右します。

CPU の動き

CPU の最大の特徴は、その柔軟性にあります。ノイマン アーキテクチャのおかげで、あらゆる種類のソフトウェアを読み込み、何百万という種類のアプリケーションに対応できます。例えばワードプロセッサで文書を書いたり、ロケット エンジンを制御したり、銀行の口座振替を扱ったり、ニューラル ネットワークで画像認識したりと、多種多様な用途に CPU を使えます。

しかしこの柔軟性のために、CPU 自体は「次にどのような計算をすべきか」を前もって知ることができません。毎回、ソフトウェアからの指示を受ける必要があります。そのため CPU は、個々の計算のすべてにおいて、計算結果を CPU 上のメモリ（レジスタや L1 キャッシュ）に保存しなければなりません。このメモリの読み書きにともなう制約はノイマン ボトルネックと呼ばれます。例えば大規模なニューラル ネットワークの計算では、実行すべき乗算と加算の手順はあらかじめ決まっています。それでも、CPU の ALU（乗算器と加算器を備える演算装置）は、個々の計算をひとつずつ実行し、その結果を毎回メモリに保存します。これが CPU による計算のスループットを制限し、また多大な電力を消費しています。

GPU の動き

この GPU のアーキテクチャは、ニューラル ネットワークの行列演算のような大規模並列計算にぴったり適しています。実際に、ディープラーニングの学習を GPU で行うと、CPU に比べて桁違いの性能が得られます。そのため、ディープラーニング用途では今もっともポピュラーな計算環境として GPU が利用されています。

しかし、GPU もやはり、数百万もの異なるアプリケーションやソフトウェアに対応する汎用プロセッサとして設計されています。そのため、ノイマン ボトルネックの本質的な制約からは逃れられません。GPU の ALU が個々の計算を行うたび、レジスタやシェアードメモリなどのメモリに計算結果を保存し、また読み出し直す必要があります。GPU は数千個の ALU を備えるため、このメモリの読み書きによって生じる電力消費もそれに比例して巨大なものとなり、また複雑な配線のためにチップ面積の多くを専有します。

TPU の動き

この成功のカギは、ノイマン ボトルネックの大半を解消できたことにあります。TPU のやるべきことは行列の乗算と加算だけなので、TPU の設計者には必要なすべての計算の手順があらかじめ分かっています。そこで、数千個の乗算器と加算器を大規模なマトリクス状に並べ、それらを配線によって直接結んでいます。これはシストリック アレイと呼ばれるアーキテクチャです。Cloud TPU v2 の場合、1 つのプロセッサは 128 x 128 のシストリック アレイを 2 つ搭載し、合計で 32,768 個の 16 ビット ALU を集約しています。

このシストリック アレイがニューラル ネットワークの計算をどのように実行するのか見てみましょう。まず、TPU はメモリからパラメータを読み込み、乗算器と加算器のマトリックスに読み込みます。

つづいて、TPU は入力データをメモリから読み込みます。個々の乗算が進むにつれて、その結果は隣の乗算器へとバケツリレー式に渡され、同時に加算も行われます。その結果、マトリックスからはすべてのデータとパラメータ同士の乗算結果を加算した値が出てきます。この大規模な計算処理が進む間、メモリへの読み書きは一切発生しません。

この設計によって、TPU ではニューラル ネットワークの計算においてきわめて高い計算スループットを達成しつつ、消費電力とチップ面積をより抑えることが可能になりました。

コストが 1/5 に

スタンフォード大学が公開する DAWNBench では、ディープラーニングの学習と推論を対象とした各種のベンチマーク結果が掲載されています。さまざまな用途やモデル、計算環境の組み合わせで、それぞれの学習や推論に要したコストや時間が発表されています。

2018 年 4 月に DAWNBench のコンテストが締め切られた時点で、TPU 以外のプロセッサによる学習コストの最低額は $72.40 でした（ResNet-50 を ImageNet データセットとスポットインスタンスで 93% 精度まで学習させた場合）。一方、Cloud TPU v2 の Preemptible 価格を適用した場合、同じ学習作業を $12.87 で終えることができます。これは TPU を使わない場合の 1/5 です。この差こそが、ニューラル ネットワーク向けにドメイン特化アーキテクチャのプロセッサを開発することで得られるメリットです。

次のステップ

謝辞