AHC延長戦の伝説"saitodevel01"の機械学習手法を再現して1位を奪い返した話

Posted at 2026-08-18

はじめに

以前、Introduction to Heuristics Contestの延長戦1位解法を解説する記事を書きました。ところが、その後にsaitodevel01というユーザがさらに高いスコアを出し、筆者の提出は2位になりました。

Introduction to Heuristics Contestでは、コンテスト1位の解法でもwriterによる解説でも、焼きなまし法が主軸となっています。そのため、延長戦上位提出コードの多くは焼きなまし法を主軸としたものが占めています。一方で、saitodevel01氏のコードは他の上位提出コードと異なり、ニューラルネットワーク（以降NNと表記）による行動候補の評価を主軸としています。saitodevel01氏はAHC015など、Introduction to Heuristics Contest以外でもNNを用いて延長戦1位をとっており、個人的に注目していました。saitodevel01氏の活躍についてはAtCoder World Tour Finals 2026 Heuristic Japanese streamでも触れられています。

saitodevel01氏のコードは、学習済み重みを提出コード容量制限に収めるための圧縮や、ライブラリリンクのための工夫によって、かなり解読が難しいです。とはいえ、がんばればモデルに入力している特徴量や、出力をどう利用しているかといった構造は解読できます。しかし、学習済みの重みがあっても、その重みをどう学習したかまでは提出コードだけでは判断できません。

この、AHCではあまりみかけない面白い方法を再現したいと考え、以下3点の達成を目標に再現実験をしました。

再現しやすい(提出スコアに含まれない学習方針なども、サンプルコードで補う)
わかりやすい（学習部分以外の余計な部分をなるべくそぎ落とす）
強い（提出スコアで上回る）

最終的には、独自に作った65,536入力の教師データからモデルを学習し、読みやすさを優先したC++のビームサーチへ載せました。延長戦順位表では、saitodevel01氏の提出 128,280,488 に対し、128,286,311 を記録して1位を更新できました。

saitodevel01氏のコードの構造

まずsaitodevel01氏のコードの特徴として、提出言語がPythonであることが挙げられます。これは、探索本体までPythonで書かれていることを意味しません。Pythonの提出ソースには、文字列に圧縮されたC++コードを復元し、利用できる状態にして起動する外側の役割を担当します。NNの推論、ビームサーチ、後処理の本体は、圧縮文字列から復元されるC++部分で行われます。

Python部分と、その中に格納されたC++部分の役割は次のように分かれています。

処理	Python部分	圧縮文字列から復元されるC++部分
C++部分の格納と起動	圧縮文字列を提出ソース内に保持し、実行時に復元してC++部分を起動します。	Python部分から起動され、問題入力を受け取ります。
モデルの準備	C++部分を利用できる状態にします。	encoderとdecoderのAOTI packageを復元し、LibTorchから読み込みます。
入力の準備	C++部分へ標準入出力を引き継ぎます。	問題入力と、それまでに選んだ行動をtensorへ変換します。
推論	C++部分の起動後は、探索本体をそちらへ任せます。	`AOTIModelPackageLoader`からencoderとdecoderを呼び出します。
探索と後処理	C++部分の実行を制御します。	ビームを更新して行動列を作り、残り時間で後処理を行います。

C++部分の処理を整理すると、全体は「NNで候補へ優先順位を付けるビームサーチ」と「残り時間を使う後処理」の2段構成です。

NNは、365日分の行動列や最終スコアを一度に出すものではありません。encoderは問題入力全体を最初に1回だけ読み、各日・各コンテストの特徴を作ります。decoderは、行動を決める日と同じ位置のencoder特徴、直前に選んだコンテスト、GRUのhidden stateを受け取り、次に選ぶ26種類のlogitと次のhidden stateを返します。このencoder特徴は前後の日の入力からも影響を受けますが、日ごとに異なる特徴です。提出コードでは、このlogitを「現在までの選択を踏まえ、次の1手としてどれが有望か」を表す値としてビームサーチに利用しています。

推論から出力までの流れは次のようになります。

ビームの途中ではmodelの累積log確率を基準に候補を残しますが、最後はmodelの確率ではなく、問題本来のraw scoreで最良の1本を選びます。そのため、NNは解を直接決めるのではなく、限られたビーム幅でどの行動列を残すかを案内する役割です。

