galilai-group/stable-pretraining

stable‑pretraining – 自己教師学習のための柔軟なフレームワーク

何であるか

• PyTorch Lightning（実験的JAX/Flax‑NNXバックエンド付き）の上に構築されたオープンソースPythonライブラリ。研究者が自己教師学習（SSL）手法を最小限のボイラープレートでプロトタイピング、トレーニング、評価できるようにする。
• コアアイデアは、loss（トレーニング用）とログに記録したい任意のテンソルを含む辞書を返すforward関数のみを実装すること。ライブラリが通常のLightningの処理（データロード、チェックポイント、再開、ログ記録、ローカル実験レジストリ）をすべて処理する。

主要コンポーネント

コンポーネント	役割
データ	辞書形式のデータセットを扱う。HuggingFaceデータセット、標準PyTorchデータセットからの変換、カスタム辞書返却データセットをサポート。Korniaを介してGPUに重いマルチビュー変換を移動するgpu_transformフックを提供し、より高いスループットを実現。
モジュール	ベースライン、プロジェクター、ユーザー提供のforwardコールバック（例：forward.simclr）をラップする薄いラッパー。training_step/validation_stepを書く必要なし。ラッパーは返された辞書からlossを抽出する。
コールバック	オンライン評価（線形プローブ、K‑NN、RankMe、LiDAR、CLIPゼロショット、画像検索、潜在空間可視化）、教師-生徒EMA、重み減衰スケジュール、チェックポイント処理（WandB、HuggingFace、sklearnモデル）、システムログ（SLURM情報、環境ダンプなど）のための即時利用可能なユーティリティ。
トレーナー / マネージャー	LightningのTrainerの薄いフェース。Managerはトレーナー、モジュール、データモジュールをバンドルし、トレーニングループを実行、ログを書き込み、実行をローカルcache_dirレジストリに登録する。
JAXバックエンド（実験的）	stable_pretraining.jaxでPyTorch APIをミラー。Flax‑NNX、Optax、JAXネイティブプリミティブを使用しながら、同じ設定ファイル、コールバック、ログフォーマットを維持。

組み込みSSLメソッド

• SimCLR、BYOL、VICReg、Barlow Twins、SimSiam、SwAV‑Sinkhorn、NT‑Xent、InfoNCE、および新しいLeJEPAとJet逆変換エンコーダ。カスタムforward関数を提供することで、新しいメソッドを追加可能。

バックボーンと最適化手法

• MLP、ResNet‑{9,18,34,50,101,152}、ConvMixer、ViT‑{tiny,small,base,large}。
• Adam、LARSなどの最適化手法；CosineAnnealingLRなどのスケジューラをサポート。

クイックスタート

# 最新リリースをインストール
pip install stable-pretraining

# または開発版をインストール
git clone https://github.com/galilai-group/stable-pretraining.git
cd stable-pretraining
pip install -e .
python -m stable_pretraining.quickstart   # 仮想データ上でSimCLRの1CPUエポックを実行

このコマンドは小さなトレーニング実行を開始し、ローカルWebビューア（spt web ./spt-quickstart/runs）を開き、損失曲線とオンライン線形プローブの精度を表示する。

一般的なワークフロー

1. データの定義 – spt.data.HFDatasetを使用するか、PyTorchデータセットをラップ；必要に応じてgpu_transformを追加して高速なマルチビュー変換を実現。
2. メソッドの選択 – 事前に実装されたforward（例：spt.forward.simclr）を選ぶか、独自に作成する。
3. Moduleの作成 – ベースライン、プロジェクター、損失、最適化設定、forwardを渡す。
4. コールバックの追加 – 例：OnlineProbeで下流性能を監視。
5. 実行 – spt.Manager(trainer=trainer, module=module, data=dm)()。

なぜ重要か

• 安定性：すべてのトレーニングループがLightningを経由するため、堅牢なチェックポイント、マルチGPU（DDP）サポート、自動再開が可能。
• 柔軟性：SSL目的をforwardコールバックを交換するだけで変更可能。パイプラインの残り部分（データ処理、ログ、評価）は変更なし。
• 高速性：GPU側の変換は現代のハードウェアでスループットを劇的に向上させる。リポジトリのベンチマーク表で実証済み。
• フレームワーク間の同等性：実験的JAXパスにより、同じ設定とログ形式でTorchとJAXの実装を直接比較可能。

ドキュメントとコミュニティ

• 完全なドキュメント： https://galilai-group.github.io/stable-pretraining/
• ベンチマーク： https://github.com/galilai-group/stable-pretraining/tree/main/benchmarks
• サポート用Discord： https://discord.gg/adzpqWKM25

誰が使うべきか

• 新しいSSL目的を構築したい研究者で、信頼できるトレーニング基盤を必要とする人。
• 大規模な事前学習に再現性があり、チェックポイントに優れたパイプラインを必要とするエンジニア。
• 同じメソッドのTorchとJAX実装を比較したい人。

---

上記のすべての詳細はリポジトリのREADMEから直接取得；追加の機能は推測されていない。