driftax — JAXによる「Generative Modeling via Drifting」の実装

AI/ML
driftax — JAXによる「Generative Modeling via Drifting」の実装

driftax — JAXによる「Generative Modeling via Drifting」の実装

本リポジトリは、論文 "Generative Modeling via Drifting"（arXiv:2602.04770）の手法をJAXで実装した軽量なリポジトリです。コアとなるドリフト過程（drifting）の数値的実装、サンプリングループ、学習ステップの素地が含まれており、研究再現や拡張のための出発点となります。最小限の例（examples）とパッケージ設定（pyproject.toml）、ライセンスが付属しています。引用情報もREADMEに明記されています（研究利用時は原著を参照してください）。

https://github.com/wrkhard/driftax

概要

driftaxは、Mingyang Dengらの論文「Generative Modeling via Drifting」（arXiv:2602.04770）をJAXで実装したリポジトリです。リポジトリ自体は小規模で、コア実装とサンプルコードが含まれています。JAXの自動微分、JITコンパイル、乱数管理（PRNG）を活かした関数型スタイルの実装により、研究実験の再現や手法の拡張が容易になることを意図しています。READMEに論文の引用情報が記載されており、研究用途での利用を促しています。ライセンスファイルも同梱されているため、利用条件の確認も可能です。

GitHub

リポジトリの統計情報

  • スター数: 1
  • フォーク数: 0
  • ウォッチャー数: 1
  • コミット数: 2
  • ファイル数: 8
  • メインの言語: Python

主な特徴

  • 論文「Generative Modeling via Drifting」のJAX実装（研究再現向けの最小実装）
  • JAXの利点（JIT、vmap、autograd）を活かした高効率な数値計算と再現性
  • examplesディレクトリに動作例を収録し、実験スタートが容易
  • pyproject.tomlによるパッケージ設定とLICENSEによる利用条件の明示

技術的なポイント

driftaxはJAXを用いることで、ドリフトベースの生成モデル（論文で提案されたdrifting手法）のコア演算を簡潔かつ高速に記述することを目指しています。主な技術ポイントは次のとおりです。

  • 関数型プログラミングスタイル：状態を明示的に取り回すことで、JAXの純関数性に適合した設計になっていると推察されます。これによりjit/grad/vmapを安全かつ効果的に適用できます。
  • 数値的なドリフト・サンプリング：ドリフト過程（時間発展や逆過程）の数値解法が実装され、ステップごとの状態更新とサンプリングループが書かれています。JAXのベクトル化機能によりバッチ処理が効率化されます。
  • 自動微分による学習：目的関数（論文で定義される損失）に対する勾配計算はJAXのautogradで処理され、最適化ループが高速に実行されます。PRNG管理はJAXの乱数キーで扱うため、実験の再現性が高く保たれます。
  • 軽量で拡張しやすい構成：コードベースは小規模（数ファイル）で、コア実装と例が分かれているため、論文アルゴリズムの検証やモデル部位（ニューラルネットワークの構造、スケジューリング、数値積分手法）の切り替え・改良がしやすいです。
  • パッケージングと利用準備：pyproject.tomlが含まれているため、プロジェクトの依存管理や配布準備が容易です。LICENSEファイルにより利用条件が明示されています。
  • 研究利用への配慮：READMEに論文の引用フォーマットが記載されており、学術的な利用を念頭に置いた設計です。

注意点として、現時点のリポジトリは非常にシンプルでコミット数も少なく、フルスケールの実験セットアップやトレーニング済みモデルは含まれていない可能性が高いです。従って、実験再現や大規模ベンチマークには追加実装やチューニング、依存ライブラリ（例えばニューラルネットワークライブラリやデータローダ）の組み込みが必要になるでしょう。とはいえ、JAXでの純粋実装は学術実験や手法の理解を深める上で有益であり、コードを起点に自分のデータセットやモデルへ適用することができます。

プロジェクトの構成

主要なファイルとディレクトリ：

  • .DS_Store: file
  • LICENSE: file
  • README.md: file
  • examples: dir
  • pyproject.toml: file

…他 3 ファイル

リポジトリはコアの実装ファイル群とexamplesディレクトリ、パッケージ設定・ライセンスのみのシンプルな構成です。大規模な実験スクリプトやデータセットは含まれていないため、実行にはJAX環境の整備と必要に応じた依存追加が必要です。

まとめ

JAXで書かれた論文再現向けの軽量実装。拡張と検証の良い出発点。

リポジトリ情報：

READMEの抜粋：

driftax

JAX implementation of https://arxiv.org/abs/2602.04770

If you use this library please cite the original authors here:

@article{deng2026drifting, title={Generative Modeling via Drifting}, author={Deng, Mingyang and Li, He and Li, Tianhong and Du, Yilun and He, Kaiming}, journal={arXiv preprint arXiv:2602.04770}, year={2026} }