ビームサーチの後も、時間が余れば大近傍法で最後の仕上げをしています。現在の割り当ての一部を乱数で固定し、固定しなかった部分を最小費用流でまとめて再配置する処理を制限時間まで繰り返しています。最小費用流は、大近傍法で作り直す部分を解くために使われています。筆者が以前の記事で紹介した大近傍法を用いた焼きなましとは異なり、saitodevel01氏の後処理には温度に応じた確率的な受理はありません。

ここまでで、PythonとC++の役割、modelの評価でビームサーチをする仕組みまで分かりました。次に、そのmodelへ何を入力し、内部でどのように情報を混ぜているのかを見ていきます。

入力特徴とモデル

modelへ渡す情報は、問題入力と、それまでに選んだ行動だけです。以下では問題文と同じく、日を $d$、コンテストタイプを $i$、開催による満足度の増加量を $s_{d,i}$、満足度の下がりやすさを $c_i$、$d$ 日目に選ぶタイプを $t_d$ と書きます。

特徴	shape	内容
満足度の下がりやすさ $c_i$	26	0〜100をembeddingします。
満足度の増加量 $s_{d,i}$	$D \times 26$	141で割った商と余りを別々にembeddingします。model内部では $26 \times D$ に転置します。
日 $d$	$D$	$d-1$をembeddingします。
直前に選んだタイプ $t_{d-1}$	stepごとに26	one-hotを2次元へembeddingします。$d=1$では全要素を0にします。

embeddingは、整数ごとに学習可能な短いvectorを割り当てる表です。たとえば$c_i$には0〜100の101種類があるため、101行×32列の表を用意します。$c_i=37$なら37行目の32次元vectorを取り出し、その値をmodelの学習によって更新します。

入力特徴の中では、$s_{d,i}$の変換が特徴的です。$s_{d,i}$を連続値として正規化せず、141進法の商と余りに分けてembeddingします。具体的には、$s_{d,i}=141q_{d,i}+r_{d,i}$として、商$q_{d,i}$と余り$r_{d,i}$を別々のembeddingへ入れます。商は値の大まかな大きさ、余りはその中の細かな違いを表します。

0〜20,000をそのまま32次元のembeddingにすると、$20,001\times32=640,032$個のparameterが必要です。141進法なら、商は142種類、余りは141種類なので、必要なのは$(142+141)\times32=9,056$個です。元の値は商と余りの組から一意に決まり、parameter数は約71分の1になります。141は$\sqrt{20,001}$に近いため、商と余りの種類数もほぼ同じになります。

この構造にparameter数を抑える効果があることは確認できます。ただし、なぜこの表現を選んだのかというsaitodevel01氏本人の意図までは、提出コードだけでは分かりません。

直前に選んだタイプ$t_{d-1}$は、選んだタイプだけを1、残りを0にした26要素のone-hotで表します。各要素の0または1を2次元へembeddingするため、結果のshapeは$26\times2$になります。

4種類のembeddingを加算する処理は、日$d$とタイプ$i$の組ごとに行います。各組の特徴を$\mathbf{x}_{i,d}$とすると、次の32次元vectorです。

\mathbf{x}_{i,d}=E_q(q_{d,i})+E_r(r_{d,i})+E_c(c_i)+E_d(d-1)\in\mathbb{R}^{32}

$\mathbf{x}_{i,d}$自体は32次元ですが、$(d,i)$の組は$D\times26$個あります。それらを並べたEncoderへの入力全体のshapeが$26\times D\times32$です。$E_c(c_i)$は同じタイプの全$D$日で使い、$E_d(d-1)$は同じ日の全26タイプで使います。

model全体のlayer構造は次の通りです。

一般にEncoderは、入力された生データを、後続の処理が利用しやすい特徴表現へ変換する部分です。このmodelでは、$c_i$、$s_{d,i}$、日$d$をまとめ、各日・各コンテストタイプの特徴を作ります。この計算は行動列$t_1,\ldots,t_D$に依存しないため、ビームサーチを始める前に1回だけ実行できます。

