AlphaTensorのゲームの部分

昨年10月、DeepMindが強化学習を用いて行列の積を従来の手法よりも高速に計算できるアルゴリズムを発見したと発表した。このニュースを見たとき強化学習と新たなアルゴリズムの発見が自分の中で結びつかなかったので、元論文を読んで社内の勉強会で発表した。その内容をこちらにもまとめ直してみる。強化学習をアルゴリズムの発見に適用する過程が特に気になっていたので、その部分に注目した内容になっている。

元論文

Discovering faster matrix multiplication algorithms with reinforcement learning

既存の行列演算アルゴリズム

定義

行列積の愚直な計算方法は、各行と列に対して定義通りに計算していくというものである。

$\displaystyle \left( \begin{array}{rr} a_{11} & a_{12} \\ a_{21} & a_{22} \\ \end{array} \right) \left( \begin{array}{rr} b_{11} & b_{12} \\ b_{21} & b_{22} \\ \end{array} \right) = \left( \begin{array}{rr} a_{11}b_{11} + a_{12}b_{21} & a_{11}b_{12} + a_{12}b_{22} \\ a_{21}b_{11} + a_{22}b_{21} & a_{21}b_{12} + a_{22}b_{22} \\ \end{array} \right)$

この方法では、 $N\times N$ 行列同士の積を計算するために乗算が $N^3$ 回、加算が $N^2(N-1)$ 回必要となる。この例では $2\times2$ 行列なので、乗算が8回と加算が4回となっている。この方法はどんな大きさの行列にも適用できるから、これよりも遅いアルゴリズムは意味がない。

Strassenのアルゴリズム

定義による計算方法では $2\times2$ 行列の積を求めるために8回の乗算を要したが、これを7回の乗算で計算できる方法がある。それは、以下のようなものである。

$\displaystyle \left( \begin{array}{rr} a_{11} & a_{12} \\ a_{21} & a_{22} \\ \end{array} \right) \left( \begin{array}{rr} b_{11} & b_{12} \\ b_{21} & b_{22} \\ \end{array} \right) = \left( \begin{array}{cc} p_1 + p_4 - p_5 + p_7 & p_3 + p_5 \\ p_2 + p_4 & p_1 + p_3 - p_2 + p_6 \\ \end{array} \right)$

ここで、

$\displaystyle \begin{align*} p_1 &= (a_{11} + a_{22})(b_{11} + b_{22}) \\ p_2 &= (a_{21} + a_{22})b_{11} \\ p_3 &= a_{11}(b_{12} - b_{22}) \\ p_4 &= a_{22}(b_{21} - b_{11}) \\ p_5 &= (a_{11} + a_{12})b_{22} \\ p_6 &= (a_{21} - a_{11})(b_{11} + b_{12}) \\ p_7 &= (a_{12} - a_{22})(b_{21} + b_{22}) \\ \end{align*}$

である。このアルゴリズムでは各 $p$ を計算する際に1度だけ乗算を行っているので、合計で7回の乗算を用いている。代わりに加算の回数は愚直に計算した場合よりも大幅に増えて、18回になっている。

ブロック行列の積

行列の積を計算する際、行列 $A, B$ それぞれをブロック（部分行列）に分け、それらを行列の成分とみなして積を考えてもよい。たとえば

