Scalable Meta-Learning via Mixed-Mode Differentiation

We thank the Reviewer for the thorough feedback and for pointing out the parts of the paper that weren’y sufficiently clear. We will incorporate out responses below into the final version to improve clarity.

I was not aware that HVP can be more efficiently computed with the forward-over-reverse mode…I still don't fully understand the technical details

Using FMD for HVPs avoids having to store some intermediate activations (which standard RMD uses). This yields the exact same gradient, but often with much less memory and compute spent, especially for the case of gradient checkpointing (Sections 2.2, 3, A.1). We will refine the narrative in the final version to consolidate all the details in one place.

Just in case, what if you add a simple experiment (in the appendix) showing that the proposed algorithm produces the exact meta-gradient?

The proposed algorithm is theoretically identical (up to machine precision) to the well-studied BPTT methods. Assuming that JAX and PyTorch correctly compute meta-gradients, we double-check this by validating that all our MixFlowMG runs produce the same numerical results as the baselines, with up to 1e-6 absolute and 1e-4 relative errors, which are common for 32-bit floating point arithmetic.

Q1

Indeed, Eq.(6) reveals the terms (VHPs and VMPs) that are calculated using FMD in the proposed algorithm. The crucial part here is that the rest of the terms are still calculated using backprop , or RMD, which makes MixFlowMG a mixed-mode differentiation method (hence the name – Mixed-Mode Meta-Gradients).

fwdrev_grad in Alg.(2) is the result of calling get_fwdrev_grad_fn in f A.4: it’s a function for calculating the gradient of inner loss $L$ w.r.t. inner parameters $\theta$, but with custom differentiation rules that constitute MixFlowMG. These rules instruct JAX to compute VHPs and VMPs as more optimal HVPs and MVPs – this is why MixFlowMG can be implemented in just a few lines of code, yet yielding substantial performance wins in practice. Precisely the same mechanism is used in the Pytorch implementation for MixFlowMG, which we will add to the final version of the paper.

Q2

To the best of our knowledge, highly-rated publicly available libraries for meta-learning (e.g. [1], [2]), indeed, use the parameterisation from Eq.(3) and implement meta-gradients as outlined in Alg.(1).

In our paper, we demonstrate that by introducing a simple reparameterisation Eq.(4) and using a slightly modified yet theoretically equivalent Alg.(2), those libraries can easily adopt MixFlowMG and benefit from significant compute and memory gains. To re-emphasize, Eq.(3) + Alg.(1) is theoretically equivalent to Eq.(4)+Alg.(2), as shown in the paper.

Essentially, the proposed re-parameterisation in Eq.(4) allows for revealing the blue and red terms in Eq.(6), which are implicitly calculated using RMD in the autodiff frameworks. Explicitly surfacing these terms in the update equations allows for calculating them via FMD, which is the core idea of MixFlowMG. Hope this clarifies the contribution of the paper!

[1] Arnold, Sébastien MR, et al. "learn2learn: A library for meta-learning research." arXiv preprint arXiv:2008.12284 (2020), https://github.com/learnables/learn2learn

[2] Metz, Luke, et al. "Practical tradeoffs between memory, compute, and performance in learned optimizers." Conference on Lifelong Learning Agents. PMLR, 2022, https://github.com/google/learned_optimization

Q3

Yes, this is possible. The primary goal of MixFlowMG is to reduce memory footprint per single inner step (i.e. the dynamically allocated memory), while increasing T demands more statically allocated memory (Section 4 and Figure 2). That being said, dynamic memory gains are directly convertible to extra static memory: for example, if the size of parameters and states is $O(P)$, and the dynamic memory reduction thanks to MixFlowMG is of order $O(WP)$, then T can be scaled as $O(W)$ in a naive way. Typically, for smaller $P$ and/or higher compute intensity per parameter this win factor is higher, as e.g. in Section 3.2.

We described more efficient ways to scale static memory in A.3, all of which are fully compatible with MixFlowMG. Those include distributing static memory over $d$ devices, effectively reducing it by $d$ times, which worked for us.

Q4

