Lightmatterは、Hot Chips 32において、「Mars」と呼ぶAI推論用アクセラレータを発表した。

このアクセラレータがその他のAIアクセラレータと大きく異なるのは、中心となるAI計算が電子的に行われるのではなく、光で行われるという点である。

ただし、AI処理の中には光での計算に向いている処理と、光でやるよりも電子的に処理する方がやりやすい処理があり、Marsでは、前者の処理はシリコンフォトニックスチップで実行させ、後者の処理は通常のシリコンチップで行わせている。

  • Mars

    Lightmatterのフォトニックにマトリクス計算を行うMarsアクセラレータボード。光で計算するので、計算は光速で、かつ消費電力が小さい (このレポートのすべての図は、LightmatterのCarl Ramey氏の発表スライドのスクリーンキャプチャである)

半導体素子の微細化に横たわる熱問題

1990年頃からトランジスタのスイッチエネルギーの低減のペースが鈍ってきている。常温で動作させる素子のスイッチエネルギーは常温の熱エネルギーより小さくはできないので、スイッチエネルギーには下限があり、これは微細化のスローダウンのだけのせいではない。

  • Mars

    トランジスタのスイッチのエネルギーは微細化に伴い減少してきたが、1990年ころからは低減の速度が低下してきている。スイッチエネルギーは熱エネルギーよりは大きい必要があり、トランジスタのスイッチ効率が悪化するのはやむを得ない

スイッチのためのエネルギーが小さくならないで、素子の密度が高くなると、チップは熱くなる。その結果、熱のせいで全部のトランジスタを動作させられないDark Siliconの問題、クロック周波数を上げられないとか、特別なエキゾチックな冷却が必要となるという問題が出てくる。

1つのパッケージに入るチップの消費電力としては400W程度が実用的には限界と考えられる。

  • Mars

    シリコンチップの発熱が大きくなると、全体を使えないダークシリコン、クロック周波数の飽和、エキゾチックな冷却などが必要になる

AIの増加する計算量を光で処理するメリット

次の図は、AIの計算量の推移を示すものであるが、AI計算の増加ペースはムーアの法則の5倍であり、エレクトロニクスでは、このペースについて行けない。

  • Mars

    AIの計算量の増加はムーアの法則の5倍の速度で、エレクトロニクスでは追いつけない

光情報を処理するシストリックアレイは、データの伝送に必要となるエネルギーが小さくて済み、光を使うテンソルコアは電子回路より高速で動作し、小さいエネルギーで動作する。そして、光は波長多重や偏波多重という技術を使うと、1つの回路で多重度倍の並列処理を行うことができる。

電子回路の場合のレーテンシは100ns程度であるが、光を使えば100ps程度と1000倍のスピードで処理ができる。信号処理の速度も電子処理では2GHz程度であるが、光を使えば10倍の20GHz程度で処理ができる。消費電力も電子回路では回路当たり1mW程度であるが、光なら1μWで済む。それで、処理回路のチップ面積は同程度であり、フォトニクスでの処理には大きなメリットがある。

  • Mars

    オプティカルなデータ伝送は少ないエネルギーで高速である。オプティカルなテンソルコアは動作クロックが高く、消費エネルギーも小さい。波長多重や偏波多重を使えば並列処理ができるなどのメリットがある

電気での信号伝送は信号線の寄生容量の充放電にエネルギーが必要で、エレクトロニクスでの信号伝送には数十Wの電力が必要になる。一方、光での伝送なら電子回路のような信号線のRCによる遅延は発生せず、伝送に必要なエネルギーも小さくて済む。

  • Mars

    光でのデータ伝送はRCに起因する遅延も無く、低エネルギーで伝送ができる。

光を用いた演算手法

光回路では、Mach Zehnder Interferometer(MZI)が計算を行うための基本コンポーネントである。MZI素子では、次の図のように入力を2つに分岐し、分岐したパスに可変のフェースシフトを行う素子を入れる。その後、分岐した信号を纏めて出力を作る。

それぞれのフェーズシフト量φ1とφ2が等しければ入力と同じ信号が出力される。一方、フェーズシフト量が180度異なっていると、φ1とφ2側の信号は打ち消しあって出力はゼロになる。右下の棒グラフはφ1とφ2のフェーズの差を横軸にとり、出力の信号をプロットしたもので、位相差が180度となるグラフの中央では、信号が打ち消されて出力がゼロになっている。

  • Mars

    Mach Zehnder干渉計はフォトニック回路の基本コンポーネントで、分岐したパスの遅延の差を180度にすると、信号が打ち消しあって出力がゼロになる