Decoderは、Encoderが作った特徴と、それまでに出力した内容を使って、実際の出力を順番に決める部分です。このmodelでは、$d$日目位置の特徴、直前に選んだタイプ$t_{d-1}$、それまでの履歴を表すhidden stateから、次に選ぶ26種類のlogitを出します。

GRUはGated Recurrent Unitの略で、時系列を扱うrecurrent neural networkの一種です。過去から受け取ったhidden stateのうち、どの情報を残し、どの情報を新しい入力で更新するかをgateで調整します。このmodelでは、コンテストタイプごとに32次元のhidden stateを持ち、$d$日目から$d+1$日目へ引き継ぎます。そのため、Decoderへ直接渡す行動は$t_{d-1}$だけでも、それ以前の選択はhidden stateを通して次のlogitへ影響します。

このmodelのEncoderでは、4種類のembeddingを32次元で加算します。7段のMixing Blockは、まずTimeConvでコンテストごとに$D$日方向を混ぜ、次にTypeTransformerで日ごとに26種類方向を混ぜます。TimeConvはchannelを32から64へ広げ、kernel size 99のdepthwise convolutionを通して32へ戻す構成です。

したがって、Decoderへ渡す「$d$日目位置のEncoder特徴」は、$s_{d,i}$だけをembeddingしたものではありません。kernel size 99のTimeConvを7段重ねたときの理論上のreceptive fieldは$1+7\times(99-1)=687$日で、$D=365$日全体を覆います。このため、どの日の入力も各位置の特徴へ影響できます。

ただし、各位置が問題入力全体をそのまま保持するという意味ではありません。日$d$のembeddingと$s_{d,i}$が異なるため、Encoderは$D$日それぞれに異なる特徴を出力します。Decoderは行動を決める日$d$に対応した$26\times32$だけを、その都度取り出します。

全日程を1個の共通vectorへまとめる構造も設計できますが、その場合は日ごとの違いを1個のvectorへ圧縮し、Decoder側で必要な日を取り出す別の仕組みが必要です。このmodelはEncoderを1回実行して$D$日分をまとめて保存し、Decoderから日ごとに参照する構成です。

このmodelのDecoderでは、$d$日目位置のEncoder特徴32次元と、直前の行動を表すembedding 2次元を種類ごとに結合します。GRUは各種類のhidden stateを日方向へ引き継ぎます。その後、TypeTransformerを1段使って26種類間の関係をもう一度混ぜ、最後のlinear layerで各種類を1個のlogitへ変換します。学習パラメータ数は165,248です。

logitは、softmaxへ入れる前の正規化されていない点数です。26種類のlogitを$z_1,\ldots,z_{26}$とすると、softmaxは各種類の確率を次のように計算します。

$$
p_i=\frac{\exp(z_i)}{\sum_{j=1}^{26}\exp(z_j)}
$$

このmodelのDecoderはlogitまでを出力し、ビームサーチ側でlog_softmaxを適用して$\log p_i$へ変換します。行動列全体の確率は各日の確率の積になるため、logを取れば各日の$\log p_i$の和として計算できます。ビームサーチは、この累積log確率が高い行動列を残します。

一方、現在までの累積満足度や$\mathrm{last}(d,i)$を、正確な数値としてDecoderへ直接渡しているわけではありません。過去に選んだ$t_1,\ldots,t_{d-1}$は、毎日のone-hot入力を通してGRUのhidden stateへ圧縮されます。Decoderは、前後の日の入力から影響を受けながら日ごとの違いも残したEncoder特徴とhidden stateを組み合わせ、次の行動を選ぶために必要な履歴を近似的に表現します。このmodelの役割は現在のスコアを厳密に再計算することではなく、教師の行動列を学習し、次に選ぶタイプへlogitを付けることです。

ここまでの実行時の構造は、提出ソース内の圧縮文字列やmodel packageを展開することで確認できます。一方、そのmodelをどのように学習したのかは、学習済み重みから逆算できません。次に、提出コードだけでは分からなかった部分を整理します。

提出コードだけでは分からなかったこと

今回、saitodevel01氏のmodelについて利用したのは、公開された演算構造、parameter名、tensor shape、dtype、packageの入出力です。

一方、次の情報は分かりませんでした。