The setup from Section 3.2 is essentially MAML: it is a BLO because the outer parameters $\eta$ define the initialisation for the inner parameters $\theta_0 = \eta$; after updating $\theta$ for T steps, we calculate the validation loss on $\theta_T$ and propagate the gradients for $\partial V / \partial \eta$ back to $\theta_0$, to perform one outer update.

Parameter $M$ controls the size of a single inner update in this toy example Eq.(9): for larger $M$, one inner step requires storing more intermediate activations for RMD. Varying $M$ in Figure 1 demonstrates how the gains from using MixflowMG scale with the size of the inner updates.

We thank the Reviewer for the actionable feedback. We add clarifications and respond to the questions below.

The submission claims that the reverse-mode differentiation performed by the default autodiff can be suboptimal in practice. However, the underlying reasons for this drawback are not clearly explained … Equations (7) and (8) claim that the transport versions of VHP and VMP lead to more efficient computation; however, this is not clearly explained.

Just to be sure, we would like to clarify that the claim about inefficiency of the reverse-mode differentiation is only related to the bilevel gradient optimisation (BLO) setups. We thank the Reviewer for pointing out that this wasn't sufficiently clear.

Essentially, the transposing of the Hessian allows us to use forward-mode differentiation here, which avoids having to store intermediate activations (which standard reverse-mode differentiation uses). This results in the exact same gradient, but often with much less memory and compute spent, especially for the case of gradient checkpointing. The reason we can do it here is that the BPTT-based BLO methods always happen to end up with symmetric Hessian there, whose transpose is trivial.

Although we discuss each of these parts in the paper (Sections 2.2, 3, A.1), we will refine the narrative in the final version to consolidate all the details in one place to improve clarity.

Only three bilevel optimization-based methods are evaluated using the proposed hypergradient computation strategy… Lack of comparison with the recent methods focusing on LLM…

We agree that trying the proposed algorithm in more specific instances of BLO is an important direction. In particular, the combination with LLMs is especially interesting, because these setups already are often scaled all the way to what could possibly fit on device memory.

In this paper, we showed that MixFlowMG yields exactly the same gradients (up to machine precision) as the existing Backpropagation-Through-Time (BPTT)-based algorithms in BLO, thus ensuring that their experimental results are fully transferable to MixFlowMG. Taking this into account, we decided to devote this paper to the practical aspects of MixFlowMG, i.e. its computational and memory benefits.

Our method can be used to already reduce the computational requirements of existing highly performant BPTT-based approaches, as well as to unlock previously-infeasible setups which would otherwise have required too much memory. For such examples (scaling inner-loop length, LLMs), we’d like to kindly ask the Reviewer to have a look at our other responses due to the size limit.

The implicit gradient-based algorithms are not discussed in the paper.

Implicit gradients are another BLO strategy with different pros and cons that are worth research in their own right and which we consider orthogonal to our specific contribution. For the readers who would like to learn more about the broader scope of the BLO, we included references to the broad surveys, e.g. [1], which have detailed sections on implicit gradients. We will add a separate paragraph mentioning alternatives to BPTT to the final version and include the provided references.

Equation (5) is not rigorous, as it sets the second derivative term to zero.

Typically, the validation loss $V(\theta_T)$ does not depend on the internal optimiser’s state $\upsilon_T$, hence we set the mentioned derivative to zero in the corresponding equations. We will amend them to include this term in the most general form. Note that this does not affect the paper’s results and analysis, as all the derivations remain correct.

The submission focuses on efficient hypergradient computation in bilI believe there is no clear connection to broader scientific research beyond the scope of meta-learning.

Our main focus is indeed BLO and efficient hypergradients. These methods are already impactful, e.g., in robotics, but we believe they will become increasingly important in AI research and applications, and so having efficient algorithms for this is important.

That said, there are some surveys [1] which connect hypergradients with research areas beyond meta-learning (e.g. signal processing, adversarially robust training). MixFlowMG applies anywhere where a similar to meta-gradient structure appears, and similar computational gains can then be obtained there.

The main deep learning and meta-learning implementation trend is still PyTorch. However, this submission only demonstrates the algorithm on <...> JAX.

We re-implemented the toy example in PyTorch and confirmed that the gains of MixFlowMG remain significant (up to 4x less memory on GPUs on this particular task). We will include this Pytorch implementation (~15 lines of code) in the final version.

