nanoGPTにおける進化計算の実験 / 進化計算はハイパーパラメータチューニングにも有効だった
AIサイエンティストに投げるため、進化計算(Evolutional Computing)をnanoGPTのハイパーパラメータチューニングに適用してみることにした。
極めていい加減な気持ちで書いたため、これで上手くいくのか我ながら疑問だったが、結論から言うと有効だった。
まず、どんないい加減なコードで進化計算をしたのか書く。
本来はDNAを定義して、DNAからプログラムを生成するコードを書くべきなのだが、GPTの場合、その根本的な構造に大きな変化がないのでハイパーパラメータだけをチューニングすることにした。まあそのうち拡散モデルとAttentionと状態モデルを組み合わせたりする実験をするかもしれないが、今は面倒なので最低限のことをやってみる。
AI-ScientistのnanoGPT_liteのexperiment.pyをこんな感じで改造する。
def train(dataset="shakespeare_char", out_dir="run_0", seed_offset=0):
# Genetic data initialization ここで基本パラメータを設定
num_of_genes=10 # 個体の数
num_of_genetic_factors=9 # 遺伝的影響を与えるファクターの数+1
num_of_generations=5 # 世代の数
genes = np.random.rand(num_of_genes,num_of_genetic_factors)
# logging
val_log_info = []
train_log_info = []
for generation in range(num_of_generations):
print(f"=== Generation {generation} ====")
print(genes)
for gene_id in range(num_of_genes):
# -----------------------------------------------------------------------------
# default config values designed to train a gpt2 (124M) on OpenWebText
# data
gradient_accumulation_steps = 1
batch_size = 64 if dataset == "shakespeare_char" else 32
block_size = 256 # context of up to 256 previous characters
# I/O
eval_interval = 250 if dataset == "shakespeare_char" else 1000
log_interval = 10 if dataset == "shakespeare_char" else 100
eval_iters = 200
eval_only = False # if True, script exits right after the first eval
always_save_checkpoint = False # we expect to overfit on this small dataset, so only save when val improves
never_save_checkpoint = True # never save checkpoints
# model
n_layer = 6 # baby GPT model :)
n_head = 6
n_embd = 384
dropout = 0.2 # for pretraining 0 is good, for finetuning try 0.1+
bias = False # do we use bias inside LayerNorm and Linear layers?
# adamw optimizer
learning_rate = 1e-3 if dataset == "shakespeare_char" else 5e-4
max_iters = 500 if dataset == "shakespeare_char" else 100000
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0
# learning rate decay settings
decay_lr = True # whether to decay the learning rate
warmup_iters = 100 if dataset == "shakespeare_char" else 200
lr_decay_iters = max_iters # make equal to max_iters usually
min_lr = 1e-4 if dataset == "shakespeare_char" else 5e-5
# Apply genetic factor 遺伝係数をハイパーパラメータに適用する
n_layer = int(n_layer * genes[gene_id][1]*4)
n_head = int(n_head * genes[gene_id][2]*5)+4
n_embd = (384 // n_head) * n_head #Adjust n_embd
dropout = dropout * genes[gene_id][3]
learning_rate = learning_rate * genes[gene_id][4]
beta1 = beta1 - genes[gene_id][5]/10.0
beta2 = beta2 - genes[gene_id][6]/100.0
warmup_iters = int(warmup_iters * genes[gene_id][7]*2)
#Show applied params
print(f"n_layer={n_layer}, n_head={n_head}, dropout={dropout:.2f}, learning_rate={learning_rate:.2e}, beta1={beta1:.2f}, beta2={beta2:.2f}, warmup_iters={warmup_iters}")
そして、ループの終わりに成績順にソート(選択)して交配・突然変異を加える
print("training done")
print(f"Best validation loss: {best_val_loss}")
print(f"Total train time: {(time.time() - og_t0) / 60:.2f} mins")
# record best vall loss 成績を保存する
genes[gene_id][0] = best_val_loss #store to gene[0]
# Selection 成績順に選択する
print("=== Generation End")
print(genes)
indices = np.argsort(genes[:, 0])
genes = genes[indices]
print("=== Sorted")
print(genes)
# Crossover 交配
num_genes = len(genes)
half_num_genes = num_genes // 2 #上位半数の遺伝子だけ残す
new_genes=np.zeros((num_of_genes,num_of_genetic_factors))
for i in range(num_of_genes):
# pick two parents 二つの親を選ぶ
idx1 = np.random.randint(0,half_num_genes)
idx2 = np.random.randint(0,half_num_genes)
# crossover gene
for j in range(1,num_of_genetic_factors):
if np.random.rand() > 0.5:
new_genes[i][j] = genes[idx1][j]
else:
new_genes[i][j] = genes[idx2][j]
# Mutation
if np.random.rand() < 0.05: #mutation rate 5% 突然変異
new_genes[i][j] = np.random.rand()
genes=new_genes
print("=== Crossover and Muted")
print(genes)
たったこれだけでnanoGPTの性能が上がることを確認した。
このように、実に綺麗にlossが下がっていく。
これはかなり手を抜いた進化計算なので、本来のDNAのようなエンコード/デコードはしていないが、実質的な意味において進化計算になっていることは少し詳しい人ならわかるだろう。
また、単に個体選択においてパラメータが増えれば自動的に精度が上がるような、単純な事実を示しているだけである可能性が高い(調べる気も起きないので調べていないが)。のだが、この話の本質はそこではない。なぜ、このような雑なやり方で正しい答えに近づけるかという問いだ。
今回は個体が10個で遺伝的パラメータは8個なので、初期状態で与えられた80個の乱数(遺伝子)の組み合わせだけで性能が上がるということになる。
もしも個体を100とか1000とか与えたら、もっと違った結果になるのだろうか。
AI-Scientistの実装だと、研究の「タネ」を与えてあとはAI任せ、というのが理想なのだが、全てのコードがAIderだけで書けないのと同じように現状ではまだいろいろと制約がある。ソースコードを直接いじらせるよりも、操作できるパラメータをコマンドラインから与えさせて、コマンドラインを考えさせるようにした方がより少ない手数で欲しい結論が得られそうだ。
もちろんSakanaAIも、AI-Scientistが本物の研究者を置き換えるとまでは思ってないだろうが、その可能性は十二分に感じることができる。
欲しいのはこういう時に「個体の数を変えて実験し、その結果を考察せよ」みたいな短い指示を与えて論文を書かせるものだ。これに比べるとAI-Scientistはやや独創的なアイデアを出させることにこだわりすぎているようにも見えなくもない。
もうちょっとシンプルな、サイエンティストよりもインターンに近い感じのやつが欲しい。