Rustによる高速なOptuna実装の試作

はじめに

OptunaはすべてPythonで実装されたソフトウェアです。速度を必要とする処理もC拡張を使わずにすべてPythonで実装しており、このおかげで様々な実行環境にトラブルなくインストールできます。このことは多種多様な用途に利用されるOptunaにおいて非常に重要な点であり、現時点では我々はこの方針を変えるつもりはありません。

その一方、限られた一部のユースケースにおいて実行速度が問題となるケースや、既存システムへの組み込みの観点でPython以外の言語から呼び出したいケースも存在しました。近年、Ruffを始めとするいくつかのソフトウェアがRustによるPythonのツールチェインの大幅な高速化に成功していることを踏まえ、RustによるOptuna実装の試作をクローズドに開始しました。現時点で具体的な公開の予定はありませんが、本ブログにて開発の目的やRustによって置き換えることで得られた恩恵についてご紹介します。

試作を開始した目的

Rust版 Optuna実装の試作には、大きく2つの目的があります。

Optunaの高速化

Optunaは主に機械学習のハイパーパラメータ最適化をターゲットとしています。このようなユースケースではモデルの学習や評価に時間がかかることも多く、Optunaの実行時間がボトルネックになるケースはあまりありません。

しかしながらブラックボックス最適化の応用可能な事例は、機械学習のハイパーパラメータ最適化以外にも数多くあります。例えば下記の発表で紹介されているマテリアルサイエンス分野でのOptuna活用事例では、未知の結晶構造の探索にOptunaを活用しています。この事例では非常に多くの評価回数を必要とするため、Optunaが高速で動作することも重要です。探索回数が増えるほど未知の結晶構造を見つけられる確率が増えるからです。
Efficient Crystal Structure Prediction using Universal Neural Network Potential and Genetic Algorithm 

多言語対応およびOptuna Dashboardでの活用

Rust版 Optuna実装は、RustやPythonから使えるのはもちろんJavaScriptバインディングやC-APIも提供しています。特にJavaScriptバインディングについてはOptuna Dashboardでの活用を目指しています。

Optuna Dashboardは、Optunaの試行結果を分析する際に用いるWebダッシュボードです。最近ではJupyter Lab拡張やVS Code拡張も提供しています。これらの実装にあたり、Optunaに実装されている可視化関数モジュール (optuna.visualizationモジュール) やJournalFileStorageのローダーなど多くの処理をTypeScriptで再実装してきました。これらの一部は速度面の課題も抱えており、Rust版 Optuna実装が高速なJavaScriptバインディングを提供することでユーザー体験の向上が期待できます。また開発上の利点もあり、それについての詳細は後ほどご紹介します。

Rust版Optuna実装によるPythonプログラムの高速化

今回試作したRust版 Optuna実装では、Optunaのコア機能をRustで再実装するとともに、PyO3を使ってPythonバインディングを提供しています。ここではRust版Optuna実装の実装方針および現時点での実行速度についてご紹介します。

Optunaとの互換性と速度の比較

Ruffが既存のツールチェインと高い互換性を提供しているように、Rust版 Optuna実装においてもOptunaのAPIとの互換性を重視しています。ユーザーは新たにAPIを覚え直す必要がなく、既存プロジェクトから手軽に移行したり、Optunaがこれまで実装してきた資産をそのまま活用することができます。Rust版 Optuna実装はOptunaと基本的に同じAPIを提供しているため、次のようなコードはimport文を変えるだけでそのまま動作します。

import <OPTUNA_RUST> as optuna


def objective(trial: optuna.Trial) -> float:
    x = trial.suggest_float('x', -10, 10)
    y = trial.suggest_float('y', -10, 10)
    value = (x - 2) ** 2 + (y + 5) ** 2
    return value


study = optuna.create_study()
study.optimize(objective, n_trials=1000)
print(study.best_trial)

上記のコードについて、Trial数やパラメーター数を変化させた際の、現時点での実行速度の比較結果は次の通りです。

TPESampler (multivariate=False)

OptunaはデフォルトでTPESamplerを使用します。公式ドキュメントの比較表にてTPESamplerの推奨Budgetが100-1000となっているように、計算量や最適化性能の観点から多くの試行回数が確保できる状況においてTPEは必ずしも適した手法というわけではありません。評価回数を多く確保できる状況では他の手法が推奨されるケースも多くあります。例えばRandomSamplerの実行時間の比較結果は次のようになります。

RandomSampler

上記の結果で重要な点の1つは、Trial数を10倍にした際にRust版 Optuna実装の実行時間は約10倍で済んでいるのに対して、Optunaの実行時間は50-100倍にまで伸びてしまっている点です。この原因はOptunaが提供するいくつかの便利な機能に起因しています。Optunaでは利便性と速度低下のバランスを慎重に判断して開発を進めていますが、多くのユーザーにとっての利便性を追求した結果、一部のユーザーにとってこのような速度低下を招いてしまっているのも事実です。

Rust版 Optuna実装は高速化により重きを置いています。速度に大幅な改善がなければ存在意義がなくなってしまうからです。互換性を重視してOptunaの設計や機能に合わせすぎてしまうと高速化に限界が生じてしまうため、enqueue_trial()のような一部の機能については導入を慎重に判断しています。 “基本的に” 同じAPIを提供していると書いた理由の1つはここにあります。

Rust版 Optuna実装とOptunaの連携

Rust版 Optuna実装はOptunaの置き換えを目的としたものではありません。Optunaのすべての機能をRustで再実装するわけではなく、併用しながら互いに強みを補完し合うことが目的です。

