みなさん、こんにちは。AutoMuon という小さな Python パッケージを作っていて、Muon オプティマイザを、任意の PyTorch の学習パイプラインにおける AdamW のドロップイン置き換えとして使えるようにしています。
中核となる考え方は比較的シンプルです。Muon は主に隠れ状態に対する 2D 重み行列（線形射影、畳み込み層）に対して動作しますが、それでも埋め込みや正規化、バイアスなどには AdamW が必要です。AutoMuon は初期化時にモデルをスキャンし、各パラメータに対して適切なオプティマイザを自動的に割り当てます。
PR（プルリクエスト）歓迎です。特に、アーキテクチャでエッジケースに当たった場合は、モジュールタイプの除外リストを拡張することがあると思うので、その対応に協力してもらえると嬉しいです。トランスフォーマーや CNN 以外の何かで試している人がいるなら、結果をぜひ知りたいです。完全にカスタムされたアーキテクチャ、たとえば flash-linear-attention みたいなものではうまくいきにくい気がするので、その場合はユーザー側のチューニングが必要になりそうです。
時系列予測、ゲノミクス、言語モデリングなどのためのテストをさらに追加する予定です。本当に Muon はどれくらい汎用的なのかを見てみたいです！
https://github.com/SkyeGunasekaran/automuon
pip install git+https://github.com/SkyeGunasekaran/automuon.git
[リンク] [コメント]