フェーズシフトを可変にするやり方として、熱による方法、P/N接合を使う方法、機械的な変形を使う方法がある。P/N接合を使う方法は、動作速度は速いが大きな電気を食う。また、素子が大きいという問題がある。

LightmatterのMarsでは、Nano Optical Electro Mechanical System(NOEMS)を使っている。NOEMSは貯めている静電気の力でウエーブガイドを曲げて位相を変える。ウエーブガイドの変形は電荷で記憶され、消費電力は非常に小さい。そして、数100MHzという高速で動作する。

  • Mars

    フェーズシフトを可変にする方法は色々とあるが、LightmatterはNano Optical Electro Mechanical Systemを使っている。ウエーブガイドを静電気の力で曲げて位相遅延を変える。溜まっている静電気の量で遅延が変わり、静電気はほとんど逃げないので、電力を消費しない。数100MHzの速度で電荷の量を変えられるという

方向性結合器(Directional Coupler)は2つの入力信号を結合することができる。そして、2つのフェーズシフタで図中に示した位相シフトを与えてやる。次の図のように、2つの方向性結合器の間にこの2つのフェーズシフタを挟んだ回路で2×2の行列に1×2のベクタを掛けるという計算をオプティカルに実行することができる。

この行列とベクトルの乗算器は光のスピードで動作し、ほぼ電力ゼロで動作する。

  • Mars

    方向性結合器、2つのMZI、方向性結合器を直列に繋いだ光回路で、2×2の行列と1×2のベクトルの積を計算できる。オプティカルな演算なので、光の速さで動作し、消費電力もほぼゼロである

この2×2行列と1×2ベクトル乗算器を次の図のように接続すれば、より大きな行列の掛け算が実行できる。

  • Mars

    2×2の行列の掛け算器をこの図のように接続すれば、より大きな行列の掛け算が行える

この乗算器は光で演算を行い出力も光であるので、電子回路との接続を行う場合は、アナログの光信号出力をフォトダイオードで電気信号に変換し、さらにADCでアナログ信号をデジタル信号に変換する。Marsでは、64個のDACで8bitのデジタルの入力信号をアナログに変換してMZIの制御信号として与え、出力側では光ダイオードからのアナログ信号を64個のADCでデジタル信号に変換している。

この回路では、光の積和演算器はほとんど電力を消費せず、電力を消費するのは入力回路と出力回路である。このため、この回路の消費電力は積和演算器の数の平方根に比例する。

  • Mars

    演算器の性能は積和演算器の数に比例するが、フォトニックなマトリクスの部分はほとんど電力を消費せず、消費電力はマトリクスの1辺の長さに比例する。これは面白い性質である

光の特徴を活用し、多重処理を可能に

光の性質として、アナログ的な強弱だけでなく、光の色(周波数)と偏波という性質があり、周波数や偏波が異なれば、同じ回路でも別の信号を伝送し、処理することができる。実際にどれだけの多重処理ができるかは、周波数や偏波がどれだけ混信せずに送受信できるかによる。

  • Mars

    光には強度に加えて、波長(色)と偏波という性質がある。これを利用して、波長多重、偏波多重を行えば、多重度の倍数だけ伝送バンド幅を増やすことができる

次の図にMarsのフォトニックコアチップの写真を示す。Marsは64要素のベクトルの64×64のマトリクスを処理するチップとなっている。レーテンシは200psで、ベクタの処理速度は1GHzで動作する。入力信号と言うか入力電力というかは主観によるが、50mWのレーザー光で動作する。

チップは90nmの標準的なフォトニックプロセスで作られており、チップサイズは150mm2である。チップは64×64のフォトニックなマトリクスが最大の面積を占め、重みのモジュレータとベクタのエンコーダらがそれに次ぐ面積を占めている。

  • Mars

    Marsのフォトニックコアのダイ写真。90nmの標準的なフォトニックプロセスで作られており、チップサイズは150mm2である

そして、光では作りにくい部分を14nmのカスタムASICで作っている。ASICで作られているのは、フォトニックコアとのアナログインタフェース、重みとアクティベーションのSRAM、I/Oインタフェース、デジタル部分などである。

  • Mars

    SoCの部分は通常の14nm フォトニックCMOSプロセスで作られており、チップサイズは50mm2。フォトニックコアとのアナログインタフェース、重みやアクティベーションを記憶するSRAM、I/Oやデジタルの部分が収容されている

デジタル部分は、アクティベーションを扱うVSliceと重みを記憶するWSliceが並んでおり、WSliceはデジタルの重みをアナログに変換するDACを持っている。VSliceはベクタ処理を行うパイプラインが含まれている。

  • Mars

    デジタルチップにはアクティベーションを供給するVSliceと重みを供給するWSliceとパイプラインのスケジューラが作られている

