Research map / Optimization / Optimizers
Scalable Meta-Learning via Mixed-Mode Differentiation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MixFlow-MG, a reparameterisation that lets bilevel/meta-gradient computations use forward-over-reverse mixed-mode differentiation instead of default reverse-over-reverse, cutting dynamic memory by up to ~10x and wall-clock time by up to 25% on Chinchilla transformers.
Problem
Gradient-based bilevel optimisation requires gradient-of-a-gradient computations through truncated inner loops, which are memory- and compute-expensive under default reverse-over-reverse autodiff at large model scale.
Approach
Reparameterise the inner update so inner gradients are a separate argument, then use Hessian symmetry (Schwarz's theorem) to rewrite vector-Hessian and vector-mixed-derivative products as Hessian-vector products computed in forward-over-reverse mode, avoiding storage of inner backward activations; also saves inner gradients as checkpoints and uses block rematerialisation.
Key results
- Peak dynamic HBM reductions reach over 10x (exceeding 90%) in a 135-configuration sweep on GPUs and TPUs, with about 75% memory reduction for 80% of configurations.
- Wall-clock step time wins reach 25% on GPU and 20% on TPU, with a median improvement of 12%.
- Scaling-ladder gains converge to 23-25x dynamic memory reduction on TPUs and 10x on GPUs for models up to 16B parameters.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace default autodiff with mixed-mode differentiation | Bilevel language-modelling setups · Language modelling task sweep, GPU and TPU | peak dynamic HBM: over 10x | Helped |
| Replace default autodiff with mixed-mode differentiation | Bilevel language-modelling setups · Language modelling task sweep, GPU | wall-clock time: 25% | Helped |
| Replace default autodiff with mixed-mode differentiation | Bilevel language-modelling setups · Language modelling task sweep, TPU | wall-clock time: 20% | Helped |
| Add mixed-mode differentiation and block rematerialisation | Chinchilla transformers · Chinchilla scaling ladder, TPU | peak dynamic HBM ratio: 23-25x | Helped |
| Add mixed-mode differentiation and block rematerialisation | Chinchilla transformers · Chinchilla scaling ladder, GPU | peak dynamic HBM ratio: 10x | Helped |
| Add saving inner gradients to mixed mode and block rematerialisation (ablation) | 489M Chinchilla transformer · GPU | HBM (G): 174.8 -> 54.8 | Helped |
| Add saving inner gradients to mixed mode and block rematerialisation (ablation) | 489M Chinchilla transformer · TPU | Time (s): 5.13 -> 4.12 | Helped |
| Add saving inner gradients to mixed mode and block rematerialisation (ablation) | 44M Chinchilla transformer · GPU | Time (s): 1.51 -> 1.19 | Helped |
| Add saving inner gradients to mixed mode and block rematerialisation (ablation) | 44M Chinchilla transformer · TPU | HBM (G): 12.7 -> 12.9 | Hurt |
| Add block rematerialisation to mixed mode and saved inner gradients (ablation) | 489M Chinchilla transformer · GPU | HBM (G): 289.2 -> 54.8 | Helped |
| Add block rematerialisation to mixed mode and saved inner gradients (ablation) | 489M Chinchilla transformer · TPU | HBM (G): 176.8 -> 46.9 | Helped |
| Add mixed-mode differentiation to block rematerialisation and saved inner gradients (ablation) | 489M Chinchilla transformer · GPU | HBM (G): 182.4 -> 54.8 | Helped |
Surprising findings
- Modern XLA compilers do not find these memory savings on their own, so the default implementation stays suboptimal despite compilation.
- Saving inner gradients is crucial for GPU memory but only matters for step time on TPUs.
Limitations
- Gains depend on compiler and hardware behaviour and vary with library versions and flags.
- Only exact gradients are considered; approximate methods are not evaluated.
- Static memory is not reduced, so total memory gains shrink to 4-6x for large models.
Experimental setup
Metrics: peak dynamic HBM ratio step time ratio
Baselines: default reverse-over-reverse autodiff implementation
Builds on: MAML truncated backpropagation through time JAX
Closest papers by meaning
- Tuning-Free Bilevel Optimization: New Algorithms and Convergence Analysis (ICLR 2025 · similarity 0.86)
- qNBO: quasi-Newton Meets Bilevel Optimization (ICLR 2025 · similarity 0.86)
- A Nearly Optimal Single Loop Algorithm for Stochastic Bilevel Optimization under Unbounded Smoothness (ICML 2024 · similarity 0.86)
- Bilevel Optimization under Unbounded Smoothness: A New Algorithm and Convergence Analysis (ICLR 2024 · similarity 0.85)
- TSP: A Two-Sided Smoothed Primal-Dual Method for Nonconvex Bilevel Optimization (ICML 2025 · similarity 0.85)
- Efficient Curvature-Aware Hypergradient Approximation for Bilevel Optimization (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.