教師の行動列をまねる模倣学習、問題のスコアを報酬にする強化学習など、学習の大方針
学習入力をどのように集めたか
行動列を教師にした場合、どの探索で作ったか
1入力につき何候補を保存したか
どのloss、optimizer、augmentationを使ったか
何件、何epoch、何seedで学習したか
どの基準でcheckpointを選んだか

提出コードに含まれているのは、学習済みmodelを使って次の行動を選ぶ処理です。学習時のlossや更新処理は含まれていないため、完成した重みだけを見ても、模倣学習と強化学習のどちらを使ったかは判別できません。modelの構造だけでなく、学習の枠組みそのものが未知でした。

学習方法の工夫

最初に判断したのは、modelの学習方法です。今回は、時間をかければ高得点な行動列を作れる自分の探索を持っていたため、その出力を教師としてmodelにまねさせる模倣学習を選びました。

saitodevel01氏の学習済み重み、checkpoint、logit、行動列、モデル出力は、データ生成、学習、推論のいずれにも使っていません。公開されたモデル構造を固定した上で、学習可能な値を乱数初期化し、自分の教師だけから学習しました。

模倣学習には、modelがまねるための高品質な行動列が必要です。教師には、以前の記事で紹介した自分の提出コードを使いました。

この解法は焼きなましが主体です。2秒の提出では探索回数が限られますが、学習データ作成には時間制限がありません。そこで、提出時のおよそ100倍の計算時間を要する水準まで探索を増やし、「遅いが高品質な教師」として使いました。

実際にこの探索量でローカル実行し、提出の時間制限を無視すれば、saitodevel01氏の解法よりも高いスコアを出せることを確認しました。この結果から、その行動選択へmodelを十分に近づけ、2秒以内で再現できればsaitodevel01氏に勝てるという確信を持って進められました。

最終データの条件は次の通りです。

項目	値
問題入力数	65,536
教師候補数	1入力につき2
教師乱数seed	0、1
教師探索反復数	1 callにつき298,529,882
問題生成seed	1,000,000からの連番

greedyだけで使うmodelなら、各入力に教師を1本だけ用意し、その行動を正解として学習する形でも成立します。しかし、最終的にはビームサーチで使うため、最良の1手だけを当てるmodelにはしたくありませんでした。ビームサーチでは、2番手以降にも良い枝が残っていることが重要です。

そこで同じ問題入力をseed 0と1で2回探索し、2本の高得点な行動列を保存しました。2本が異なる分岐を含む場合も、片方を誤りとして押し下げず、教師スコアに応じた重みで両方を学習できます。

教師は時間ではなく反復数で停止します。同じsource、問題seed、教師seed、反復数なら同じ出力になるため、数日かかる生成を中断・再開しやすくなります。各sampleには入力と行動列だけでなく、raw score、seed、source hashも保存し、収集後に全件を再採点しました。

教師の実装は cpp/teacher.cpp、dataset生成はsrc/neural_schedule/generate_dataset.py です。

これで、未知だった教師データの生成方法を独自の方法で埋めました。次は、1入力に2本ずつ用意した行動列を、どのようにmodelへ学習させたかです。

2本の教師をどう学習するか

ここからは、もう1つの未知だった学習方法を説明します。

まず、2本の教師についてsequence cross entropyを別々に計算します。これは、365日それぞれで教師が選んだ行動の負のlog確率を足したもので、2本のlossを$L_0,L_1$とします。次に、教師のraw scoreを$s_0,s_1$として、各lossへ掛ける仮の重み$a_0,a_1$を次のsoftmaxで作ります。

$$
a_k=\frac{\exp(s_k/T)}{\exp(s_0/T)+\exp(s_1/T)},\qquad T=5{,}000
$$

temperatureの$T$は、教師のスコア差を重みの差へどの程度強く反映するかを調整する値です。$T$が小さいほど高得点な教師へ重みが集中し、大きいほど2本の重みが50％ずつに近づきます。$T=5{,}000$では、2本のスコア差が5,000なら、softmaxへ入る値の差が1になります。このとき、高得点な方の重みは約73.1％、低い方は約26.9％です。

このsoftmax重みをそのまま使わず、さらに20％の一様分布を混ぜます。候補が2本なので一様分布は$(0.5,0.5)$であり、最終的な重みは次のようになります。

