Faster Matrix Multiplication: Unveiling Insights from Strassen to FBHHRBNRSSSHK, Continuing with KM, and Advancing to AIH
普段はプレプリントの段階ではあまり研究公開の話はしないのですが、今回だけはちょっと特別とします。たぶんプレスリリースとかもしないだろうし、でも滅多にない世界記録更新した研究なので、少し経緯をここで解説します。
現在、卒業研究で受け入れている荒井くんの研究成果をプレプリントとして公開しました。
Adaptive Flip Graph Algorithm for Matrix Multiplication
Yamato Arai, Yuma Ichikawa, Koji Hukushima
arXiv:2312.16960
昨年、2022年10月に話題になったNature論文がこれです。
Discovering faster matrix multiplication algorithms with reinforcement learning
この論文では、50年以上更新されていなかった行列積のアルゴリズムの最短手記録をAlpha Tensorを使った強化学習が更新した研究が報告されています。高校の教科書で習う2×2行列の積は、行列要素の掛け算を8回必要とします。一般にnxnの正方行列の場合にnの3乗の計算が必要です。しかし、1969年にStrassenが7回で済む方法を提案しました。これが行列積アルゴリズムの研究の発端といえます。このNature論文では4×4行列に対して、素朴には64回の掛け算が必要で、Strassenのアルゴリズムを繰り返し使うと49回のところ、それよりも更に少ない47回で済む方法を発見しました。さらに、5×5行列では96回の方法を提示しました。50年以上の時間があいたこと、そしてDeepMindのAlpha Tensorがみつけたことがこの研究のインパクトを高めました。
０と1の数しか扱わない行列積の問題はrank 1テンソルの分解問題として定式化され、高速な行列積の方法を探す問題はそのテンソル分解のある種のランクの最小化問題に帰着されます。Strassenの見出したアルゴリズムもランク７のテンソル分解とまとめることができ、発見法的なアプローチに加えて、最適化問題として探索的アプローチが可能になり、強化学習によるAlpha Tensorの研究につながりました。
国内のネットで話題になったのはここまでくらいかと思いますが、これで終わりではありませんでした。そのNature論文が公開された10月5日から3日後の10月8日にプレプリント・サーバーに次の論文が投稿されました：
The FBHHRBNRSSSHK-Algorithm for Multiplication in ℤ5×52 is still not the end of the story
Manuel Kauers, Jakob Moosbauer
FBHHRBNRSSSHKはAlpha Tensor論文の著者たちの頭文字です。この中のHには、Demis HassabisのHも含まれています。この論文ではAlpha Tensorの5×5の記録を更新して95回の方法を発見しました。具体的な探索アルゴリズムは彼らの別の論文で、Flip graphと呼ばれる探索空間を適切に定義して、そのグラフ上のランダムウォーク探索を提案しています。彼らの方法を適当な初期条件から適用しても97回しか見つからず、Alpha Tensorの96回の方法を初期条件としたときにのみ95回の方法を見つけました。このKauers-Moosbauerの研究に関しても、いくつかのネット上の解説をみつけることができます。たとえば、QuantaMagazineの解説やほぼ同じ内容のYoutubeがあります。解説にMoosbauerさんが登場しますが、主題はAlpha Tensorの研究のようです。このような記事もあります： Humans beat DeepMind AI in creating algorithm to multiply numbers
これらの研究は、10月6日に福島研究室のslackに卒業生の高橋惇くんが解説してくれて、その最後に交換法でもアタックできるかもと指摘しています。さらに10月15日には当時D３だった長野くんがKauers-Moosbauer論文を教えてくれました。統計力学やモンテカルロ法を知っていて探索問題に帰着できるのであれば、もうちょっと賢くできるだろうという感覚は研究室で共有している感じです。強化学習もよいのですが、Alpha Tensorはある意味で汎用的なアプローチになっているので、モンテカルロ法や最適化の専門家がその気になれば勝てるだろうという楽観的な感じもあります。まあ、温度をちょっと上げればいいんじゃないのってことです。ちょうどそのころに2023年度の統合自然科学科の７学期実験（数理演習II）の課題公開の時期だったので、そこで示したのがこのページのトップ画像の「Nature誌に掲載された高速行列積アルゴリズムを打倒するプロジェクト」です。そして、そこに反応したのが荒井くんでした。
荒井くんの研究は、2023年4月からはじまりました。最初は敵の手の内を知っておこうということで強化学習の勉強をしながら、Flip graphでの探索を並行して始めました。強化学習の勉強のところから市川くんも合流しました。5月2日には3×3行列の23と4×4行列の47が追試できました。重要な点はこれらをノートPCでできたことです。先行研究では巨大なコンピューターで計算していましたが、荒井くんはノートPCだけで追いついたので、手応えを感じていました。ただ、5×5行列は手強く、壁がありました。それはKrauers-Moosbauerが直面した困難点を再認識したことでもあります。物理的な言い方だと、エントロピー的な障壁で、場合の数が多すぎて当たりを引けない状況です。ゴルフホール的な難しさという言い方もできるかと思います。そこでの改良ポイントは２つあって、Flip graphのアルゴリズムに手数を増やすPlus transitionと探索空間を制限するEdge constraintsと名付けた手続きを導入したことにあります。前者は温度を上げる感じ、後者はちょっと賢いコショウをふりかけた感じです。そして、6月15日に94を見つけました。この計算も、5^3=125回の通常方法を初期条件として、ノートPCで完了です。そこからちょっと時間をかけてしまいましたが、論点を整理してこのプレプリント公開に至りました。
荒井くんは大学生になってから最適化問題に興味を持ったそうで、交換法も当然？知っていました。ただ、今回は交換法を使っていないのです。そこはとても良いことだと思っています。必要なら使うけど、必要ないなら使わないのは当たり前ですね。この研究は、AI assisted scienceなのかもしれないです。少なくともAI motivatedは確かです。ただ、上の解説にあるようにAIが見つけた初期条件をつかって人間の考えたアルゴリズムで新しい方法を見つけたってことでウケる話にしたかもしれませんが、別にそんなストーリは必要ないってことです。まあ、Kaggle Masterの荒井くんがちょっと本気出したら、大型汎用機も強化学習もふっ飛ばしちゃいましたって感じです。
荒井くんによるオンラインセミナー＠統計物理と統計科学のセミナーが３月５日に行われます。上のような浮ついた話ではなくて、実際にテンソル分解での表現やらアルゴリズムの説明が聞けるかと思います。
セミナーの様子はこちらから。