Scalable Meta-Learning via Mixed-Mode Differentiation
This work addresses scalability issues for researchers and practitioners in meta-learning, hyperparameter optimization, and related fields, representing an incremental improvement in computational efficiency.
The paper tackles the computational inefficiency of gradient-based bilevel optimization in meta-learning by proposing Mixed-Flow Meta-Gradients (MixFlow-MG), which uses mixed-mode differentiation to achieve over 10x memory reduction and up to 25% faster wall-clock time compared to standard implementations.
Gradient-based bilevel optimisation is a powerful technique with applications in hyperparameter optimisation, task adaptation, algorithm discovery, meta-learning more broadly, and beyond. It often requires differentiating through the gradient-based optimisation itself, leading to "gradient-of-a-gradient" calculations with computationally expensive second-order and mixed derivatives. While modern automatic differentiation libraries provide a convenient way to write programs for calculating these derivatives, they oftentimes cannot fully exploit the specific structure of these problems out-of-the-box, leading to suboptimal performance. In this paper, we analyse such cases and propose Mixed-Flow Meta-Gradients, or MixFlow-MG -- a practical algorithm that uses mixed-mode differentiation to construct more efficient and scalable computational graphs yielding over 10x memory and up to 25% wall-clock time improvements over standard implementations in modern meta-learning setups.