2018年にOptunaを公開して以降、Optunaは機械学習のハイパーパラメータ最適化にとどまらず多様な分野で利用されるソフトウェアになりました。様々な分野での活用が進む一方で新たに見えてきた課題もあります。大規模な最適化における実行速度や消費メモリの課題などがその代表例です。

近年RuffやPolars、Pydantic V2を始めとする様々なソフトウェアが、Rustによる大幅な高速化を実現していることを踏まえ、Rustによる高速なOptuna実装である「Rustuna」を開発・公開しました。本記事では、Rustunaの特徴や使い方、今後の展望を紹介します。

• GitHub: https://github.com/optuna/rustuna/
• 公式ドキュメント: https://rustuna.readthedocs.io/

 なおOptunaの5回目のメジャーリリースである「Optuna 5.0」も同時公開しています。Optuna 5.0のリリースに関する詳細はOptuna 5.0 リリースブログをご確認ください。

Rustunaの使い方

RustunaはOptunaとの高い互換性を重視して開発を進めてきました。そうすることでOptunaを使ったことのある人であれば、新たにAPIを覚え直す必要はなくスムーズにRustunaを動かすことができます。使い方は次のとおりです。

$ pip install rustuna

import rustuna

# 1. 目的関数を定義 (ここでは最小化問題を仮定)
def objective(trial: rustuna.Trial) -> float:
    # 2. Suggest APIを通して探索空間を定義
    x = trial.suggest_float("x", -10, 10)
    y = trial.suggest_float("y", -10, 10)
    return (x - 2) ** 2 + (y + 5) ** 2

# 3. Studyを作成
study = rustuna.create_study()

# 4. パラメーターを変えながら目的関数を100回評価
study.optimize(objective, n_trials=100)

# 5. 最も良い評価履歴を表示
print(study.best_trial)

 Rustunaは現時点ではOptunaほど多くの機能を備えたソフトウェアではありませんが、それでも多くのユーザーが必要とする基本的な機能はすでに備えています。また一部の機能を除いて、RustunaはPythonの依存関係を持ちません。これにより近年急増しているサプライチェーン攻撃のリスクも緩和されます。Optunaの実行速度や消費メモリについて特に課題を感じていなかったというユーザーのみなさんも、ぜひ一度お試しください。

Rustunaがサポートしていない機能については、Optunaの機能を呼び出すことも可能です。使い方によってはRustunaの高速化の恩恵が得られない場合もありますが、例えば探索終了後にOptunaの可視化関数で最適化履歴の分析などが可能です。詳細は Rustuna公式ドキュメントをご覧ください。

Rustunaの実行速度

Rustunaは高速化に重きを置いて開発しています。Rustunaが提供する各種サンプリングアルゴリズムについてOptunaとの実行速度を比較した結果を図1に示します（M4 ProのMac Bookにて計測）。

 TPEおよびMOTPE（多目的TPE）は、Optunaがデフォルトで使用するベイズ最適化のアルゴリズムです。TPEは探索パラメータをサンプリングするために必要な計算時間がTrial数に比例して増加します。Optunaでは目安として 探索回数が1000 Trials以下 の場合に利用を推奨していましたが、Rustunaの高速なTPE実装であればその10倍となる1万 Trialsでも1分以内で実行できます。マルチスレッド使用時はこの差がさらに顕著になります（表1）。 

 CMA-ESやNSGA-IIも同様に大幅な高速化を達成していますが、単にRustで実装されているからではありません。これらのアルゴリズムは本来サンプリングに必要な計算時間がTrial数に対して一定な手法ですが、Optunaではいくつかの設計上の理由からTrial数の増加に伴い探索時間が悪化します。Rustunaは一部の設計を見直すことで高速化を実現しました。

 このようにRustunaでは、実装と設計の両方のアプローチにより大幅な高速化を達成しています。

省メモリなストレージ実装

Rustunaは消費メモリの改善にも注力しています。設計や実装面での細かいテクニックによる改善もありますが、Optunaと大きく異なる点として探索に不要な試行履歴を破棄する機能を提供します。

TPEやガウス過程ベースのベイズ最適化アルゴリズムとは異なり、CMA-ESやNSGA-IIのような最適化アルゴリズムは探索時に評価履歴の一部しか参照しません。しかしOptunaではすべての評価履歴を常にメモリに保持しており、探索が進むにつれ消費メモリや実行時間が悪化します。Rustunaでは各種ストレージが提供する apply_discard オプションを利用することで、探索に不要な試行履歴をメモリから破棄して探索を実行します。

