GPTのシステム開発を行っている渋谷のLEAGENCEです。今回はAIが強化学習によってアルゴリズム開発を行うGoogle DeepMind「AlphaDev」について解説したいと思います。

記事はこちら

https://www.nature.com/articles/s41586-023-06004-9

エンジニア向けの記事になります

簡単に言うと、「アルゴリズムを強化学習によりブラッシュアップするアルゴリズム」です。

「AlphaDev」の概要1

（元記事の日本語訳）

基本的なアルゴリズム（例えば、ソートやハッシュ関数）は、毎日何兆回も使用されています。計算需要が増す中で、これらのアルゴリズムができるだけ高性能であることが重要となっています。これまでにも大きな進歩がありましたが、これらの機能をさらに向上させることは、人間の科学者や計算アプローチにとって難しい課題でした。この研究では、これまで知られていないアルゴリズムを発見することで、人工知能が現状を超える方法を示しています。

具体的には、より良いソートアルゴリズムを見つけるタスクを1人用ゲームとして定式化しました。そして、新しい深層強化学習エージェントであるAlphaDevを、このゲームをプレイするように訓練しました。AlphaDevは、これまでの人間が設定した基準を上回る、ゼロから生成された小さなソートアルゴリズムを発見しました。これらのアルゴリズムは、LLVM標準C++ソートライブラリに統合されました。この変更は、強化学習を使用して自動的に発見されたアルゴリズムで部品を交換したものです。また、他のドメインでの結果も提示し、アプローチの汎用性を示しています。

つまり、この研究では、人工知能がゲームをプレイさせて新しいソートアルゴリズムを発見することで、従来のアルゴリズムよりも効率的なものを開発し、実際のライブラリに統合することができました。この手法は、一般的であるため、他の分野やアルゴリズムにも適用可能であることを示唆しています。

この研究では、人工知能（AI）がゲームを使って新しいソート（並べ替え）アルゴリズムを発見し、それが従来のアルゴリズムよりも効率的であることがわかったということです。

この手法は汎用性が高く、他の分野にも応用できる可能性があります。つまり、AIが新しいアルゴリズムを発見することで、より効率的なコンピュータ処理や、仕組みを考えることができます。

「AlphaDev」の概要2

（元記事の日本語訳）

人間の直感やノウハウは、アルゴリズムの改善に重要な役割を果たしてきましたが、多くのアルゴリズムは人間の専門家でさえこれ以上最適化できない段階に達しており、計算上のボトルネックが増大しています。古典的なプログラム合成分野の研究では、正確なプログラムの生成や、レイテンシ（遅延）に関連する指標を使ったプログラムの最適化を目指しています。

ディープ・リフォースメント・ラーニング（DRL）を使うことで、CPU命令レベルで実際のレイテンシを最適化しながら、正確で高性能なアルゴリズムを生成することが可能になります。
また、正確で高速なプログラムの空間を効率よく探索し、考慮することができます。

コンピューターサイエンスの基本的な問題の1つは、シーケンスのソートです。これは、世界中の初等コンピュータサイエンスの授業で教えられており、さまざまなアプリケーションで幅広く使用されています。数十年にわたるコンピュータサイエンスの研究は、ソートアルゴリズムの発見と最適化に注力してきました。実用的なソリューションの重要な要素は、短い要素シーケンスに対する小さなソートで、これが大きな配列を分割統治アプローチでソートする際に繰り返し使用されます。この研究では、2つのタイプの小さなソートアルゴリズムに焦点を当てています。(1) 固定ソートと (2) 可変ソートです。固定ソートアルゴリズムは固定長のシーケンスをソートし、可変ソートアルゴリズムはさまざまなサイズのシーケンスをソートできます。

新しい効率的なソートアルゴリズムの発見の問題を、AssemblyGame と呼ばれるシングルプレイヤーゲームとして定式化しました。このゲームでは、プレイヤーは低レベルのCPU命令（アセンブリ命令）を選択して、新しく効率的なソートアルゴリズムを作成します。これは、プレイヤーがアセンブリ命令の組み合わせの空間を検討し、証明可能で速いアルゴリズムを生成する必要があるため難しい課題です。AssemblyGameの難しさは、検索空間のサイズだけでなく、報酬関数の性質にも起因しています。AssemblyGameでは、1つの誤った命令がアルゴリズム全体を無効にする可能性があり、この空間での探索が非常に困難です。

