Gooogleの人工知能研究所DeepMindが開発したAlphaシリーズAIは、人間を超えることが難しいと言われてきた囲碁において、世界チャンピオンを打ち負かす「AlphaGo」によって一躍有名になり、その後もタンパク質の構造を解析するAI「AlphaFold」に派生するなど、その活躍の場を広げている。もともとゲームをプレイするために訓練されたAlpha AIだが、他のタスクに取り組まされ、驚くべき能力を発揮しているのだ。

プログラミングの基礎的な授業を受けたことがある人なら、ソートアルゴリズム（順序のないリストを昇順または降順に並べるコード）を学んだことがあるだろう。ソートにはさまざまな方法があり、また、できるだけ効率的にソートを行う事が重要になると習ったはずだ。

ソートは非常に基本的なものなので、プログラミング言語の標準的なライブラリのほとんどにアルゴリズムが組み込まれている。そして、LLVMコンパイラで使用されているC++ライブラリの場合、そのコードは10年以上変わっていないのだ。

しかし、GoogleのAIグループ「DeepMind」は今回、人間のコード例でまず学習させなくても、極めて最適化されたアルゴリズムを開発できる強化学習ツール「AlphaDev」を開発した。ここにおいて、ゲームプレイを得意とするAlphaシリーズを用いるために、プログラミングをゲームとして扱うように設定して解決させたのだ。

解決したい問題を“ゲーム”と捉える

DeepMindは、とりわけ、ゲームの遊び方を自ら学ぶソフトウェアを開発したことで注目されている。チェス、囲碁、スタークラフトなど、さまざまなゲームを攻略し、その効果を実証している。詳細はゲームによって異なるが、自分でプレイすることで学習し、スコアを最大化するための選択肢を発見していく。

このAIは、人間がゲームプレイの様子で学習していないため、人間が思いつかないようなゲームへのアプローチを発見することができるのだ。もちろん、常に自分自身と対戦しているため、人間が利用できる盲点をついているケースもある。

このアプローチは、プログラミングにも大いに関係する。大規模言語モデルが効果的なコードを書くのは、人間の事例をたくさん見てきたからだ。しかし、そのため、人間がこれまでにやったことのないようなことを開発することはまずないだろう。ソート関数のような、よく理解されているアルゴリズムを最適化するのであれば、既存の人間のコードをベースにしても、せいぜい同等のパフォーマンスを得ることができる程度だ。しかし、AIに真に新しいアプローチを発見させるには、どうすればよいのだろうか？

DeepMindは、チェスや囲碁と同じようなアプローチを取り、コードの最適化をゲームにした。AlphaDevシステムは、コードのレイテンシをスコアとして扱い、コードがエラーなく完了まで実行されるようにしながら、そのスコアを最小化しようとするx86アセンブリアルゴリズムを開発した。強化学習によって、AlphaDevは次第にタイトで高効率なコードを書く能力を身に付けていく。

AlphaDevの仕組み

システムがレイテンシを最適化すると言っても、それがどのように動作するかを説明するのとは全く違う。他の多くの複雑なAIシステムと同様に、AlphaDevもいくつかの異なるコンポーネントで構成されている。その1つが表現機能で、開発中のコードの全体的な性能を追跡する。これには、アルゴリズムの一般的な構造だけでなく、x86レジスタやメモリの使用も含まれる。

このシステムは、モンテカルロ法によるツリー検索によって、アセンブリ命令を個別に追加していく。この手法の「ツリー」という側面は、システムが膨大な数の潜在的な命令の中から限られた領域を素早く絞り込むことを可能にし、モンテカルロはその分岐から選ばれる正確な命令にある程度のランダム性を持たせている。(ここでいう「命令」には、有効で完全なアセンブリを作成するために選択される特定のレジスタなどが含まれることに注意)

そして、システムはアセンブリコードの状態をレイテンシと有効性で評価し、前回のスコアと比較しながらスコアを割り当てる。さらに、強化学習によって、プログラムの状態に応じて、異なる枝を下っていくことがどのように機能するかという情報を保持する。やがて、「ソート完了」というゲームの勝利条件を、最大スコア、つまり最小レイテンシで達成する方法を「学習」する。

