HireRoo TechBlog

HireRooの技術ブログ

AlphaCodeの概要を理解する

本ブログでは、DeepMind社から発表された、自動でプログラミングコンテストの問題を解くAlphaCodeと呼ばれるシステムの仕組みを説明していきます。調査は、先月2月19日に公開されたPreprintおよび公式の解説記事をもとに行いました。

AlphaCodeはシステムの名称であり、その中には細かな工夫が大量に散りばめられています。それらを全て紹介していくと記事が長くなり過ぎてしまうため、本ブログでは全体像をざっくりと理解できるようになるところをゴールに設定して解説させていただければと思います。

システム概要

AlphaCodeとは、競技プログラミングのような形式で問題を入力すると、その問題の回答を出力してくれるようなシステムの名称です。

f:id:KKosukeee:20220330192651p:plain
図1: AlphaCodeのシステム全体像

図1は、AlphaCodeのシステムの全体像です。この図をみてみると、1. データ取得 (DATA), 2. モデル訓練 (LEARNING), 3. 回答出力 (SAMPLING&EVALUATION)という大きく3つの区画に分けてシステムが構成されていることがわかります。 ざっくりいうと、GitHubやプログラミングコンテストのサイトからデータを収集し機械学習モデルを学習させ、モデル出力に対してフィルタリングを行なうことで最終的な出力を得ているようです。

データ取得

GitHubおよびプログラミングコンテストのサイトからデータは獲得されています。GitHubのソースコードは競技プログラミングに限らないものであり、後述するモデル訓練の際においてはPre-trainingと呼ばれる事前準備に使用されます。 プログラミングコンテストのサイトのデータは、モデル訓練の際にプログラミングコンテストを解くことに特化させるための学習であるFine-tuningに使用されます。

獲得したデータはそのまま使用されているわけではなく、前処理を行うことで、機械学習モデルの学習難易度を下げる工夫が取り入れられているようです。例えば競技プログラミングにおいて問題を解く速度は重要なので、参加者の中には少しでもコーディング量を削減するためにマクロを使用する人が多く存在します。例えば、よく使われるマクロの例は以下のようなものがあります。

#define rep(i,n) for(int i=0;i<n;i++)
typedef vector<pair<ll,ll» vpll;

このようにマクロが多用されたデータが収集されるわけなのですが、マクロにより置き換えが大量に発生しているデータをそのまま使用してしまうと機械学習モデルの学習の難易度が高くなってしまうことは想像できるのではないかと思います。そこで、学習前に前処理として、マクロによる置換をあらかじめ元に戻しておく処理を行なったりしているそうです。また、スペースや改行の位置などがユーザーごとに違う場合もありますので、前もって全てのソースコードをフォーマッタにかけて整形した上で学習するようにしたりすることで、性能が向上するそうです。

モデル訓練

訓練の際には、Pre-trainingおよびFine-tuningと呼ばれる2種類の学習が行われています。Pre-trainingにおいては競技プログラミングに限らずプログラミング言語そのものの特性(文法など)を把握させるための学習がなされ、Fine-tuningにおいてプログラミングコンテストを解くことに特化させるための学習が行われます。

1. Pre-training Pre-trainingの目的は、競技プログラミングに限らずプログラミング言語そのものの特性(文法など)を把握することです。そのためここでは、競技プログラミングの問題文や回答例などは用いず、GitHubに存在するソースコードのみを使用した学習が行われています。

今回は事前学習を行う上でのタスク設定のひとつとして、ソースコードの中身を一部マスクして、そのマスクされている部分が何なのかを予測するmask predictionというものが行われているようです。

f:id:KKosukeee:20220330192851p:plain
図2: Pre-trainingの手法例

後述するようにFine-tuningにおいてはnext-token predictionを元にした仕組みが使用されるのでmask predictionは冗長な学習のようにも思えてくるのですが、Fine-tuning後のモデルの精度の観点では、mask predictionによるPre-trainingは必須だったそうです。

このようにソースコードのみで完結した簡易的なタスクをまず解かせて機械学習モデルに全体的な文法などの情報を教え込んでいきます。繰り返しにはなりますが、この時点では競技プログラミングのソースコードや問題文の情報に関しては使用されていません。

2. Fine-tuning AlphaCodeで使用されているモデルは、何かしらの文章を途中まで入力すると、その文章の続きを出力してくれるような機能を持つ機械学習モデルとして解釈できます。

f:id:KKosukeee:20220330192929p:plain
図3: 再帰的な出力を行うモデルの挙動例

図3では、「現在、私は機械学習モデルの」という書き出しから次にくる単語を予測しています。次にくる単語が「訓練」であると予測されたのであれば、その次はそこまで含めて次の単語を予測していきます。イメージとしては、スマホなどの入力補完を無限に繰り返していくようなものに近いと思います。

学習用データとして競技プログラミングの問題文とその回答のコードがセットで存在していれば、問題文を入力したらそこからソースコードを書き上げていってくれるようなモデルが作れるというわけです。

