PyTorch learns to differentiate through its own optimizers

What it does

TorchOpt is a PyTorch add-on that makes optimization steps themselves differentiable. It provides functional optimizers (think Optax for JAX, but PyTorch-native) and three differentiation modes—explicit, implicit, and zero-order—for bilevel optimization where an outer loss depends on the result of an inner optimization. The library also bundles CPU/GPU-accelerated optimizer kernels, RPC-based distributed training, and fast tree operations for manipulating nested parameter structures.

The interesting bit

The functional API is the real pivot: by treating optimizer state and updates as pure functions rather than in-place mutations, TorchOpt lets you backpropagate through entire training loops without hacking PyTorch’s autograd. You can write meta-learning code that looks like ordinary PyTorch or like JAX-style functional programming, depending on your tolerance for parentheses.

Key highlights

• Three differentiation modes for bilevel optimization: explicit gradient over unrolled steps, implicit differentiation via fixed-point conditions, and zero-order gradient estimation.
• Functional optimizer API compatible with functorch, plus a familiar torch.optim-style wrapper for traditional PyTorch users.
• Differentiable by default: pass inplace=False and the optimizer emits non-mutating updates that preserve the computation graph.
• Distributed training support and custom CPU/GPU kernels for the tree operations that parameter flattening/unflattening requires.
• Includes visualization tools and tutorial notebooks for meta-learning patterns like MAML.

Caveats

• The README is heavy on API examples and light on performance benchmarks; claims of “efficiency” are stated but not quantified.
• Implicit differentiation and zero-order modes are mentioned but the detailed mechanics are truncated in the source; users will need to consult the paper or tutorials for the full picture.

Verdict

Researchers and practitioners in meta-learning, hyperparameter optimization, or learned optimizers should look here—especially if they want JAX-like functional patterns without leaving PyTorch. If your training pipeline is standard supervised learning with no nested optimization, this is overkill.