このシステムの最大の利点は、学習する際にコード例が不要であることです。このシステムの最大の利点は、学習用のコード例を用意する必要がなく、自らコード例を生成し、それを評価することである。その過程で、ソートに効果的な命令の組み合わせの情報が蓄積される。

便利なコード

複雑なプログラムにおけるソートは、大規模で任意のアイテムのコレクションを扱うことが出来る。しかし、標準ライブラリのレベルでは、1つまたは少数の状況を処理する高度に特殊な関数の大規模なコレクションから構築されている。例えば、3項目、4項目、5項目のソートにはそれぞれ別のアルゴリズムがあり、さらに、任意の項目数を制限値まで扱える関数もある（つまり、4項目までならソートできるが、それ以上は扱えないということだ）。

DeepMindはこれらの関数のそれぞれにAlphaDevを設定したが、その動作は大きく異なる。特定の項目数を扱う関数では、変数の状態に応じて異なるコードを実行する分岐のないコードを書くことが可能だ。その結果、このコードの性能は、一般に必要な命令数に応じてスケールする。AlphaDevは、sort-3、sort-5、sort-8でそれぞれ1命令、sort-6とsort-7でさらに1命令を削減することができた。ただ1つ（sort-4）だけ、人間のコードを改善する方法を見つけられなかった。このコードを実際のシステムで繰り返し実行したところ、少ない命令数で性能が向上することが分かった。

可変数のエントリーをソートする場合、コードに分岐が生じるが、プロセッサーによって分岐を処理するためのハードウェアの容量が異なるのだ。そこで、100台の異なるマシンで動作させた場合の性能を評価した。ここでもAlphaDevは、さらなる性能を引き出す方法を発見した。ここでは、最大4つの項目をソートする関数について、その方法を紹介する。

C++ライブラリの既存の実装では、コードは一連のテストを行い、ソートする必要があるアイテムの数を確認し、その数だけ専用のソート関数を呼び出す。修正後のコードでは、もっと奇妙なことをする。項目が2つあるかどうかをテストし、必要に応じて別の関数を呼び出してソートするのだ。項目が2つより多い場合は、最初の3つの項目を並べ替えるために呼び出される。3つの項目がある場合は、そのソートの結果を返す。

しかし、ソートする項目が4つある場合は、3つのソートされた項目のセット内の適切な場所に4つ目の項目を挿入するのに非常に効率的な特殊コードを実行する。これは奇妙なアプローチに聞こえるが、既存のコードを常に凌駕していたのだ。

製品化にあたって

AlphaDevはより効率的なコードを作成したため、チームはこれらをLLVM標準C++ライブラリに再び組み込もうと考えた。ここで問題になったのは、そのコードがC++ではなくアセンブリであったことだ。そこで、同じアセンブリを生成するC++のコードを逆算する必要があった。それが終わると、このコードはLLVMツールチェーンに組み込まれた。

その結果、AlphaDevのコードは1日に何兆回となく実行されるようになった、と研究者は推測している。

論文

• Nature: Faster sorting algorithms discovered using deep reinforcement learning

参考文献

• Google DeepMind: AlphaDev discovers faster sorting algorithms

研究の要旨

ソートやハッシュなどの基本的なアルゴリズムは、一日に何兆回も使用されている。計算機に対する要求が高まるにつれ、これらのアルゴリズムが可能な限り高性能であることが重要になってきている。過去には目覚ましい進歩があったが、これらのルーチンの効率をさらに向上させることは、人間の科学者と計算機によるアプローチの両方にとって困難であることが判明した。ここでは、人工知能が未知のルーチンを発見することで、現在の技術水準を超えることができることを示す。これを実現するために、我々は、より優れたソートルーチンを見つけるというタスクをシングルプレイヤーゲームとして定式化した。そして、このゲームをプレイするために、新しい深層強化学習エージェントであるAlphaDevを訓練した。AlphaDevは、これまで知られていた人間のベンチマークを凌駕する小さなソートアルゴリズムをゼロから発見した。これらのアルゴリズムは、LLVM標準のC++ソートライブラリに統合された。ソートライブラリのこの部分への変更は、強化学習を使って自動的に発見されたアルゴリズムへのコンポーネントの置き換えを意味する。また、本アプローチの汎用性を示すため、他のドメインでの結果も紹介する。