[1] Zhang, Yihua, et al. "An introduction to bilevel optimization: Foundations and applications in signal processing and machine learning." IEEE Signal Processing Magazine 41.1 (2024): 38-59

We thank the Reviewer for the thorough feedback and address the raised concerns below.

While the experiments demonstrate the practical benefits of the proposed algorithm, the paper could be strengthened by including more experimental performance of the proposed method solving specific problems. … However, the paper does not show the concrete problem formulation that can be solved by the proposed method.

We agree that trying the proposed algorithm in specific instances of BLO is an important direction. In this paper, we showed that MixFlowMG results in exactly the same gradients (up to machine precision, Section 3 and Appendix A.2.) as the existing Backpropagation-Through-Time (BPTT)-based algorithms, thus ensuring that their experimental results are fully transferable to MixFlowMG. Taking this into account, we decided to devote this paper to the practical aspects of MixFlowMG, i.e. its computational and memory benefits.

We would like to highlight that the proposed algorithm can be combined with a wide variety of meta-learning algorithms, with the same guarantee of equivalence (up to numerical precision), as well as with future algorithms that have not yet been discovered.

To share our enthusiasm about the experimental direction, we would like to highlight that several works, e.g. [1], [2], [3], reported positive effects of increasing the number of inner steps (up to 10 in [1], [2], and up to 100 in [3]) in the BPTT-based meta optimisation. Those results suggest that increasing this parameter by over 10 times with MixFlowMG is likely to follow this positive trend potentially leading to qualitatively new outcomes. That being said, to ensure the well-defined scope of the paper, we leave such investigations to future works and instead devote the experimental section mainly to the thorough technical analysis of the proposed algorithm.

[1] Figure 5 in Finn, Chelsea, Pieter Abbeel, and Sergey Levine. "Model-agnostic meta-learning for fast adaptation of deep networks." International conference on machine learning. PMLR, 2017

[2] Table 2 in Antoniou, Antreas, Harrison Edwards, and Amos Storkey. "How to train your MAML." arXiv preprint arXiv:1810.09502 (2018)

[3] Figure 2a in Rajeswaran, Aravind, et al. "Meta-learning with implicit gradients." Advances in neural information processing systems 32 (2019)

The paper can be strengthened by including the experiment result on the application of BLO on LLM.

To the best of our knowledge, gradient-based bi-level optimisation for LLMs and especially practically strong BPTT-based algorithms still remain under-explored because of the prohibitive computational costs. Recently, it has become an emerging trend: for example, [1] uses meta-gradients with LLMs in the inner loop for the task of dynamic token-level loss reweighting at the fine-tuning stage. Despite using only 6 inner-loop steps, the authors reported significantly improved F1 scores on downstream QA tasks.

That setup can be seamlessly scaled further with MixFlowMG, in terms of both model size and inner-loop length. We believe that its memory- and compute efficiency gains is a critical step towards the adoption of meta-gradients for LLM-related tasks. That is why, in the experimental part of the paper, we conduct a detailed analysis of the properties of MixFlowMG in common meta-learning workflows enhanced with real-world LLMs (the Chinchilla scaling ladder).

[1] Hu, Nathan Zixia, et al. "Meta-Learning Online Adaptation of Language Models." The 2023 Conference on Empirical Methods in Natural Language Processing.

In application of MAML, how many inner-level steps are applied? If it is one, what is the benefit of the proposed approach?

We benchmarked the method on the number of inner-level steps T=2,4,6,8 for all tasks, including MAML. The results can be found in Figure 5(c), which shows that for the largest (4.5B) model, our method results in approximately a 10x reduction in memory requirements, and this relative gain stays fairly constant across different numbers of inner steps. We expand on this in Sections 4 and 5.3 of the paper (i.e. static and dynamic allocation patterns).

We did observe significant memory and compute gains for T=1 as well, as 1-step hypergradients, despite not having vector-by-Hessian products (VHPs), still contain vector-by-mixed-derivatives-matrix products (VMPs) which are optimised in MixFlowMG (Proposition 3.1.). We decided to not include it in the plots because the corresponding points look like outliers compared with T>1, which may be confusing for the readers; we will comment on T=1 in the final version of the text.

