arXivに2026年3月26日付で公開された学術論文「Second-Order, First-Class: A Composable Stack for Curvature-Aware Training」（Mikalai Korbit、Mario Zanon著）は、Optaxベースのスタックとして「Somax」という二次最適化ライブラリを構築しています。同論文は「Optaxのような最適化ライブラリは学習を変換の連鎖として構成するが、典型的な二次最適化の実装はソルバー中心またはフレームワーク固有のインターフェースを通じて機能を公開する」と指摘し、Optaxの変換連鎖という設計思想自体は二次最適化にそのまま適用しにくいという制約を明らかにしています。SomaxはOptaxの勾配変換フレームワークと組み合わせて曲率スケーリングされた方向を計算した後、Optaxの変換チェーンへ通す設計とすることで、この制約を回避しつつOptaxとの互換性を保っています。これは第三者による独立した学術研究であり、Optax自体の公式ドキュメントによる主張ではありません。