■解説
・つまり、 AIにゲームをプレイさせることによって強化学習を行われるアルゴリズム
・トランスフォーマーモデルがアセンブリ言語を改善して最適化する
・簡単に言うと、従来の深層強化学習に対して、入力変数を自然言語（アセンブリ言語）にしたものである
・応用可能性としては、Python等のプログラム、プロンプト、その他自然言語を利用ししたものなど多岐にわたる。

アルゴリズムを低レベルCPU命令として表現する方法

（元記事の日本語訳）

例えば、C++のような高レベル言語からマシンコードに変換する際には、まずアルゴリズムがアセンブリに変換されます（図1b参照）。その後、アセンブラがアセンブリプログラムを実行可能なマシンコードに変換します。この研究では、アセンブリレベルでアルゴリズムを最適化しています。

典型的なアセンブリプログラムでは、メモリからレジスタへ値がコピーされ、レジスタ間で操作され、そしてメモリに書き戻されます。サポートされるアセンブリ命令のセットは、プロセッサアーキテクチャに依存します。この研究では、AT＆T構文を使用したx86プロセッサアーキテクチャでサポートされるアセンブリ命令のサブセットに焦点を当てています。各命令はOpcode⟨OperandA, OperandB⟩という形式です。命令の例としてmovがあります。これは、ソース（A）から宛先（B）に値を移動することを意味します。また、比較（cmp）、条件付き移動（cmovX）、ジャンプ（jX）などの命令定義は、Extended Data Table 1で見つけることができます。

図1bの例では、%eax、%ecx、%edx、%ediは、4つの異なるレジスタ位置に対応しており、(%rsi)、4(%rsi)は、2つの異なるメモリ位置に対応しています。$2は、この例ではベクトルの長さに対応する定数値のプレースホルダです。この研究では、「アセンブリプログラム」と「アセンブリアルゴリズム」を同じ意味として使っています。これは、AlphaDevがAssemblyGameをプレイする度に、最初は順不同の命令セットからスタートしてアセンブリプログラムをゼロから構築し、新しく効率的なアルゴリズムを定義しているからです。

図1：C++とアセンブリプログラムの関係
figure 1
a, 2つまでの要素を持つ任意の入力シーケンスをソートできる変数のsort 2関数をC++で実装したもの。b, aのC++実装が、この低レベルアセンブリ表現にコンパイルされます。

→解説省略

「AlphaDev」のアルゴリズムについて

（元記事の日本語訳）

DRL（深層強化学習）を用いて高速なアルゴリズムを発見する方法について説明します。この問題では、環境を一人用ゲームとしてモデル化し、AssemblyGameと呼びます。ゲームの各状態はベクトルSt = ⟨Pt, Zt⟩で定義され、Ptはゲームで生成されたアルゴリズムの表現、Ztは実行中のアルゴリズムが事前に定義された入力セットに対して実行された後のメモリとレジスタの状態を表します。

例えば、図2aに示すように、時刻tで、プレイヤーは現在の状態Stを受け取り、アクションatを実行します。これにより、合法的なアセンブリ命令（例：mov）がこれまで生成したアルゴリズムに追加されます。正確さとレイテンシの両方を含む報酬rtが与えられます。アルゴリズムの正確さ（図2b）は、現在のアルゴリズムPtにN個のテストシーケンスを入力してN個の出力を生成することで計算されます。これらの出力は期待される出力と比較され、正確さ報酬rtが計算されます。レイテンシ報酬は、（1）アルゴリズムの長さを増やすことでエージェントを罰する（長さとレイテンシが高度に相関している場合）、あるいは（2）アルゴリズムの実際のレイテンシを測定することで生成できます。ゲームは限られたステップ数で実行され、その後ゲームが終了します。アセンブリ命令を使って正確でレイテンシの低いアルゴリズムを生成することがゲームの勝ちです。不正確なアルゴリズムや正確だけど効率が悪いアルゴリズムを生成するとゲームに負けます。

この一人用ゲームをプレイするエージェントをAlphaDevと呼びます。エージェントの主要な学習アルゴリズムは、AlphaZeroエージェント32の拡張版であり、深層ニューラルネットワーク33,38を使用してモンテカルロツリー探索（MCTS）プランニング手順を案内します。ニューラルネットワークへの入力は状態Stであり、出力は方針と価値の予測です。方針予測はアクションの分布であり、価値関数は現在の状態Stから受け取るべき累積リターンRの予測です。ゲーム中にエージェントは現在の状態Stを入力として受け取ります。次に、エージェントはMCTS手順を実行し、これを使用して次のアクションを選択します。生成されたゲームはネットワークのパラメータを更新し、エージェントが学習できるようにしています。