We appreciate the feedback of the Reviewer and address the raised concerns below.

The only issue I see is a lack of the meta learning performance analysis. As far as I know, different meta learning algorithms lead to (sometimes vastly) different final results.

We agree with the Reviewer that the landscape of meta learning algorithms is truly vast. In this paper, we focus on the class of gradient-based backpropagation through time (BPTT) optimisation methods and, in particular, on their computational and memory aspects. The proposed algorithm MixFlowMG helps to scale the existing practical setups up by more than an order of magnitude, whilst remaining theoretically identical (up to machine precision) to the well-studied standard BPTT methods, as we demonstrated in Section 3 and Appendix A.2.

In particular, this implies that the relevant results from the works that use this class of algorithms are fully transferable to MixFlowMG. To double-check this, we validated that all our benchmarks produce the same numerical results as the baselines, with up to 1e-6 absolute and 1e-4 relative errors (which is within the machine precision for 32-bit floating point arithmetics).

We would like to highlight that the proposed algorithm can be combined with a wide variety of meta-learning algorithms, with the same guarantee of equivalence (up to numerical precision), as well as with future algorithms that have not yet been discovered.

… it would be nice if meta-learning algorithm implemented with their method actually improves the final performance.

We would like to highlight that several works, e.g. [1], [2], [3], reported positive effects of increasing the number of inner steps (up to 10 in [1], [2], and up to 100 in [3]) in the BPTT-based meta optimisation. These results, combined with the theoretical equivalence of MixFlowMG to the presented algorithms, keeps us optimistic that increasing this parameter by over 10 times with MixFlowMG is likely to follow this positive trend, potentially leading to qualitatively new outcomes. That being said, to ensure the well-defined scope of the paper, we leave such investigations to future works and instead devote the experimental section mainly to the thorough technical analysis of the proposed algorithm.

[1] Figure 5 in Finn, Chelsea, Pieter Abbeel, and Sergey Levine. "Model-agnostic meta-learning for fast adaptation of deep networks." International conference on machine learning. PMLR, 2017

[2] Table 2 in Antoniou, Antreas, Harrison Edwards, and Amos Storkey. "How to train your MAML." arXiv preprint arXiv:1810.09502 (2018)]

[3] Figure 2a in Rajeswaran, Aravind, et al. "Meta-learning with implicit gradients." Advances in neural information processing systems 32 (2019)

I believe this paper sits between ML and systems paper. Given that ICML has been traditionally closer to ML conference, I am not sure how system contributions are received in this avenue.

We consider the main contribution of this paper to be a novel algorithm to reduce the computational requirements of meta-gradients based on mixed-mode differentiation.

ICML conference lists the following category in the call-for papers https://icml.cc/Conferences/2025/CallForPapers: Machine Learning Systems (improved implementation and scalability, hardware, libraries, distributed methods, etc.). We believe that this work falls under this category and that it can be of interest to the ML community, as it can be implemented and used by people right away. We will also include an implementation of MixFlowMG in PyTorch in the final version of the paper.

Some examples of similar conference papers, including ones from the past ICML conferences:

• Archer, Aaron, et al. "Practical Performance Guarantees for Pipelined DNN Inference." International Conference on Machine Learning. PMLR, 2024. https://proceedings.mlr.press/v235/archer24a.html

• Biderman, Stella, et al. "Pythia: A suite for analyzing large language models across training and scaling." International Conference on Machine Learning. PMLR, 2023. https://proceedings.mlr.press/v202/biderman23a.html

• Patil, Shishir G., et al. "POET: Training neural networks on tiny devices with integrated rematerialization and paging." International Conference on Machine Learning. PMLR, 2022. https://proceedings.mlr.press/v162/patil22b/patil22b.html

• Dao, Tri, et al. "Flash Attention: Fast and memory-efficient exact attention with io-awareness." Advances in neural information processing systems 35 (2022): 16344-16359. https://openreview.net/pdf?id=H4DqfPSibmx

• Stanton, Samuel, et al. "Kernel Interpolation for Scalable Online Gaussian Processes." International Conference on Artificial Intelligence and Statistics. PMLR, 2021. https://proceedings.mlr.press/v130/stanton21a.html