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 process 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.
View on arXiv@article{kemaev2025_2505.00793,
title={ Scalable Meta-Learning via Mixed-Mode Differentiation },
author={ Iurii Kemaev and Dan A Calian and Luisa M Zintgraf and Gregory Farquhar and Hado van Hasselt },
journal={arXiv preprint arXiv:2505.00793},
year={ 2025 }
}