$$
w_k=0.8a_k+0.2\times\frac{1}{2}=0.8a_k+0.1
$$

たとえばスコア差が5,000の場合は、softmaxで得た$(0.731,0.269)$が、最終的に約$(0.685,0.315)$になります。ここでいう「20％の一様分布」は、低い方の教師へ20％を与えるという意味ではありません。重み全体の20％を2本へ均等に配るため、どれほどスコア差が大きくても各候補には最低10％の重みが残ります。スコアが同じなら、最終的な重みも$(0.5,0.5)$です。

学習に使うlossは$w_0L_0+w_1L_1$です。高得点な行動列をより強く学習しながら、もう一方を完全には無視しません。これは、1本の正解だけを当てる分類器ではなく、複数の有望な枝をビームサーチへ供給するpolicyとして学習するためです。

最終学習条件は次の通りです。

項目	値
初期値	ランダム初期化
採用run	seed 0
epoch	120
batch size	4
step/epoch	16,384
総step	1,966,080
optimizer	AdamW
learning rate	0.0003
weight decay	0.0001
warmup	51,200 step
gradient clipping	1.0
augmentation	contest番号の置換、時間反転
計算	float32、決定的実行

checkpointはvalidation lossではなく、学習に使っていない固定128入力のgreedy raw scoreで選びました。最終的に採用したのはepoch 101、step 1,654,784です。

モデルは src/neural_schedule/model.py、学習器はsrc/neural_schedule/train.py にあります。

これで、提出コードだけでは分からなかった教師データの生成方法と学習方法を、どちらも独自の方法で埋められました。残るのは、得られたmodelを2秒以内の提出へ載せる工程です。

探索コードの整理

最終提出では、NNの推論だけでなく、その出力を使うビームサーチもC++で実装しました。ただし、saitodevel01氏の探索部をそのまま移すのではなく、AtCoder Heuristic First-step Vol.2 Trainingのサンプルコードに近い、処理の流れを追いやすい形へ書き直しています。

サンプルコードでは、1本の行動列をStateで表し、StateにisDone()、advance()、legalActions()を持たせています。汎用的なbeamSearch()は、問題固有の状態を直接操作せず、このinterfaceだけを使って候補の展開、評価順の並べ替え、ビーム幅までの絞り込みを行います。今回の提出コードでも、同じ関数名と役割分担を採用しました。

NNの推論だけは、サンプルコードと同じ方法では効率良く実行できません。同じ日に残っているビーム全体をまとめて推論するため、共有するSearchContextと、バッチ内の位置を表すbeam_index_を追加しています。Stateが公開するinterfaceとbeamSearch()の流れはサンプルに近いまま、modelの累積log確率が高い状態を残せるようにしています。

探索コードのインターフェース templates/submission.cpp.in

class State
{
public:
    bool isDone() const;
    void advance(const int action);
    vector<int> legalActions() const;
};

vector<int> beamSearch(const State& state, const int beam_width);

この構成は役割が分かりやすい一方、saitodevel01氏の提出に含まれる問題専用の探索実装ほど、状態管理や候補展開を速度のために突き詰めてはいません。同じ幅140ではAtCoderの実行時間制限に対する余裕を安定して確保できなかったため、最終版ではビーム幅を104まで減らしました。

さらに、saitodevel01氏の最終提出はビーム幅140の後に、大近傍法で後処理を行います。今回の最終版では、コードを読んだときにmodelとビームサーチの関係だけを追えるように、この後処理もすべて外しました。ビーム幅104が選んだ行動列を、そのまま出力します。

したがって、探索量だけを比べると、今回の提出はビーム幅が小さく、後処理もないため不利です。それでも延長戦順位表では、saitodevel01氏の128,280,488に対して128,286,311となり、5,823上回りました。探索を複雑にして差を埋めたのではなく、読みやすい探索へ載せても上回れるmodelを学習できたことが、今回の結果で重視している点です。そのため、探索部をさらに改善すれば、まだスコアを伸ばす余地はあります。

PythonのモデルをC++提出へ載せる

前述のとおり、saitodevel01氏はPython提出の中に圧縮したC++部分を格納していました。今回は、推論処理と探索を読みやすくするため、提出ソース自体をC++として構成しました。