VSliceはアクティベーションの出力にDACを持つパイプラインが多数個あり、それぞれの入力に対応している。

  • Mars

    デジタルパイプラインはオペランドをローカルに記憶し、SRAMのアクセスを最小化している

Marsの電力はどこで消費されるのか?

重みはMZIに近い場所に電荷として蓄えられており、記憶のための消費電力はほぼゼロである。また、64×64のフォトニック乗算器の消費電力もほぼゼロである。

重みのアップデートのためのダイナミックな電力消費は必要であるが、この電力は電子的な回路の場合と同程度の消費電力である。電子回路と同様に、入力を纏めて処理するバッチ化を行えば重みのアップデートの電力は少なくて済むようになる。

したがって、Marsの場合の電力消費の大部分はデータ変換に伴うものであり、Marsは計算のためには電力を使っていない。

  • Mars

    Marsでは重みの記憶はMZIの近くに電荷として記憶しており、記憶のための電力はほぼゼロである。また、フォトニックな計算の電力もほぼゼロであり、電力消費の大部分はデータ変換のために使っている

次の図はMarsの電力消費の内訳を示すものである。そして、 ResNet-50で99%の精度を実現したケースで、Marsの全体の消費電力はレーザー光源の電力を含めて3W以下で、その大部分がデータ移動のための電力である。

電力消費の内訳は、48.7%がデジタル部のデータ移動などのダイナミックな消費電力、ベクタのDACが11.3%、重みのDACが10.3%、レーザーが8.8%、ADCが12.7%といったところが大口の電力消費者である。

  • Mars

    Marsの電力消費の内訳。デジタル部分のダイナミックな動作に全体の48.7%の電力が使われている。また、DACとADCで、合計34.4%の電力を使っている

そして、Marsチップは、次の写真のように、90nmプロセスで作られたフォトニックチップの上に、14nmプロセスのASICチップを搭載した構造になっている。フォトニックチップの電力が小さいので、チップを重ねても放熱の問題は生じない。そして、チップのスタッキングでデータ移動の距離が減り、エネルギーも減るという効果が出ている。

  • Mars

    Marsはオプティカルチップの上にデジタルのSoCを積み重ねて搭載している。このスタッキングでデータの移動距離を縮めて、さらに消費電力を減らしている

MarsアクセラレータはPCI Expressカードの形態で作られており、ホストCPUに接続される。マシンラーニングのフレームワークやコンパイラはホストCPUで実行される。

フレームワークとしては、PyTorch、TensorFlowとONNXをサポートしている。Lightmatterのソフトウェア環境で、モデルのシミュレーションを行い、コンパイルや実行を行うことができる。そして、実行中のモデルの内部状態をモニタできるデバッガなども提供している。また、実行中のプロファイルを取り、チューニングを行えるツールも提供している。

  • Mars

    LightmatterはPyTorchやTensorFlowなどのフレームワークで書いたモデルを入力として使う。Lightmatterは、シミュレータやコンパイラ、デバッガ、プロファイラなどのソフトウェアツールを提供している

Marsはフォトニクスを使ってAI計算を行う推論アクセラレータである。このアクセラレータは、推論に関しては汎用のアクセラレータである。

  • Mars

    Marsはフォトニクスで計算を行う推論用の汎用アクセラレータである

Marsは実用的な技術となるのか?

結論であるが、オプティカルなAI計算が現実的なものとなった。LightmatterのMarsボードはマシンラーニングの中心となるマトリクス×ベクトルの計算をフォトニクスで実行し、アクティベーションやI/Oなどは電子回路で実行している。

オプティカルチップの上に電子回路のチップを搭載するスタック構造を取っており、アクティベーションとコンピュートの間の距離が短くなっている。オプティカルチップの消費電力は小さく、より大きなマトリクス演算器を搭載したり、より多くのSRAMを搭載することが可能である。

  • Mars

    Marsは現実に製品として販売されているフォトニクスで計算を行う推論用のアクセラレータである。フォトニクスで動作するので、高速で、消費電力が少ない。現在のMarsは64×64の行列処理であるが、より大型のチップを作るのに電力的な障害は無い

計算にはほとんど電力を消費せず、そして、計算速度も電子回路に比べて10倍速いという夢のような技術である。しかし、ニューラルネットの計算とは言え、アナログ計算でどの程度の規模まで実用的な計算ができるのかは気になるところである。

デジタル技術で補正してアナログ回路の計算精度を上げることは可能であろうが、動作速度の低下や消費電力の増加が起こることも考えられ、実用的な技術に育つのかどうかは注意深く見ていく必要があろう。