galilai-group/stable-pretraining

何を解決するか

"ファウンデーションモデル"（CLIP、DINO、SimCLR などの一般的な視覚特徴を学習するモデル）用の大規模ニューラルネットワークの訓練は、非常に手間がかかることが知られています。研究者はデータロード、オーギュメンテーション、ログ記録、評価、クラスタジョブ管理を同時に管理しなければならず、小さなミスが長時間の実行を無効にする可能性があります。このプロジェクトは、PyTorch Lightning をラップしたフレームワークであり、その全体的なワークフローをより簡単で柔軟かつ安定したものにします。研究者はモデルと損失に集中でき、フレームワークが下請けの作業を処理します。

どう動くか

このフレームワークは、Python の辞書を相互に渡す4つのコンポーネントで構成されています：DataModule（データ）、Module（モデル＋フォワードパス）、Callbacks（モニタリング/評価フック）、Trainer/Manager（オーケストレーション）。重要なアイデアは、フォワード関数が単に {'loss': ..., 'embedding': ...} のようなテンソルの辞書を返すだけで、ログ記録、評価プローブ、チェックポイントなどはその辞書から読み取ることで、トレーニングループを変更する必要がないということです。多くの自己教師学習手法（SimCLR、DINO、MAE、BYOL、CLIP など）の事前構築された「フォワード」レシピが付属しており、OnlineProbe（固定された特徴量上で小さな線形分類器を訓練して精度をリアルタイムで追跡）や OnlineKNN（訓練なしの最近傍評価）などのコールバックも提供されています。また、kornia を使ってデータオーギュメンテーションを GPU 上に移動し、CPU のボトルネックを解消しています。さらに、Manager は SLURM クラスタ機能（ジョブの再キュー/再開、アトミックチェックポイント、照会可能な実行レジストリ）を提供します。実験的な JAX/Flax-NNX バックエンドも同じ設計を反映しています。

対象ユーザー

ファウンデーションモデルや自己教師学習の研究を行う研究者やエンジニア — 特に GPU クラスタ上で大規模な視覚モデルを訓練し、リアルタイム評価、堅牢なチェックポイント、最小限のボイラープレートを求める人。また、SimCLR、DINOv2、MAE などの SSL メソッドを最小限のコードで素早くプロトタイピングまたはベンチマークしたい人にも役立ちます。

特徴

• 辞書ベースの設計：フォワード関数が状態辞書を返すため、任意の中間テンソルが自動的にログ可能になり、コールバックはトレーニングループを触らずにアタッチできます。
• 30 以上の組み込みレシピ：SSL、教師あり、マルチモーダルの前処理をカバー（SimCLR、DINO/DINOv2、MAE、BYOL、VICReg、Barlow Twins、LeJEPA、CLIP など）。
• リアルタイム評価コールバック：OnlineProbe や OnlineKNN など、トレーニング中に表現品質を監視します。
• GPU 上のバッチ化オーギュメンテーション（kornia による）：バッチ全体でオーギュメンテーションをベクトル化し、モデルサイズや精度に応じて測定されたスループット向上を実現。
• SLURM 水準のオーケストレーション：Manager がプリエンプション/再キュー、アトミックチェックポイント、照会可能な実行レジストリを処理。
• 実験的な JAX/Flax-NNX バックエンド：torch の設計を反映し、数値的主張に対して同等性の回帰テストを実施。