例えばOptunaにおいて利用者の多い機能の1つにRDBStorageがあります。試行結果の永続化や分散最適化を可能にする重要な機能ですが、高速化の観点で考えるとI/Oが主なボトルネックとなるこの機能をRustで再実装するメリットはほとんどありません。Rust版 Optuna実装ではこのような機能を再実装しないで済むように、変換用のユーティリティーを提供します。

import <OPTUNA_RUST> as optuna
from <OPTUNA_RUST> import FromOptunaStorage
from optuna.storages import RDBStorage


def objective(trial: optuna.Trial) -> float:
    ...


storage = FromOptunaStorage(RDBStorage("sqlite://db.sqlite3"))
study = optuna.create_study(storage=storage, study_name="example")
study.optimize(objective, n_trials=10)

上記の例ではRust版 Optuna実装からOptunaの機能を部分的に使用していますが、その逆も可能です。例えば次に示すコードはOptunaから部分的にRust版 Optuna実装を利用しています。

import optuna

# Import Rust-Optuna's FanovaImportanceEvaluator.
from <OPTUNA_RUST>.optuna.importance import FanovaImportanceEvaluator


def objective(trial: optuna.Trial) -> float:
    ...


study = optuna.create_study()
study.optimize(objective)

importance = optuna.importance.get_param_importances(
    study, evaluator=FanovaImportanceEvaluator()
)

こちらの方法は先程の方法と比較して高速化の観点で見ると一部限界がありますが、既存のコードを大きく書き換えることなくRust版 Optuna実装を試しやすく、なにか問題が発生したり速度面の恩恵が得られなかった際には戻す、という切り替えがより手軽になります。

Optuna Dashboardの技術的課題への対処

最後に紹介するのはOptuna Dashboardでの活用ビジョンです。Optuna Dashboardは、Optunaの試行結果を閲覧するWebアプリケーションです。このWebダッシュボードの開発には大きく2つの技術的課題が存在します。1つは実行速度の課題、もう1つはユニットテストの難しさです。

実行速度の改善

Optuna DashboardはPythonのAPIサーバーとTypeScript/Reactで実装したSingle Page Applicationからなります。APIサーバーの主な役割はStudyやTrial一覧をWebフロントエンドに伝えるためのJSON APIの提供のみで、Dashboardの重要な処理のほとんどはWebフロントエンドで実行されます。このことはOptuna Dashboardの設計のシンプルさやユーザー体験の向上に大きく寄与するとともに、VS Code拡張の提供にも大いに役立っています。

しかしながら、Trial数が数万を超えるケースにおいて現状のOptuna Dashboardは残念ながらかなり遅く使いづらいのが現状です。これにはいくつかの根本的な要因がありますが、Rust版 Optuna実装が高速なJavaScriptバインディングを提供すれば、現状のシンプルな設計を保ちながらも処理が高速になりユーザー体験が向上します。

メンテナンス性の高いユニットテストの記述

Optuna Dashboardのフロントエンドのロジックの多くはOptunaのTrial一覧を引数として受け取ります。Trialにはparamsやdistributionsなど必要な情報が多く、複数生成しようとするとセットアップのためのコードが非常に冗長になります。フィクスチャー生成用のユーティリティーを導入することも検討できますが、それ自体が肥大化して何をテストしているのかわからないという負債になる懸念もあります。

Rust版Optuna実装のJavaScriptバインディングを用意したことで、テストで使うためのStudyやTrialの情報をNode.js側で容易に生成できます。Trialを用意するためにDistributionやinternal_params等をハードコードする必要がなくなり、内部構造の変化にも強くユニットテストのメンテナンス性が向上します。

import * as optuna from '<OPTUNA_RUST>.js';

const study = optuna.create_study();
const objective = (trial) => {
    const x = trial.suggest_float("x", -10.0, 10.0);
    const y = trial.suggest_int("y", -10, 10);
    const z = trial.suggest_categorical("z", ["foo", "bar"]);
    return value = (x - 5) ** 2 + (y + 2) ** 2
}

study.optimize(objective, 10)
console.log(study.best_trial())

これまでOptuna Dashboardでの開発は、E2Eテストに強く依存してしまっている状況でした。実際にOptunaを実行して、試行結果をSQLite3に保存し、実際にPythonのAPIサーバーを起動して問題なく表示されているかpytestとPlaywrightでテストしています。この方法は実際にOptunaを動かすことで容易に様々な目的関数の最適化履歴を短いコードで用意できますが、コードの正しさを網羅的に検証することは困難でした。上記のようにPythonを経由せずにJavaScript側で手軽にOptunaの試行履歴が生成できればメンテナンス性の高いユニットテストの記述が容易になります。

おわりに

本記事では、Optuna開発チームで新たに試作を開始したRust版Optuna実装の目的とその活用ビジョンについて紹介しました。冒頭でお伝えした通り、Rust版 Optuna実装はあくまで試作段階にあり、現時点では具体的な公開の予定はありません。引き続きコミッタ内部で開発や有用性の検証を継続していく予定です。

このようにOptuna開発チームでは、より多くのユーザーにより便利にOptunaをお使いいただけるよう、様々な可能性を模索し日々改善を続けています。一緒にOptunaを開発するパートタイムエンジニアを随時募集しています。本記事を通して、ご興味を持っていただけたという方はぜひ下記ページをご確認ください。

Software engineer (Optuna) / ソフトウェアエンジニア（Optuna） / 株式会社Preferred Networks