この方式が最初に使われたsaitodevel01氏の2024年6月24日の提出は、Python（CPython 3.11.4）で行われました。当時のAtCoder公式の言語・ライブラリ一覧を見ると、C++のGCC 12.2にはLibTorchのinclude pathやlibrary pathがなく、compile commandにも-ltorch、-ltorch_cpu、-lc10は含まれていません。一方、Python（CPython 3.11.4）にはtorch 1.13.1が用意されていました。当時は通常のC++提出からTorchを利用するための設定が用意されていなかったため、saitodevel01氏は提出言語にPythonを利用していた可能性があります。

現在は状況が異なります。Introduction to Heuristics Contestのルールページで「C++23 (GCC 15.2.0)」の項を見ると、Torchのheaderを読むinclude pathとlibrary pathに加え、-ltorch、-ltorch_cpu、-lc10がcompile optionへ最初から入っています。現在の公式言語・ライブラリ一覧にも、C++で利用できるlibtorch 2.8.0と同じcompile commandが記載されています。そのため、提出者がcompile optionを変更できないAtCoder上でも、現在はTorchを使うコードをC++として提出できます。

今回の学習済みモデルはencoderとdecoderに分け、PyTorch 2.8.0のAOTI packageへ変換します。C++側では torch::inductor::AOTIModelPackageLoader から呼びます。ローカルcompileでは、PyTorchのinclude・library path、rpath、ABI指定に加えて、次のlibraryを明示的にlinkしました。

-ltorch -ltorch_cpu -lc10

AtCoderの言語環境は今後も変わり得るため、この方法を再利用するときは、ルールページに掲載されている対象言語のcompile commandを確認する必要があります。

AOTI packageはbinary fileなので、そのままC++のsourceへ貼り付けることはできません。さらに、AtCoderのルールページでは提出ソースの上限が512 KiB、つまり524,288バイトと定められています。encoderとdecoderを埋め込んだsourceを、この上限へ収める必要があります。

最初に生成したAOTI版は、base94を使っても約556 KiBあり、AtCoderへの提出時にsourceが長すぎるとして拒否されました。そこで、AOTI archiveに含まれる実行時には使わない生成C++ memberの中身を短い文字列へ置き換え、encoderとdecoderを連結してからXZで圧縮しました。AOTIのloaderがmember名を参照するため、file自体を消すのではなく、名前を残して中身だけを削っています。

圧縮後のbinaryをsourceへ埋め込むときは、一般的なbase64ではなくbase94を使いました。base64は3バイトを4文字で表すため、403,880バイトの最終圧縮データは538,508文字になります。この時点で512 KiB上限を14,220バイト超えており、decoderやビームサーチのC++コードを加える前から提出できません。一方、base94は、ASCIIの!から~までの94文字を使い、2文字へ13 bitまたは14 bitずつ詰めます。同じ403,880バイトを494,061文字で表せるため、base64より44,447文字短くなりました。最終的にAtCoderへ提出したsourceは511,065バイトで、上限までギリギリですが制約に収まりました。

C++実行時には、base94の文字列をbinaryへ戻してXZを展開し、encoderとdecoderのAOTI packageを一時ファイルへ復元します。

推論時にはencoderを最初に1回だけ実行します。decoderは365日分呼びますが、同じ日のビーム全体をまとめて1回のdecoder呼び出しへ渡しています。

checkpointからAOTI変換、圧縮、埋め込み、C++ compileまでをsrc/neural_schedule/export_submission.py の1コマンドで実行できます。最終提出を作った後に、C++を人手で修正する工程はありません。

AWSの利用

筆者のローカルPCのCPUはIntel Core i7-14700F、GPUはRTX 5060 Ti 16 GBです。

教師データを作るときは、学習方法の工夫でも説明した通り、1つの問題入力に対してseed 0と1で教師となる探索を2回実行し、2本の行動列を保存します。このため、教師探索の実行回数は入力数の2倍になります。ローカルでは単位時間あたりに完了できる教師探索の数を測り、14並列で段階的に入力数を増やして実験しました。

