Scalable Meta-Learning via Mixed-Mode Differentiation

2025-05-01

AI Summary

Gradient-based bilevel optimization—widely used in meta-learning—suffers from high computational overhead and scalability bottlenecks due to frequent computation of second-order and mixed derivatives. To address this, we propose MixFlow-MG, the first method to systematically integrate mixed-mode automatic differentiation into the gradient construction process for meta-learning. Our approach combines computational graph rewriting with memory-aware backpropagation, significantly improving the efficiency of higher-order derivative computation. Compared to standard implementations, MixFlow-MG reduces memory consumption by over 10× and accelerates end-to-end training by 25%, while preserving accuracy on mainstream meta-learning benchmarks. By enabling efficient, scalable differentiation for structured bilevel optimization, MixFlow-MG establishes a new paradigm for differentiating nested optimization problems in deep learning.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Mixed Discrete/Continuous Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

Abstract

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.

Problem

• Efficiently compute gradient-of-a-gradient in bilevel optimization
• Reduce computational cost of second-order and mixed derivatives
• Improve scalability and performance in meta-learning setups

Innovation

• Mixed-mode differentiation for efficient gradients
• Scalable computational graphs for meta-learning
• Memory and time improvements in optimization