$\displaystyle \begin{align*} \left( \begin{array}{rr|r} a_{11} & a_{12} & a_{13} \\ a_{21} & a_{22} & a_{23} \\ \hline a_{31} & a_{32} & a_{33} \\ \end{array} \right) \left( \begin{array}{rr|r} b_{11} & b_{12} & b_{13} \\ b_{21} & b_{22} & b_{23} \\ \hline b_{31} & b_{32} & b_{33} \\ \end{array} \right) &= \left( \begin{array}{r|r} A_{11} & A_{12} \\ \hline A_{21} & A_{22} \\ \end{array} \right) \left( \begin{array}{r|r} B_{11} & B_{12} \\ \hline B_{21} & B_{22} \\ \end{array} \right) \\ &= \left( \begin{array}{r|r} A_{11}B_{11} + A_{12}B_{21} & A_{11}B_{12} + A_{12}B_{22} \\ \hline A_{21}B_{11} + A_{22}B_{21} & A_{21}B_{12} + A_{22}B_{22} \\ \end{array} \right) \\ &= \left( \begin{array}{rr|r} a_{11}b_{11} + a_{12}b_{21} + a_{13}b_{31} & a_{11}b_{12} + a_{12}b_{22} + a_{13}b_{32} & a_{11}b_{13} + a_{12}b_{23} + a_{13}b_{33} \\ a_{21}b_{11} + a_{22}b_{21} + a_{23}b_{31} & a_{21}b_{12} + a_{22}b_{22} + a_{23}b_{32} & a_{21}b_{13} + a_{22}b_{23} + a_{23}b_{33} \\ \hline a_{31}b_{11} + a_{32}b_{21} + a_{33}b_{31} & a_{31}b_{12} + a_{32}b_{22} + a_{33}b_{32} & a_{31}b_{13} + a_{32}b_{23} + a_{33}b_{33} \\ \end{array} \right) \end{align*}$

のように計算できる。ブロック行列の性質を用いれば、 $2\times2$ 行列の積のアルゴリズムを任意の大きさの行列積に対して適用することができる。たとえばStrassenのアルゴリズムを用いたとすると、ブロック行列どうしの積が7回発生するが、それぞれの積に対してStrassenのアルゴリズムを再帰的に適用することによって大きな行列の積を計算することができる。このとき、必要な乗算の回数は $N^{\log_27}$ 回となり、愚直に計算した場合よりも高速になる。行列の和と積の計算は常に後者のほうが重いため、アルゴリズムを再帰的に適用することを考えたとき、乗算の回数が少ないアルゴリズムのほうが計算量が小さく優秀である。

ここで重要なのは、何らかのヒューリスティックな方法を用いてある大きさの高速な行列積アルゴリズムを発見できれば、ブロック行列の性質を適用してそれ以上の大きさの行列積も高速化できるということである。 AlphaTensorでは強化学習を用いて特定の大きさの行列積の高速なアルゴリズムを発見することを目的としている。簡単のために正方行列で説明しているが、行列が長方形であってもまったく問題ない。

TensorGame

AlphaTensorでは強化学習を行列積のアルゴリズム発見に適用するために、「クリアしたら行列積のアルゴリズムを発見したことになるゲーム」を設計し、これをTensorGameと呼んでいる。

Matrix multiplication tensor and algorithms. （元論文より引用）

状態

$N\times N$ 行列どうしの積のアルゴリズムを発見したいとき、TensorGameの盤面として $N^2\times N^2\times N^2$ のテンソル $T$ を用いる。テンソルの要素 $T_{ijk}$ は、行列 $A, B$ の積を $C$ としたとき、その要素 $c_k$ における $a_ib_j$ の係数を表している。つまり、TensorGameの盤面のバリエーションは、 $C$ の各要素が任意の $a_ib_j$ の線形結合であるような状態を表現できる。

引用画像aは $2\times2$ 行列どうしの積のアルゴリズムを発見するTensorGameの盤面である。この図では $T$ のうち0の要素を半透明に、1の要素を不透明に表示している。つまり

$\displaystyle \begin{align*} c_1 = a_1b_1 + a_2b_3 \\ c_2 = a_1b_2 + a_2b_4 \\ c_3 = a_3b_1 + a_4b_3 \\ c_4 = a_3b_2 + a_4b_4 \\ \end{align*}$

という状態を表している（これは正しい行列積の表示になっている）。TensorGameでは各要素を任意の整数にすることができるが、正しい行列積を表している状態は1つだけである。

TensorGameには2つの特別な状態が存在すると考えることができて、それは $T$ のすべての要素が0の状態と、 $T$ が「正しい行列積の表示」になっているような状態である。このうちの一方の状態からスタートして、ルールに従って $T$ に決められた操作を行うことでもう一方の状態に遷移させることができればクリアとなる。AlphaTensorは、このゲームを最短手数でクリアすることを目指す。

アクション

TensorGameでは、プレイヤーが取れるアクションは $N^2$ 要素のベクトル3つで表される。たとえば、