AlphaDevは、複雑なアルゴリズム構造を効率的に探索するために、命令の空間を効果的に表現できる表現39,40が必要です。このために、AlphaDev表現ネットワーク（Extended Data Fig. 1a）を導入しました。このネットワークは、（1）アルゴリズム構造の表現を提供するトランスフォーマエンコーダネットワーク、および（2）アルゴリズムがメモリとレジスタのダイナミクスにどのように影響するかを予測するために役立つCPUの状態エンコーダネットワークの2つのコンポーネントで構成されています。CPU状態エンコーダネットワークは、入力セットごとの各レジスタおよびメモリ位置の状態を入力として受け取るマルチレイヤーパーセプトロンで構成されています。これらのネットワークは、それぞれエンベッディングを出力し、そしてこれらの出力を結合してAlphaDev状態表現を生成します。

したがって、AlphaDevは複雑なアルゴリズム構造を効果的に学習し、高速なアルゴリズムを発見することができます。

AlphaDevのアルゴリズム解説

・St = ⟨Pt, Zt⟩について
Ptは[[0.54, 0.33, ・・・][・・・]]といったアセンブリ言語で書いたコードをベクトル表現したものであり、
ZtはPtのアルゴリズムを実行した際のメモリとレジスタの状態をベクトル表現にしたもの。
Stは、PtとZtのペアを表現するため、Stの形式は[[ Ptの要素 ], [ Ztの要素 ]]という2つのリストまたは配列が含まれたネストされた構造になります。
Stはプログラムとその実行によって得られるメモリとレジスタの状態。
・rt（報酬）について
rtは正確性とレイテンシ。
・Ztは、アルゴリズムPtが実行された後のメモリとレジスタの状態を表す指標であり、アルゴリズムの正確性を評価するために使用される。
・プレイヤーがStを受け取り、そのStの中にあるPt(アセンブリ言語のプログラム)に対して新たなコードを追記する。
・そしてその追記から得られた正確性とレジスタrtが与えられるということ
・記事の中で詳細な合格基準については触れられていないが恐らく既存のソートアルゴリズムを合格基準としてそれを上回った場合にゲームに勝ったというふうに定義している
・合格基準を上回るまでwhileする
・方針予測と価値予測とは囲碁で表すと方針予測が”次の手”であり価値予測が”次の盤面の勝率”にあたる
・方針予測はトランスフォーマーエンコーダーによって行われる
・深層強化学習モデルで学習する

「alphadev」の重要なポイント

・出力予測（選択肢の作成）をトランスフォーマーモデルで行うことによってその後のアルゴリズムは囲碁のalphagoのアルゴリズムとほぼほぼ同じようなアルゴリズムになるということです。
・つまり、複雑なアルゴリズムの構造を効果的に捉えるトランスフォーマーモデルの能力を活用しながら、選択肢が限られた変数（命令や手）を評価・選択する深層強化学習のタスクに適用しているということです。

①トランスフォーマーモデルでアセンブリ言語の改善案を予測し、それらの選択肢をリストアップ。
②それらの選択肢を数値（1、2、3、4など）に変換。この数値は、改善案それぞれに対応する
③深層学習モデルによって生成された方策（policy prediction）に基づいて選択。

方策予測について説明

深層学習モデル（AlphaDev）の方策予測（policy prediction）とは、現在の状態から次に実行すべきアクション（改善案）の確率的な分布を表すものです。具体的には、アセンブリ言語の各命令に対して、それぞれ実行される確率が割り当てられます。

例えば、あるアセンブリ言語のアルゴリズムを最適化したいとします。AlphaDevが選択肢として以下のアセンブリ命令を提示しました。

1. mov A, B
2. add A, C
3. sub D, A
4. jmp E

AlphaDevの方策予測に基づくと、それぞれの選択肢が選ばれる確率は次のようになります。

1. mov A, B: 0.7 (70%)
2. add A, C: 0.1 (10%)
3. sub D, A: 0.1 (10%)
4. jmp E: 0.1 (10%)

この場合、”mov A, B”の命令が最も高い確率（70％）で選ばれるため、AlphaDevはこの改善策を選択することが最も適切であると判断します。そして、この改善策をアセンブリアルゴリズムに適用して最適化を進めます。

方策予測は、アルゴリズムの現在の状態（アセンブリ命令の並びやレジスタやメモリの状態など）と、学習済みの深層ニューラルネットワークの知識を基にしています。最適なアクションを選択することで、アセンブリアルゴリズムが最適化されていきます。

以上、解説終了。