import rustuna

def objective(trial: rustuna.Trial) -> float:
    x = trial.suggest_float("x", -10, 10)
    trial.set_user_attr("large_text", "a" * 10000)
    return (x - 2) ** 2

# Set apply_discard=True to enable feature to discard trials.
storage = rustuna.storages.JournalFileStorage("./rustuna-journal.db", apply_discard=True)

sampler = rustuna.samplers.CmaEsSampler()
study = rustuna.create_study(sampler=sampler, storage=storage, study_name="example")
study.optimize(objective, n_trials=100000)

探索時とは異なり、最適化結果の確認や分析の際には、すべての評価履歴を参照したいケースもあります。その際は apply_discard オプションを指定せずにストレージを読み込むと、破棄されていたトライアルを含むすべての評価履歴が読み込まれます。

storage = rustuna.storages.JournalFileStorage("./rustuna-journal.db")
study = rustuna.load_study(storage=storage, study_name="example")
print(study.best_trial)

本機能の効果を実サービスにおいて検証した結果もあわせて紹介します。Matlantis株式会社が提供する結晶構造探索サービス「Matlantis CSP」では、Optunaを活用したNSGA-IIベースの手法を独自に設計・実装しています（解説記事: Matlantis CSPにおけるOptunaを使った結晶構造探索）。Rustunaに置き換えることで消費メモリや実行速度にどのような影響があるかを検証した結果を図2に示します。 

 Optuna版では探索が進むにつれて評価履歴が蓄積されることで消費メモリが大幅に増加します。またNSGA-IIベースの最適化アルゴリズムを実装するにあたって評価履歴の一覧を順にチェックし必要な値を取り出す処理が含まれているため、評価履歴が10～20万件に膨れ上がるにつれて処理時間が悪化します。

 一方で、Rustunaでは apply_discard オプションの利用により、不要な評価履歴が破棄されることで、探索回数の増加に伴う消費メモリの増加は小さく、処理時間の悪化もほとんどありません。これにより探索回数のさらなる増加や高い並列数での実行が可能となります。本検証結果を踏まえて、現在はRustunaの本格導入に向けた開発が進行しています。

今後のロードマップ

ここからはRustunaの今後の開発ロードマップについても触れたいと思います。Rustunaはすでに各種最適化アルゴリズムやストレージ、ハイパーパラメーター重要度計算機能など多くの基本機能がサポートされており、Optunaほど多くの機能は提供していませんが大多数のユーザーにとって十分な機能が揃いつつあると考えています。

 今後はさらなる利便性の向上や用途の拡大に応えられるように、次のような改善および機能開発を予定しています。

• 安定性の向上：Rustunaはまだ成熟したソフトウェアではありません。今回のリリースからしばらくは新たなバグの発見も予想されるため、ユーザーの皆さまからのご報告やフィードバックをもとに改善を進めていきます。
• よりシームレスなOptuna連携：RustunaはOptunaの置き換えを目的として開発されたソフトウェアではありません。Optunaのように高機能なソフトウェアを目指すのではなく、協調的に利用されるソフトウェアになることを目指しています。より手軽にOptunaの機能を呼び出せるように連携強化を進めてまいります。
• 多言語サポート：Rustで実装されたソフトウェアはC-APIやWebAssemblyを経由することで、比較的手軽に多言語対応が可能です。これまではPythonバインディングにフォーカスして開発を進めてきましたが、今後はNode.jsやその他の言語への展開も検討しています。

 また、Rustunaの開発を通して得られた知見を活かし、Optunaの改善にも努めたいと考えています。多くのユーザーを抱えるOptunaは後方互換性を重視して開発を進めているため、新機能や大規模な内部構造の変更には慎重なプロセスが必要です。しかし、今回Rustunaで採用したBulk InsertベースのStorage APIやTrialQueueの導入、前節で紹介した apply_discard オプションなどは、Optunaの将来的な性能改善においても有益です。今後も互換性を慎重に維持しながら、これらの成果を段階的にOptunaへフィードバックしていきたいと思います。

おわりに

本記事では、Rustunaの開発背景から使い方、Optunaと比べた利点について紹介しました。Rustunaが成熟したソフトウェアに成長するためには、多くの方に使っていただきフィードバックをいただくことが重要だと考えています。ぜひRustunaをお試しいただき、なにか問題を見つけた際は気軽にGitHub issueよりご報告をお願いします！