$\displaystyle u_1= \left( \begin{array}{rr} 1 \\ 0 \\ 0 \\ 1 \end{array} \right),\; v_1= \left( \begin{array}{rr} 1 \\ 0 \\ 0 \\ 1 \end{array} \right),\; w_1= \left( \begin{array}{rr} 1 \\ 0 \\ 0 \\ 1 \end{array} \right)$

はそのうちの一つである。 TensorGameにおける「一手」は「行列 $A, B$ の各要素の線形結合をそれぞれ1つずつ作り、それをかけ合わせたうえで $C$ の各要素に好きな係数で足す」操作を表している。

引用画像bは前述のStrassenのアルゴリズムを表しているが、先に例示した $u_1,v_1,w_1$ はこの図における $m_1$ に対応している。つまり、 $u_1 = (1, 0, 0, 1), v_1 = (1, 0, 0, 1)$ に対応する $(1a_1 + 0a_2 + 0a_3 + 1a_4)(1b_1 + 0b_2 + 0b_3 + 1b_4)$ を作り、それを $c_1, c_2, c_3, c_4$ に係数 $w_1=(1, 0, 0, 1)$ で加算しているということになる。

ここで重要なのは、TensorGameにおける一手は乗算をかならず1回だけ含むという点である。このことから、ゲームをクリアするまでに要した手数がそのまま完成したアルゴリズムに含まれる乗算の数となり、ゲームを最速でクリアする手順を求めることが最速のアルゴリズムの発見と等価であるということができる。

引用画像cはStrassenのアルゴリズムをTensorGameのアクションで表したものであり、ボードゲームにおける棋譜に対応するものである。これによると、Strassenのアルゴリズムは「7手」のアルゴリズムであると考えることができる。

成果

AlphaTensorはAlphaZeroをベースにした強化学習アルゴリズムによってさまざまなサイズのTensorGameの解法を学習し、その過程で生み出された解法には既存のアルゴリズムを上回る性能のものがあった。その中には $4\times4$ における47手（既存の最速は49手）や、 $5\times5$ における96手（既存の最速は98手）などが含まれる。

その他興味深かったトピック

元論文中には学習の方法や手法の応用についても書かれており、興味深いものがいくつかあった

上で紹介したTensorGameは時間計算量の考え方にしたがって乗算の回数が最小になるアルゴリズムを求めるものだったが、理論上最速のアルゴリズムがGPUやTPUの実機上でも最速とは限らない。TensorGameではスコアを「乗算の回数」ではなく「実機でそのアルゴリズムを用いたときにかかった時間」とすることで各デバイスに最適化した行列積のアルゴリズムを発見することができた。
AIの文脈では囲碁や将棋といったボードゲームは大きな状態空間をもっていることがよく言及されるが、TensorGameは状態空間だけでなく行動空間も非常に大きい。実際の学習では各ベクトルの要素が取り得る値をある範囲の整数に限定しているが、それでも各アクションの取り得る範囲は膨大であり、既存手法のようにすべてのアクションの評価値を求めて比較するといった方法を取ることができない。これを解決するためにモデルのPolicy HeadにTransformerのアーキテクチャを使用し、文章を生成するようにアクションを生成した。

など。ほかにも学習上の工夫のうち別のゲームに応用できそうな部分があった。

感想

囲碁や将棋をはじめとして、強化学習で訓練されたエージェントがボードゲームにおいて人間を上回るパフォーマンスを発揮する事例は多く知られているので、今回はアルゴリズムの発見をゲームに落とし込む過程にとくに注目した。「AIが強化学習で新しいアルゴリズムを発見した」というニュースを聞いたとき、さまざまなアルゴリズムに適用できる汎用的な手法が開発されたのかな、と思ったが、実際に論文を読んでみると問題設定の部分で人間がかなりサポートを行っており、他のアルゴリズムにそのまま適用できるものではなさそうだった。

環境を用意できるがその詳細が完全に分かっていない場合に強化学習は特に力を発揮するが、GPUの実機で高速なアルゴリズムを発見するということは机上の研究ではできないため、非常に興味深かった。今後発売されるデバイスの中には専用のアルゴリズムを強化学習でチューニングして性能を上げたモデルが出てくるかもしれない（？）