到達した入力数	新たに生成した入力数	教師探索の実行回数	生成時間
128	128	256	1時間55分
512	384	768	6時間28分
2,048	1,536	3,072	17時間35分
4,096	2,048	4,096	33時間47分
8,192	4,096	8,192	70時間59分
合計	8,192	16,384	130時間44分

各段階の生成時間を合計すると5日10時間44分です。この時間には、modelの学習、評価、gate、次の実験を始めるまでの待ち時間は含めていません。

データを増やす判断には、学習に使っていない512入力で、増加前後のmodelをビーム幅96で比較したローカル評価を使いました。学習seed 0、1、2の3本を対にして比べた結果は次の通りです。評価に使った512入力は段階ごとに異なるため、入力の難しさに左右される絶対値ではなく、同じ入力上で測った増加前後の差を示します。値は1入力あたりのraw scoreの平均差で、正の値が改善を表します。

教師データの入力数	直前の段階からの平均改善量	95％信頼区間
512 → 2,048	+42,586.0	+39,331.1 ～ +46,852.7
2,048 → 4,096	+10,803.7	+9,943.3 ～ +11,696.4
4,096 → 8,192	+6,124.8	+5,453.8 ～ +6,853.1
8,192 → 16,384	+4,040.6	+3,578.4 ～ +4,503.2

増加幅は小さくなっていましたが、すべての段階で3本の学習seedがそろって改善しました。ただし、8,192入力で学習したmodelは、saitodevel01氏と同じAOTI wrapperへ載せて別途比較しました。学習に使っていない256入力をビーム幅96・後処理なしで解いたところ、1入力あたりのraw score平均は5,655.7下回りました。勝てたのは256入力中28入力です。データを増やすほど改善している一方で、8,192入力ではまだ足りないことがローカルでも分かりました。

この結果から、さらにデータを増やす方針に踏み切りました。一方、8,192入力までをローカルで生成するだけでも約5日11時間かかっています。データ生成の速さを左右するのはGPUではなく、教師探索1回を処理するCPUの速さと、同時に実行できる探索数です。そこで、続きは多数の教師探索を並列に実行できるAWSへ移しました。AWSで16,384入力まで増やした後の評価でも改善が続いたため、最終的に65,536入力まで増やしました。128入力から512入力までと、16,384入力から65,536入力までについては、同じ条件で比較できるローカル評価が残っていないため表には含めていません。

今回の教師探索は、長時間のCPU計算中に外部と通信しないため、日本からの距離による遅延はほとんど影響しません。そこで、調査時のSpot価格とquotaを比較し、同じc7i.16xlargeのSpot価格中央値が東京より約26％安く、必要なSpot vCPU quotaを確保できた米国西部（オレゴン）リージョン（us-west-2）を選びました。

最終追加分では、このリージョンにAWS Batchのmanaged compute environmentを作り、c7i.16xlargeのSpot Instanceだけを使いました。On-Demandへのfallbackは設定せず、同時に利用する上限を256 vCPUとしました。1つの実行タスクには1 vCPUと2 GiBを割り当て、教師探索を8回ずつ直列に実行させています。

既存の16,384入力へ49,152入力を追加し、各入力についてseed 0と1で教師探索したため、新たな教師探索は98,304回です。1つの実行タスクが8回を処理するように12,288個へ分け、AWS Batchの上限に合わせて10,000タスクと2,288タスクの2組として投入しました。開始から収集と監査の完了までは約67時間14分でした。

AWS上では、固定した教師探索programとworkerをcontainer imageにしてECRへ保存し、問題入力と結果をS3へ保存しました。教師探索が1回終わるたびに行動列、実行記録、完了markerを保存し、再実行時にはhashが一致する完成済みの結果をskipします。Spotの中断や一部jobの失敗があっても、完了したデータを作り直さずに再開できるようにしました。収集時には全98,304回分のhashとscoreを監査し、49,152入力分のデータがそろったことを確認しています。

AWSはCPUを使う教師データの生成だけに利用しました。一方、教師データを読み、NNのparameterを更新する学習ではGPUが重要です。65,536入力をそろえた後のseed 0学習はローカルのRTX 5060 Ti 16 GBで行い、約31時間かかりました。AWS Batchを再現する場合は、対象regionで利用できるSpot InstanceとStandard Spot vCPU quotaを事前に確認し、ECR、S3、AWS Batch、CloudWatch Logsを利用できる最小限の権限、費用budgetとalert、Spot中断後に再開できる保存方式を用意する必要があります。

