仕事でOptunaを使う機会があったので、論文を軽く読んでまとめてみる。

Optunaとは

OptunaはPFNが開発したハイパーパラメータのチューニングツールであり、近年、勾配ブースティングやニューラルネットワークのチューニングによく使われている。
既存のツールに対するOptunaの特徴として、下記３点が挙げられている。

Define-by-runで動的にパラメータ探索空間を設定可能
効率的なsampling, pruningアルゴリズム
設定が簡単で、軽量な実験から分散処理による大規模な計算まで対応可能

1. Define-by-run

define-by-runは深層学習フレームワークの分類に使われる用語だが、筆者らはパラメータチューニングツールにも同様の概念を提唱している。
従来のツールはdefine-and-run、すなわちパラメータの探索範囲を最初に設定した上で、最適なパラメータの探索を行うものだった。

一方、define-by-runでは動的に探索範囲を設定することができる。これは例えばニューラルネットワークのレイヤー数とユニット数を同時にチューニングする際などに役に立つ。

2. 効率的なsampling, pruningアルゴリズム

パラメータ最適化の効率は探索戦略（探索すべきパラメータを決める）と性能評価戦略（学習曲線から現在探索中のパラメータの良さを判断し、捨てるかどうか判断する）により決まる。

論文には探索（サンプリング）アルゴリズムの詳細は記載されていないが、TPE (tree-based Parzen estimator)*1, CMA-ES(Covariance Matrix Adaptation - Evolution Strategy)*2, GP-BO(Bayesian optimization) *3などを用いているとのこと。サンプリング手法にはパラメータの相関を利用するrelational samplingと独立にサンプルするindependent samplingがある。TPEはindependent sampling、CMA-ES, GP-BOはrelational samplingである。

また、性能評価については、見込みの悪い試行を途中で終わらせること（pruning、枝刈り）が重要である。 OptunaではAsynchronous Successive Halving (ASHA)*4と呼ばれるアルゴリズムを採用している。ASHAでは試行中のパラメータセットの暫定的な順位に基づき、分散処理の個々のworkerが非同期的にpruningを行う手法とのこと。

使い方

Optunaではパラメータの最適化をobjective functionの最小化・最大化の問題と捉える。
最適化のプロセスをstudy、objective functionの１回の評価をtrialと呼ぶ。パラメータの探索空間はtrialオブジェクトのメソッドにより動的に構築される。

以下のコードでは（モデルのチューニングではないが）、簡単な例として２次関数の最小化を行っている。

import optuna

def objective(trial):
    x = trial.suggest_uniform('x', -10, 10)
    return (x - 2) ** 2

study = optuna.create_study()
study.optimize(objective, n_trials=100)

study.best_params  # E.g. {'x': 2.002108042}}

以下はより実践的な、TensorFlowでニューラルネットワークのパラメータチューニングを行う場合のコード抜粋。完全なコードは下記リンク先参照。 https://colab.research.google.com/drive/1MQp5rn6Dxhdv4r9PuJOyGf1sdsS2PHF1

ここでは、隠れ層の数、ユニット数、weight decayといったモデルのパラメータに加え、optimizerの種類とそのパラメータも最適化している。従来のツールでこうしたチューニングを行うためには、最初にかなり複雑なパラメータ探索範囲を記述する必要があるが、Optunaのdefine-by-runの思想により、簡潔に実装することができる。

def create_model(trial):
    # We optimize the numbers of layers, their units and weight decay parameter.
    n_layers = trial.suggest_int("n_layers", 1, 3)
    weight_decay = trial.suggest_loguniform("weight_decay", 1e-10, 1e-3)
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Flatten())
    for i in range(n_layers):
        num_hidden = int(trial.suggest_loguniform("n_units_l{}".format(i), 4, 128))
        model.add(
            tf.keras.layers.Dense(
                num_hidden,
                activation="relu",
                kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
            )
        )
    model.add(
        tf.keras.layers.Dense(CLASSES, kernel_regularizer=tf.keras.regularizers.l2(weight_decay))
    )
    return model


def create_optimizer(trial):
    # We optimize the choice of optimizers as well as their parameters.
    kwargs = {}
    optimizer_options = ["RMSprop", "Adam", "SGD"]
    optimizer_selected = trial.suggest_categorical("optimizer", optimizer_options)
    if optimizer_selected == "RMSprop":
        kwargs["learning_rate"] = trial.suggest_loguniform("rmsprop_learning_rate", 1e-5, 1e-1)
        kwargs["decay"] = trial.suggest_uniform("rmsprop_decay", 0.85, 0.99)
        kwargs["momentum"] = trial.suggest_loguniform("rmsprop_momentum", 1e-5, 1e-1)
    elif optimizer_selected == "Adam":
        kwargs["learning_rate"] = trial.suggest_loguniform("adam_learning_rate", 1e-5, 1e-1)
    elif optimizer_selected == "SGD":
        kwargs["learning_rate"] = trial.suggest_loguniform("sgd_opt_learning_rate", 1e-5, 1e-1)
        kwargs["momentum"] = trial.suggest_loguniform("sgd_opt_momentum", 1e-5, 1e-1)

    optimizer = getattr(tf.optimizers, optimizer_selected)(**kwargs)
    return optimizer


def objective(trial):
    # Get MNIST data.
    train_ds, valid_ds = get_mnist()

    # Build model and optimizer.
    model = create_model(trial)
    optimizer = create_optimizer(trial)

    # Training and validating cycle.
    with tf.device("/cpu:0"):
        for _ in range(EPOCHS):
            learn(model, optimizer, train_ds, "train")

        accuracy = learn(model, optimizer, valid_ds, "eval")

    # Return last validation accuracy.
    return accuracy.result()