f:id:KKosukeee:20220330193024p:plain
図4: プログラミングコンテストの問題を解くイメージ

このような機能を、Transformerベースのモデルを用いて訓練することで実現しています。

回答出力

プログラミングコンテストにおいては参加者がソースコードを何度も提出することが可能であることが多いです。なお、不正解のソースコードを提出するとペナルティがつくことも多いので、今回の問題設定としては、10回の提出の中に正解となるコードが含まれているかどうかをAlphaCodeは机上評価の指標として考えています。

上述したコンテストの仕組み上、機械学習モデルに複数個の回答候補を出力してもらうことは重要です。そこで、出力される言語(C++, Python)の切り替えや入力問題文の書き換えなどにより、出力候補を問題ごとに大量(>100万個/問題)に獲得するような戦略をとっています。

f:id:KKosukeee:20220330193108p:plain
図5: 提出サンプルの選定方法

次は、このようにして得られた大量の回答候補の中から、提出する回答を選んでいく必要があります。その際にも、競技プログラミングの仕組みを利用した工夫が取り入れられています。例えば、競技プログラミングにおいては問題に対して簡易的なテストケース(入力と出力の例)が提供されることが多いので、そのテストケースに通過するかどうかでのフィルタリングが全ての生成例に対してが行われています。なお、この時点で、99%の生成結果はテストケースを超えられずフィルタリングされてしまうそうです。

テストケースを用いたフィルタリングを行っても、問題ごとに数千、数万のプログラム候補が残る可能性があります。この中からランダムに10件を選ぶと、記法は異なるが表現している内容的には等価なプログラムに対して提出枠を浪費してしまう可能性があります。では、どうすれば等価なプログラムになっているかどうかを判定することができるでしょうか?答えはシンプルでして、”入力に対して同じ結果を出力するかどうか”を調べられれば良いのです。複数の入力例(対応する出力例は不要)があればプログラムの挙動が等価かどうかを判定できるわけなのですが、AlphaCodeにおいてはそのような入力例も機械学習モデルを用いて別途生成しているようです。このようにして得られた複数個の挙動の異なる生成されたコードが、最終的なAlphaCodeの出力となります。なお提出の際には、出力が同じグループのサイズ順にランダムにひとつが選ばれ提出されていきます。

f:id:KKosukeee:20220330193134p:plain
図6: 出力候補数(横軸)と性能(縦軸)の関係

図6は、生成されたソースコード候補の数に応じて、最終的な正解率がどのように変化するかを示したような図です。元となる候補の数が多ければ多いほど、最終的な性能も高いような傾向が見て取れるかと思います。

結果

1. 性能

f:id:KKosukeee:20220330193215p:plain
図7: 10回のCodeforcesにおけるコンテストに参加した場合の推定順位(値が低いほど良い)

Codeforcesというプログラミングコンテストのプラットフォームにて、ペナルティ数や回答に要した時間などを踏まえ仮想的にコンテストに参加してレーティングを計算したところ、平均的には全ユーザーの平均程度(上位54.3%)の順位につくそうです。約半分のユーザーよりはレートが高いわけなので、インパクトのある結果なのかなと思います。

2. 計算効率

f:id:KKosukeee:20220330193252p:plain
図8: (左)訓練に要する時間および(右)出力候補数と性能(縦軸)の関係

図8は、訓練および回答の出力に要する時間を使用するモデル規模ごとにplotしたものです。 左図の訓練時間とのトレードオフを見てみると、TPUを1つ使えたとしても、10000日(=約27年)ほどの時間を要するということが見て取れます。TPUをたくさん持っていれば並列化による高速化は可能ではありますが、一般的な用途においては現実的ではありません。現状、自分たちでAlphaCodeと同様のものを独自につくるのは、かなり大きな工夫をしないと難しいのではないかと思われます。

また右図から見られる通り、回答を出力する時間も、回答候補の数が増えるにつれて伸びていきます。あまりに長すぎると、プログラミングコンテストの時間中には終わらないこともあり得ます。実際、最高精度のものは、TPUを1つだけ使えるような状況を仮定すると、ひとつの問題の回答を得るのに106秒(=約12日)かかるようです。

とはいえ技術も進歩していくので、このようなシステムが現実的に稼働する未来がくることも、案外遠くない可能性もあるのではないかなと思います。 図を見た限りではモデルを大きくしたり出力候補数を増やすことにより性能が向上するトレンドはまだサチっていないように見えるので、処理を高速化していくことができれば、それはそのままシステムの性能向上にも直結していくのではないかなと思います。

まとめ

いかがだったでしょうか?AlphaCodeは世界的にも非常に注目されているプロジェクトですので、少しでも具体的な内容が本ブログにより理解できたのであれば幸いです。 ソースコードの生成技術は、未来のエンジニアの形にも影響を与えていくかもしれません。発展の著しい分野ですので、今後も引き続き注目していけたらなと思います。