今回のAWS利用全体の実課金額は $311.26 です。これは最終49,152入力だけの料金ではなく、事前test、それ以前のデータ生成、失敗とresumeを含む今回のAWS利用全体の金額です。

なお、筆者のように途中からAWSを利用するのではなく、最初から65,536入力すべてを同じ構成で生成する場合、実行時のSpot単価上限で見積もると約3日18時間、約$360かかる見込みです。

延長戦順位表での結果

最終結果は次の通りです。

提出	探索	後処理	延長戦スコア
saitodevel01氏の提出	ビーム幅140	大近傍法あり	128,280,488
今回の最終提出	ビーム幅104	なし	128,286,311

差は5,823です。こちらはビーム幅が小さく、後処理もありませんが、延長戦順位表の完成品として比較対象を上回りました。

サンプルコード

本記事に合わせて、サンプルコードリポジトリを公開しています。データ生成、学習、評価、checkpointからC++提出コードを生成する処理、学習済みcheckpoint、最終提出コードを収録し、この記事の最終結果に至る一連の処理を確認できる構成にします。途中で試した別のlossや探索方法は含めず、最終的に採用した方法だけを収録します。

サンプルコードの基本的な使い方にAWSは必要ありません。まずはsmoke testで、少数の入力生成、C++の教師探索、dataset作成、CPUでの1 epochの学習、評価までがつながることを確認できます。次に、生成する入力数と教師の探索回数を小さく設定し、手元のPCで一連の処理を試すことを推奨します。この設定では最終提出と同じ品質のmodelは作れませんが、課金せずにデータ生成から提出コード作成までの仕組みを確認できます。

ローカル用のサンプルとAWSで動かした処理で、教師探索、入力seed、教師seed、dataset形式、model、loss、学習手順は共通です。AWSは互いに独立した教師探索を短時間で並列実行するために利用しただけなので、理論上はローカル用のコードでも、65,536入力と各入力2本の教師を同じ条件で用意できれば、記事で紹介したmodelとスコアを再現できる見込みです。ただし、ローカルで同じ量を生成するには非常に長い時間がかかり、実行環境やAtCoderの計測には揺れもあるため、同一スコアを保証するものではありません。

データ生成を待たずに後半だけを確認したい読者向けに、最終学習済みcheckpointも収録します。このcheckpointから、AOTI packageへの変換、圧縮、C++への埋め込み、提出コードのcompileまでを1つのcommandで実行できます。最終提出コードもそのまま収録するため、学習せずに探索部やmodelの呼び出し方だけを読むこともできます。

実際のAHCで使えるか

今回の方法を短期AHCの開催中にそのまま使うのは、ほぼ不可能です。強い教師を先に用意し、数万入力を生成し、約31時間学習し、C++へ載せて提出調整する必要があります。

長期AHCでも簡単ではありません。強い既存解法、再開可能な生成基盤、quota、予算、学習とdeploymentの事前準備がそろって、ようやく使えるかどうかという水準です。ローカルで数日待てば十分だったわけでも、AWSへ投げればすぐ終わったわけでもありません。

一方で、「時間をかければ強いが、提出には遅すぎる探索」がある問題なら、その探索を教師にして速いモデルへ圧縮する考え方自体は有効です。教師の質、候補の多様性、checkpoint選択、提出runtimeまでを1つの設計として考える必要があります。

まとめ

saitodevel01氏のデータ生成方法と学習方法が未知の状態から、次の経路を独立に作りました。

元延長戦1位の焼きなましを約100倍の探索量で動かす
各入力2候補、65,536入力の教師データを作る
score-weighted lossで165,248パラメータのモデルを学習する
AOTI packageへ変換し、読みやすいC++のビームサーチへ載せる
ビーム幅104・後処理なしで延長戦スコア 128,286,311 を得る

AHCのコンテスト期間中にはあまり見かけない方針で強い解を得るまでのプロセスを理解でき、とても面白い体験でした。めちゃめちゃ苦労したので、誰かがこれを利用してコンテスト優勝してくれることを願います。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up