Scalable Meta-Learning via Mixed-Mode Differentiation
Authors: Iurii Kemaev
Dan A. Calian
Luisa M Zintgraf
Gregory Farquhar
Hado van Hasselt
Venue: ICML 2025 poster (2025)
Keywords: meta-learning
bilevel optimization
second-order optimization
automatic differentiation
autodiff
mixed-mode differentiation
gradient-based methods
scalable algorithms
Primary area: deep_learning->algorithms
TL;DR: 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.

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.

Scalable Meta-Learning via Mixed-Mode Differentiation
Turii Kemaev'
Dan A. Calian
Luisa M. Zintgraf'
Gregory Farquhar' Hado van Hasselt
Abstract
optimisation (Bengio, 2000; Franceschi et al., 2018), data weighting (Hu et al., 2023b; Calian et al., 2025), and fast task adaptation (Finn et al., 2017), to neural architecture search (Liu et al., 2018), adaptive reinforcement learning (Xu et al., 2018: Zahavy et al., 2020), algorithm discovery (Oh et al., 2020), and more.
Gradient-based bilevel optimisation is a pow- erful technique with applications in hyperpa- rameter optimisation, task adaptation, algorithm discovery, meta-learning more broadly, and be- yond. It often requires differentiating through the gradient-based optimisation process itself, leading to "gradient-of of-a aaa--rrddentt calculations with com- putationally expensive second-order and mixed derivatives. While modern automatic differentia- tion libraries provide a convenient way to write programs for calculating these derivatives, they oftentimes cannot fully exploit the specific struc- ture of these problems out-of-the-box leading to suboptimal performance. In this paper, we anal- yse such cases and propose Mixed-Flow Meta- Gradients, or MixFlow-MG a practical algo- rithm that uses mixed- mode differentiation to con- struct more efficient and scalable computational graphs yielding over 10x memory and up to 25% wall-clock time improvements standard im- plementations in modern meta-learning setups.
In gradient- bilevel optimization. the meta-parameter update requires backpropagating through the inner loop, leading second-order derivatives (gradients of gradients) notoriously computationally expensive process both in terms of memory and FLOPs. Updating outer parameters ev- ery T inner steps (truncated backpropagation through time, Truncated- -PPT:; Werbos, 1990) still results in computa- tional cost scaling linearly with T. Consequently, we are often restricted small inner and outer models 0 and n and short horizons T, limiting the exploration of the full potential of BLO. While Truncated- MBTTT can be effective for smaller meta-models v) (Xu et al., 2018 3018; Shaban et al., 2019), its applicability to large neural networks with billions of parameters (Gemini et al., 2023; OpenAl et al., 2023) re- mains an open question. Moreover, given the demonstrated impact of scale on model performance (Kaplan et al., 2020; Hoffmann al., 2022b), the trend of scaling inner mod- cls 0 is likely to continue. This necessitates more efficient BLO algorithms to support modern and future generations of models and to explore larger backpropagation horizons T whilst keeping the cost of experiments affordable.
1. Introduction
Bilevel optimisation (BLO) is a commonly used tool to solve problems in meta-learning and deep learning (Liu et al., 2021; Zhang et al., 2024). In this problem setting, an inner- loop optimisation of parameters incrementally searches for optimal values 0*, in a process that depends on (fixed) meta-parameters 1). In an outer-loop neta--ppimiistion, search for optimal meta- petaapparameters 17". For instance, 7) may include hyperparameters of the inner update (Bengio, 2000) or even their per-weight versions (Sutton, 1992).
In this paper, we first analyse standard implementations for Truncated- Muucated-BBPTT-beae BPTT-based bilevel gradients modern frame- works for automated differentiation and highlight their in- herent inefficiencies. We then propose Mixed-Fil -low Meta- Gradients, or MixFlow-MG a simple reparameterization of the inner-loop learning dynamics that exposes the under- lying symmetry of the problem and uses mixed-mode auto- matic differentiation to seamlessly exploit it. Finally, we use modern hardware and libraries for tensor programming to demonstrate that the proposed algorithmic technique, whilst requiring only minor code modifications, yields significant performance improvements in common meta-learning sce- narios. In representative setting, MixFlow-MG demon- strates reductions up to 95% in the active memory con- sumption and 25% reduction in wall-clock time, thus allowing to scale bilevel gradient setups by more than an order of magnitude in compute-effiiier way.
This framework offers a powerful approach to automating the design and optimisation of learning systems, leading to significant advancements in various machine learning domains. It has applications ranging from hyperparameter
Google DeepMind. Correspondence to: Jurii Kemaev <iuke- macy @ @ww.g@ggooooo......
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
While numerous approximations for the (Truncated-JBPTT- based numerical procedure, such as implicit (Rajeswaran et al., 2019; Lorraine et al., 2020; Blondel et al., 2022; Choe et al., 2023) and fooe gradients (Silver et al., 2021; Shen et al., 2024) have been proposed recently, we focus on calculating exact gradients to isolate and address the core computational bottlenecks. The presented ideas can be seamlessly incorporated into approximate methods as well.
and reverse. Forward mode calculates lacobian-by-vectte product (JVP) Ju with arbitrary vector a computational cost proportional to single forward pass (Baur & Strassen, 1983). By carrying out JVPs with n imput's 's basis vectors, the full Jacobian J can be recovered colunn-by-column, hence requiring O(n) forward passes in total. Reverse mode, on the other hand, computes vector-by-Jacober product (VJP) vJ and recovers the Jacobian one row time, in total requiring (mm) forward passes for computing the full Jaco- bian; however, by design, it operates in two passes - forward and backward - and requires storing all intermediate activa- tions during the forward pass to use them in the backward pass, resulting in significantly higher memory requirements.
2. Background
2.1. Bilevel Optimisation
In the general form, BLO can be posed as the following constrained optimisation problem:
For neural networks, typical objects for differentiation are loss functions L IR" R that output scalars. This is why reverse mode is the default choice, since it recovers the whole Jacobian Jz in O(1) forward passes.
min (nn), where V(n) = EE V (0)()))))
(1)
s.t. 8 (n) arg min E X L(0, 17, x)
(2)
8
Certain classes of differential programs, such as those imple- menting second -oddr optimisation in Section 2.1, require computing products with second-order derivatives of the corresponding loss functions. One example is Hessian-by- vector products (HVPs) 8 L/Ox²v. HVPs can be cheaply evaluated using repeated VJP and/or JVP products (Pearl- mutter, 1994), and there are three computationally tractable ways available in practice:
where are the outer meta- perta-prrametess are the inner model parameters, V and L are validation and train losses calcu- lated on the data points y ~ Y and x X, respectively, Note that standard network training regimes are special case, where the validation loss in Equation (1) is minimised by tuning the meta-parameters v) by hand.
Typically 8* (n) in Equation (1) is approximated with T steps of gradient- based methods (Maclaurin et al., 2015):
LL
(OL
JOL 0x )
VJP(C.JVP(L..)
min V (n) == E = v))
Ox2
dx
" (0++1: 1++1 (0,, Us, 17, 2..
(3)
reverse-over-forwarr
O
aL
JVP(VJJ((,,L).))
where U/ is an arbitrary state at step i, such as an optimiser's momentum, and Φ(0,, Us 17, 2.) is an update that involves calculating the gradient (L. 17, 200, and is differen- tiable by 7). This ensures that meta-parameters 1,, in their turn, can also be optimised with gradient methods, giving rise to quantities involving second-order derivatives of the loss function L(0,7) 17,2) 2). In particular, such schemes require computing left- or right-hand side products of the second- order derivatives with arbitrary vectors.
dx
0x
forward-over-reverse
a
0L 0x
V.P((c,VJP(e.LL)
U Ōr
reverse-over-reverse
where V R" is an arbitrary vector and e (1) is a unit vector. Note that forward-over-forwal mode was purpose- fully omitted due to its prohibitive computational cost of O(z) forward passes.
2.2. Primer on Automatic Differentiation
One crucial observation is that forward-over-revers mode avoids storing activations from the inner backward pass, of- ten making it the most memory efficient choice in practice, In addition, it has lower I/O overhead (no need to read/write activations), potentially leading reductions in wall-clock time. This advantage becomes even more apparent when calculating 0T./Ox relies on the gradient checkpointing tech- nique (Griewank & Walther, 2000), as it is effectively a no- op for forward -ode differentiation. This property forms the core of a highly efficient algorithm for gradient- based BLO which will be described further in the paper.
A convenient way to compute the quantities involving second-order derivatives in Equation (3) is provided by mod- ern automatic differentiation libraries such as JAX (Brad- bury et al., 2018) or PyTorch (Paszke et al., 2017). This section explains fundamental concepts upon which these libraries are built, which is important for understanding how to design efficient algorithms for solving equations (3).
Let us consider arbitrary f(z): R" + IM with the corre- sponding Jacobian J J= =0/āāz mxx, Autodiff provides two types of differentiation for such functions: forward
2

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
2.3. Automatic Differentiation in BLO
defined in the code. By exploiting their hidden symmetry in a non-intrusive way, we achieve substantial memory sav- ings with minimal code changes in the user programs, Our benchmarks show that modern compilers are not able to make such improvements on their own.
Several works (e.g., Franceschi et al., 2017) explore the trade-offss between forward- and reverse- mode differentia- tion for meta-pate perttataramttesss..ee 7) in gradient-basee based BLO in Equa- tion (3). In this context, the computational cost of reverse- mode differentiation for the validation loss Jv with respect to 77 is comparable to the computation of the inner- and outer losses themselves, i.c. O(brl) + ((1))), but requires storing all intermediate activations in memory that are not necessary in forward-mode differentiation. However, this memory efficiency comes at the price of the increased com- putational cost, which becomes (11) x [@1). The overall consensus is that when the number of meta-paraman parameters is small, one should consider forward-mode differentiation to avoid incurring extra memory costs. However, in many modern applications (Finn et al., 2017: Wichrowska al., 2017: Oh et al., 2020), (n] can be much larger, often even comparable in size to the number of model parameters [@] (e.g., 1) are themselves the parameters of neural network, as in Finn et al. (2017)), thus making forward- -odderd-mode mode differ- entiation prohibitively expensive. In such cases, a straight- forward implementation uses reverse mode at both levels of the corresponding program, as shown in Algorithm All of the open-source repositories for gradient-based meta- optimisation that we verified use this implementation (e.g., Arnold et al., 2020; Metz et al., 2022). We therefore re- fer to Algorithm 1 as a standard or default autodiff-based implementation further in the text.
3. MixFlow-MG: Mixed-mode Differentiation for Bilevel Gradients
We are now going to decompose equations (3) in order to expose the Hessian matrix. This will allow us to employ a more memory-efficient algorithm for calculating the outer gradients.
Firstly, we propose reparameterising equations (3) have gradients VL; = (L, 77, x00/ as a separate argument in the combined update function T:
((+1) $(0,, vi, v;, 3;) = T(V L,, O,, Us. U;,11 v.). (4)
After applying the chain rule to the gradient of the validation loss with respect to and using the fact that the validation loss V in (3) does not depend on the last-step state uy, we obtain (in vector notation)
dV 0v d0r
0V dut
= o)
dTr dr
+
(5)
dnj 00T dn dur dn 30T
Then, after unrolling one step for YH+1 (Appendix A.2), we get the following recurrent relation for the total derivatives:
The aforementioned autodiff libraries (Paszke et al., 2017: Bradbury et al., 2018) compile the computational graph that is defined by a user's program (in just-in-time or dy- namic regime) before executing it. Compilation allows for leveraging advanced optimisations and memory allocation algorithms, which can make all three differentiation modes theoretically equivalent in terms of compute- and memory efficiency in many use cases (Dagréou et al., 2024; Blondel & Roulet, 2024).
d' dn
0r..1 LLL
01
d'.
(00
00;
0v
00?
Ovs
dn
(6)
0T
LT.
0m.1
+ On
LV 01000
Equation (6) allows to unroll the loop "backwards", from = T 1 to O. According to Equation (5), for calculating dV /dn it needs to be multiplied by the vector /00r 0) from the left, hence it only requires one VJP. However, it can be noticed that Equation (6) contains explicit vector-by- hessian and vector-by- by-toe- -eriiidderivasivess the default autodiff implementation will perform them in leverse-over-reve mode, which can be suboptimal in prac- tice.
Nevertheless, their corresponding hidden constants can sub- stantially differ in practice. The practical differences stem from various factors, including: the underlying model's structure, inputs' size, autodiff framework and hardware backend, compiler's configuration and flags, custom op- timisations, etc. Moreover, these frameworks only have limited contextual information about given problem's 's na- ture, hence often failing capture and exploit inherent symmetries and structural bottlenecks of the problem at hand, which can lead to suboptimal low-level programs in practice, as we will demonstrate further in the paper.
products;
To circumvent this, we transform the relation using classical results. Assuming that the function approximator and loss function have continuous second-order derivatives, which is typically the case for neural networks, the following identi- ties hold (c.f. Schwarz's 's theorem):
In the following section we identify an algorithmic improve- ment based on the fact that standard computational graph for bilevel gradients includes symmetric matrices, such as Hessians, which are never fully instantiated explicitly
82 L
0²L
LL T
8
003
30
0n00,
00,0n
3

---

be transformed MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
Algorithm 1 Standard Miinnaated-PPT1 (Equation (3))
Algorithm 2 Mixed-mode Truncated-BPTT (Equation (4))
Input: 7, lo. vo. inputs (x:) val
Input: l. vo. inputs {xt) val
Output: @vjōn
Output: av
2: function @(0, U,,,,,
2: function T(VL,0, 2,
3: VL (O,
3: empty line
4: arbitrary operations on 0, VL
4: ...arbitrary operations on 0, 0,00
5: (0,v) + optimizer((, L))
5: (0, v0,v) + opttmizer(( 00)
6: ...biratrrryy operations on 8
6: ...ariitrary operations on 8
7: return 0,
7: return v
8:
8:
9: function VaLLoss(1, do vo. {ae)(-1. val x))
9: function VALLOss(1, 8., Vo. faelT- val x)
10: (0,0) + (fo, uo)
10:
(0, v) + (fo, vp)
for i T do
for
do
empty line
V fwdrev (O,
13: (0 (0, 2,)
13:
(0,0) T(VL, n.
14: return (0, xx)
14: return V(O, val_
16: 0V galarr Oo. un. (o val_r)
0V (aad(ValLoss)))) Oo. un, (xr) val_x)
17:
return 0V
17: return 0V
Combining them with Equation (6) can rewrite the vector-by-han by-hssian (VHP) and cector-by-mixed-deriiatiii matrix (VMP) products into their transposed versions, i.e. hessian-by-veetee (HVP) and vector (MVP) products:
of cases (Dagréou al., 2024; Blondel & Roulet, 2024), their practical performance vary remarkably, which we demonstrate in Section 5. In general case, is recommended trying all three options for choosing the best one for setup at hand; the proposed reparameterisation (4) makes this probing straightforward. In Section 5 we demonstrate how MixFlow-Me leverages for significant performance gains practice.
matrix-bya----
normal VJP
differentiation
0T (+1 OL LV 28?
L 082
oY i+1 Li
3.1. Implementation in JAX
inefficient VHP
efficient HVP
JAX (Bradbury al., 2018) is a powerful library for dif- ferential tensor programming. relies the functional paradigm, where stateless functions and returned by other functions; one of its key transforma- tions is grad(f) which accepts a scalar- valued function f(x) IR= - and returns a new function g(r) : TR R that computes the gradient of f with respect to x, i.c. g(z) = 0f/a(=).
normal VJP T
0Y LL OVL, 0n00,
L
oY
00,0m
OVL;
inefficient VMP
efficient MVP
Proposition 3.1. Reparamererisation (4) and identities (7), (8) allow for replacing the default peerrse--ve---eeers dif- ferentiation for recurrent relation (6) with more efficient forwand-over-reverse or enverse-vverffowaate alternatives.
The default autodiff-based implementation this conve- nient notation for computing second-order derivatives in the original training loop (3), as shown in Algorithm 1. This however can be highly suboptimal as it fails to exploit the problem's 's inherent symmetries, as discussed in Section 2.3. Our proposed reparameterisation (4), outlined in Algorithm 2, allows use mixed-mode differentiation via custom fw- drev -rred transformation, which defines VJP rule for cal- culating HVPs in forward-over-rewent mode. This requires changing only few lines of code; our implementation can be found in Appendix A.4.
Since mixed-mode differentiation constitutes the algo- rithmic improvement in this technique, we call it Mixed- Flow Meta-Gradients or MixFlow-MG.
While advanced autodiff compilers and memory allocation algorithms can make all three differentiation modes equiv- alent in terms compute- and memory efficiency most

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
50
mixed mode
0.5
mixed
400
HBM footprint with allocations (GiB)
default
350
Peak HBM (GiB)
6x
0.3
300
lower HBM
6x
faster
10
200
dynamic memory
07
150
20
40
60
20
40
60
Loop length
Loop length
100
Figure 1. Peak HBM and step time across the number of (per inner step) transformations M in Equation (9) (GPU).
50
static memory
0
program instruction
1e4
3.2. Motivating Example
To illustrate the effects of MixFlow-MG, we consider the following simple BLO problem (Equation (3)), similar to Finn et al. (2017): 17 defines the initialisation point On for the inner optimisation; the inner loss is a standard 1.2 loss which independent of yy the update dynamics (0,, Us, x.)) a standard stateless (u Ø) gradient step.
Figure 2. Device memory footprint for update when using
autodiff for one step of bilevel optimisation. The memory can be divided into static (checkpoints, inputs, parameters, states) and dynamic (activations for backpropagation). The dynamic memory can be reduced by exploiting the problem structure (see Figure 3). X-axis: instruction number in compiled computation graph. Y-axis: amount occupied device memory.
The inner model yau is the following M step recursive map:
(0) (2 1 1),
(9)
O(T (IA)) + + |uD). While it can be affordable for small setups, real-world models are too large to be adapted for meta-training this way due to high cost and scarcity of the high bandwidth on-device memory: typically, one inner step already uses all available on-accelerators memory.
where 3o (0, Ox. BxD, _ 0 DxD We used B = 1024 and D 4096 in our experiments and kept the number of inner updates T 2 for simplicity. The compu- tational graph gets longer as the number (per inner step) transformations M increases, meaning we can study the ef- fects on memory and runtime by adjusting M. For the sake of demonstration, minimised the effects of compiler's optimisation by disabling all loop fusions.
Gradient checkpointing (Griewank Walther, 2000) for inner steps is often used in practice (e.g. in Arnold al. (2020); Metz et al. (2022)) to bring the memory footprint down to O[ T (l0) + lu])), since only activations for the current step are kept in memory at any moment of time during meta-backpropogationn ters are getting checkpointed during the outer-loop unroll. Typically the size of activations and partial derivatives |A] is substantially larger than the size of parameters and opti- miser states [8] + Jul due to the dependency of the former both the latter and inputs' sizes. This makes gradient checkpointing instrumental for scaling, and following this common practice, we keep enabled in all our benchmarks.
Figure demonstrates how the metrics change across the number of per-step transformations M. The HBM and step- time scales much better when using mixed-mod differenti- ation, with memory and wall-clock reductions up 85% as M increases. The corresponding code and XLA--geneated computational graphs given in Appendices A.6 and A.7.
and 7 ( (u|) parame-
4. Scaling to Large Models
This section investigates device memory patterns and mem- ory optimisation techniques in gradient- based bilevel opti- misation for the case when the underlying models get larger.
One important distinction to make is that checkpoints, in- puts, parameters, and states require O(T-(10)+|o) bytes that get allocated at the beginning of a program for the whole execution time and written to only once. For this reason, we refer to this type of memory static. On the contrary, O(IA)) bytes are allocated during outer-level backpropaga- tion and re- re-purposed for new activations at every inner step, hence we refer to it as dynamic memory. Typical memory footprint for single outer step can found in Figure 2.
The standard implementation of Truncated-BPPP for BLO (Algorithm !) loops over T inner updates & to obtain 0y for calculating the outer (validation) loss. If done naively, this algorithm requires storing intermediate activations A, and outputs ee. v for each of 1..7' inner steps, hence the peak memory consumption for one meta update scales
5

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
HIM
Static
ailic
0.5
3.5
0.
3.5
25
1.0 Program Instruction
2.5
355
Program Instruction
164
Program Instruction
l6d
1ed
Program Instruction
1el
Figure 3. HBM footprints for each stage of optimisation for 489M chinchilla transformer on GPU: 1 block rematerialization, 2 saving inner gradients, 3 3-mmin mixed-mode differentiation. Each optimisation is responsible for reducing specific chunk of HBM.
In addition to enabling gradient checkpointing for inner loop unrolling, found the following two optimisations impor- tant for amplifying the gains of the proposed algorithm:
ference in the compilers' efficiency for these two backends. More details this can be found in Appendix A.8.
5. Benchmarking Language Modelling Tasks
1. Block rematerialisation: for neural networks with block-residual connections 2;+1 = fx;;)) ) z4, such residual networks (He et al., 2016) and transformers (Vaswani. 2017), gradient checkpointing can be applied to each of the blocks fi to substantially reduce memory footprint at the theoretical cost of one forward pass; this is a known optimisation, hence we keep enabled for both baseline and the proposed method avoid running out of memory even for smallest networks.
The primary goal of this section to demonstrate the bene- fits of MixFlow-MG on a representative set of BLO setups. Without limiting generalisation, we chose the language mod- elling domain for the inner-level optimisation, where the standard loss is the next-token prediction loss NTP(68,r) We use the Chinchilla family of language models (Hoff- mann et al., 2022a) with RoPE (Su et al., 2024) and the Adam optimiser (Kingma, 2014). When a meta model is present, we use the same architecture as for the inner model.
cext-token-prediction
2. Saving inner gradients: aL/30 can be saved (in addi- tion to per-inner step inputs and parameters 0) as part of inner-loop gradient checkpointing to avoid incurring one extra backward pass during the outer-level gradient propagation; we have not found this optimisation in previous works and existing libraries, hence it can be considered as an additional contribution of this paper; we enable only for MixFlow-MG.
Firstly, we explain the rationale behind choosing the perfor- mance metrics. Then, we select a diverse class of real-world problems to demonstrate possible gains in practice. Further, we investigate different properties of MixFlow-MG using various model structures and data regimes. Finally, com- bining all findings, we provide practical recommendations on efficiently scaling bilevel gradient sctups by orders of magnitude beyond any existing frameworks.
Benchmarking was performed in JAX (Bradbury al., 2018) TPUv5p and H100 using the OpenXLA back- end and libraries from DeepMind et al. (2020). While we observe consistent behaviour across setups and tasks, results may vary depending on library versions, hardware, compiler flags, and other factors beyond the scope this work.
Both these optimisations plus mixed-mode diffcrentiation, as introduced in Section 3, complement each other. Fig- 3 include the ablation study for 489M Chinchilla model (on MAML; our full benchmark setup is described in the next section). In particular, block rematerialisation under forward-over-rewerse differentiation does not require storing intermediate per-block checkpoints thanks to the forward mode at the outer level. This allows to almost completely remove block #3 in Figure 3 thus drastically reducing peak memory usage. Note that some portion of extra memory is still claimed for calculating activations and JVPs on-the-fly. fly, this is why forward mode differentiation typically requires 3 times more memory than the basic forward pass.
We listed the minimal changes required for implement- ing MixFlow-MG in Section 3.1 and included the relevant Python code for JAX and Py in Appendix A.4.
5.1. Metrics
MixFlow-MG operates on per-inner-step basis, i.e. it ad- dresses dynamic memory. In our metrics we focus solely on dynamic memory and defer to Appendix A.3 for practical recommendations on how to reduce static memory.
We also observed that, while saving inner gradients is crucial both for memory and step-time reductions on GPUs, was only important for the latter on TPUs, which shows the dif-
6

---

MixFlow-MG: Scalable Mcta-Learning via Mixed-Mode Differentiation
14
Peak dynamic HBM ratio (per task)
1.25
12
Step time ratio (per task)
1.20
10
task
8
1.15
loss weightigg
Ir r_learming
maml
6
1.10
platform
H100
1.05
TPUv5p
2
default
default
0.2
0.4
0.6
0.8
0.0
0.2
0.8
1.0
Quantile
Quantile
Figure 4. Joint sweep over tasks, models, and hyperparameters from Table I: peak dynamic HBM and step time ratio between default and mixed-mode differentiation, sorted by value in descending order (higher better, and values 1 mean that MixFlow-Me improves over the default autodiff implementation). All variations win both memory- and compute-wise. with highly correlated gains between tasks.
We consider the following three common BLO setups:
Table I. Sweep over tasks: hyperparameters and values.
Parameter
Values
Hyperparameter Learning: similar to Bengio (2000) and Maclaurin et al. (2015), are the per-pann parameter learning rates for the inner optimiser, so that
Task
(learning maml, loss, weightin)) (57, 106, 163, 217, 306)
Model size (x10%)
of inner updates (7)
2, 4, 8]
Batch size
12 4,8
Sequence length
{2048, 4096, 8192]
0,+1 aN) O,, v:).
with a function that includes optimiser's transforma- tions for converting gradients into parameter updates.
We measure peak dynamic High Bandwidth Memory (HBM) (device memory) and wall-clock step time. Where more appropriate, we report two performance metrics which defined a ratio the corresponding measurements be- tween the default implementation and the proposed changes, i.e. higher values indicate stronger gains over the baselines.
Model-Agnostic Meta-Learning (MAML, Finn et al. (2017)): 7 defines the initialisation point wi 1) for the inner optimisation and the inner loss is otherwise independent of
Peak dynamic HBM ratio is the ratio between the peak usages of dynamic HBM (Section 4)
L(O;,n,2:)= NTP(O,,
Meta-learning Adaptive Loss Weighting: inspired by Hu et al. (2023a), this setup uses calculate per-data point loss weighting factors:
HBMaetukt HBM
M.sttic static default
(10)
H B MMii HMM
static
MixFlow-MG
Step Time ratio the ratio between wall-clock time per
I(O,, n.).) a(1,2;) NT).
meta step
Taefaull
MixFloo MG-
(11)
We sweep over the hyperparameters in Table 1, totalling in 135 distinet configurations per task and sort all results by gains, descending order. Figure 4 shows memory gains and step-time wins for the runs that fit in available device memory (80 GiB for GPU and 96 GiB for TPU).
5.2. Sweeping over Bilevel Optimisation Tasks
To recap, typical setup for the gradient-based based BLO Equa- tion (3) is comprised of an inner loop that updates model parameters 0 for T' steps, and an outer loop that updates by backpropagating @V (ön through the inner loop steps by unrolling it backwards; the particular dependence of the inner-loop optimisation on 7 defines the problem setup.
MixFlow-MG delivers substantial improvements across the board. We observe that memory footprint and step time are reduced for all hyperparameter combinations. Remarkably, memory usage decreased by approximately 75% (nearly

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
10
Peak dynamic HBM ratio
8
model_sh size
106M
6
6-
587M
1.08
4.5B
2000
4000
6000
8000
2
Sequence length
Batch
Number inner updates
Figure Sweep over data regimes for Chinchilla models (GPU): peak dynamic HBM ratio between default and mixed-mode diff-n.
Pcak dynamic HBM ratlo
20
platform
H1.00
10-
Tpu
10 Number MHA heads
1000
2000
Number of blocks
Attention size
Feed-forward layer sizee
Figure Sweep over transformer components: peak dynamic HBM ratio between default and mixed-mode differentiation.
4x less memory) for 80% of configurations, with peak reduc- tions exceeding 90% (over 10x less memory) on both GPUs and TPUs. Previously, memory constraints severely limited the scale of bilevel optimization. These results open the door to training models of much larger size and complexity. Wall-clock time wins reach 25% for GPU and 20% for TPU, with a median improvement of 12% for both.
only a single memory buffer for activations, as opposed to number of blocks buffers for the default implementation.
This difference in memory scaling leads substantial reduction in peak dynamic HBM usage, quantified by the ratio:
BL(S+ L(1+kS) B + +k52)
(12)
1 Ks
Wall-clock gains are almost uniform across configurations, while memory gains vary significantly. We investigate this in the following, and disentangle factors contributing to the memory behaviour to showcase MixFlow-MG's properties.
The factor L in the enumerator ensures that MixFlow-MG is an algorithmic improvement for models with block-residual connections, such as residual networks and transformers.
5.3. Model and Data Scaling
To validate this theoretical estimate, we benchmark combi- nations of transformer models, context lengths, batch sizes, and number of updates. In the previous section observed that MixFlow-MG shows highly correlated gains across all tasks, so we report metrics only for the MAML setup here.
The dynamic memory requirements of transformer models using the default implementation scale ([BL(S+ASS)) where L is the number of layers, S the context length, B the batch size, and & compiler-dependen constant. This scal- ing arises from the self-attention blocks and holds even with block rematerialisation enabled the default implementa- tion. However, as detailed in Section 4, our proposed mixed- mode differentiation with block rematerialisation offers significant advantage: its memory usage is independent of where
Figure 5 shows the gains for different models, batch sizes B, context lengths S, and inner-loop lengths for GPUs, with similar dynamics observed for TPUs. These empirical results closely align with Equation (12): discounting minor compilation effects, the gains are constant across B and T and sub-linearly increase towards /T./ for larger S.
the number of layers, scaling only D(B(+kk:*) k represents the corresponding constant for mixed-mode gra- dients and is significantly smaller than &. This reduction stems from the forward-over-reverse mode, which requires
The impact of scaling different architectural components of Chinchilla transformer is shown in Figure 6. The memory gains scale linearly with the number of layers confirming
8

---

MixFlow-MG: Scalable Mcta-Learning via Mixed-Mode Differentiation
25
algorithm
6000
mixed-mode
Peak dynamic HBM (GIB)
default
Peak dynamic HBM ratio
20
5000
piatform
platform
H100
4000
15
TPUv5p
3000
10
2000
1000
defoult
2
10
12
14
0.0
2.5
5.0
7.5
10.0
12.5
15.0
Model size, billions
Model size, billions
Figure 7. Chinchilla scaling ladder: peak dynamic HIBM gains across transformers of various sizes.
theoretical analysis. While the gains could be expected to be near-constant for the other structural parameters, the real numbers differ in practice, especially for small models on GPUs. This can be attributed to compilation effects: the smaller a computational graph, the more memory optimi- sations low-level compiler can find limited time, e.g. GPUs may be able schedule the fixed-size thread warps more efficiently for small graphs.
meta-leaming setups. Importantly, as the domain models become larger and more sophisticated, the positive effect of using MixFlow-MG compounds, allowing to drastically reduce scaling costs. We believe that our work will help to facilitate research gradient- baadd bilevel optimisation by opening up larger class of models for experimenting whilst minimising the need for extra computational resources. We included minimalistic implementation in JAX and PyTorch for MixFlow-MG in Appendix A.4 for reference and easy adoption.
In the real world transformers simultaneously scale across all components (Hoffmann et al., 2022a). Figure 7 shows the peak dynamic HBM gains across a reduced version of the original Chinchilla scaling ladder, with models ranging from 44M 16B parameters. We observe that the gains get larger for bigger models, eventually converging 23-25x (96%) dynamic memory reductions for TPUs and 10x (90%) for GPUs. We hypothesise that the convergence happens due to the underspecified compiler's behaviour given the fact that starting from 1B transformers, the corresponding default computational graphs outgrow any available memory by more than one order of magnitude, which can too far from typical compilation targets.
A Acknowledgements
Authors would like express their deep gratitude to Jun- hyuk Oh, Matteo Hessel, Dan Horgan, the JAX, XLA, and RL teams, and David Silver for fruitful discussions and sup- throughout the project. We also thank the reviewers for the valuable feedback that helped to improve the clarity of the paper.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
6. Conclusion
In this paper, we examined the practical aspects of gradient- based methods for bilevel optimisations, identifying inef- ficiencies in default autodiff-based implementations. To address them, we proposed MixFlow-MG that uses mixed- mode differentiation for the most computationally de- manding components in the corresponding programs. We achieved this by introducing simple generic reparame- terisation technique that can be effortlessly integrated into standard implementations. We conducted detailed analysis of the proposed algorithm and identified its scaling proper- ties. Our empirical benchmarks demonstrated significant practical improvements, including up to 10x total mem- ory reductions and 25% lower wall-clock time in modern
References
Arnold, S. M. R., Mahajan, P., Datta, D., Bunner, I., and Zarkias, K. S. learn2leamm: A library for meta-learning research, 2020. URL https: httpp:s///xxxvoooo/ 2008 12284.
Baur, W. and Strassen, V. The complexity of partial deriva- tives. Theoretical Computer Science, 22(3):317-330, 1983. ISSN 0304-3975. doi: https://do..rg//0..0666 0304-3975(33)01 10-XX

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
Behrmann, J., Grathwohl, W., Chen, R. T., Duvenaud, D., and Jacobsen, J.-H. Invertible residual networks. In International conference machine learning, pp. 573- 582. PMLR, 2019.
Franceschi, L... Donini, M., Frasconi, P., and Pontil, M. Forward and reverse gradient-based hyperparameter opti- mization. In Precup, D. and Teh, Y. W. (eds.), Proceed- ings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learn- ing Research, pp. 1163-1173, PMLR, 06-11 Aug 2017.
Bengio, Y. Gradient- based optimization of hyperparameters Neural computation, 12(8):1889-13 188-1990, 2000.
Franceschi, L., Frasconi, P., Salzo, S., Grazzi, R., and Pontil, M. Bilevel programming for hyperparameter optimiza- tion and meta-learning. In International conference on machine learning, pp. 1568 1577. PMLR, 2018.
Blondel, M. and Roulet, V. The elements of differentiable programming, arXiv preprint arXiv:2403.14606 2024.
Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-López. F., Pedregosa, F., and Vert, J.-P. Efficient and modular implicit differentiation. Advances in neural information processing systems, 35:5230-5242, 2022.
Gemini, T., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang. Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: /tt//ggttttuuu com/ wtt. gowgle/ jax.
Griewank, A. Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods and software, 1(1):35-54, 1992.
Griewank, A. and Walther, A. Algorithm 799: revolve: an implementation of checkpointing for the reverse ad- joint mode of computational differentiation. ACM Trans- actions on Mathematical Software (TOMS), 26(1):19-45, 2000.
Calian, D. A., Farquhar, G., Kemaev, I., Zintgraf, L. M., Hessel, M., Shar, J., Oh, J., György, A., Schaul, T., Dean, J., van Hasselt, H., and Silver, D. Datarater: Meta-leamed dataset curation, 2025. URL httpt::t///// htt.t:ttt//aaeet httpbp:://rxxv.orrt 0bbs/2505.rm 17895.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn ing for image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Choe, S., Mehta, S. V., Ahn, H., Neiswanger, W., Xie, P., Strubell, E., and Xing. E. Making scalable meta learning practical. Advances in neural information processing systems, 36:26271-2622222
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022a.
2023.
Dagréou, M., Ablin, P., Vaiter, S., and Moreau, T. How to compute hessian-vector products? In ICLR Blogposts 2024, 2024.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbi, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. 2022b.
DeepMind, Babuschkin, I., Baumli, K., Bell, A., Bhupati- raju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hes- sel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I. King. M., Kunesch, M., Martens, L., Merzic, H., Mikulik, V., Norman, T., Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez, A., Sartran, L., Schnei- der, R., Sezener, E., Spencer, S., Srinivasan, S., Stano- jević, M., Stokowiec, W. Wang. L.. Zhou, G., and Vi- ola, F. The DeepMind JAX Ecosystem, 2020. URL http: //tpt//ggttth/////ee com. google google-deepmind deepmind.
Hu, N., Mitchell, E., Manning, C., and Finn, C. Meta- learning online adaptation of language models. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods Natu- ral Language Processing. Association for Computational Linguistics, December 2023a.
Hu, N., Mitchell, F.,, Manning, C. D., and Finn, C. Meta- learning online adaptation of language models. arXiv preprint arXivv2305. 2023b.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th Inter- national Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1126- 1135. PMLR, 06 Aug 2017.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001 083611, 2020.
10

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14 2014.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Per- formance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0444 2020.
Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. Meta-learning with implicit gradients. Advances neural information processing systems, 32, 2019.
Liu, H., Simonyan, K., and Yang, Y. Darts: Differentiable architecture search. arXiv preprint 2018.
Shaban, A., Cheng, C.-A., Hatch, N., and Boots, B. Trun- cated back-propagation for bilevel optimization. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 1723- PMLR, 2019.
Liu, R., Gao, J., Zhang. J., Meng, D., and Lin, 7. Investigat- ing bi-level optimization for learning and vision from a unified perspective: survey and beyond. IEEE Trans- actions on Pattern Analysis and Machine Intelligence, 44 (12):10045-10067, 2021.
Shen, Q., Wang, Y., Yang, Z., Li, X., Wang, H., Zhang, Y., Scarlett, J., Zhu, Z., and Kawaguchi, K. Memory- gradient unrolling for large-scale bi-level optimization. arXiv preprint arXiv:2406.. 14095, 2024.
Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing mil- lions of hyperparameters by implicit differentiation. In International conference on artificial intelligence and statistics, pp. 1540-1552. PMLR, 2020.
Silver, D., Goyal, A., Danihelka, I., Hessel, M., and van Hasselt, H. Learning by directional gradient descent. In International Conference Learning Representations, 2021.
Maclaurin, D., Duvenaud, D., and Adams, R. Gradient- based hyperparameter optimization through reversible learning. In International conference on machine learn- ing. pp. 2112-2122. PMLR, 2015.
Su, J., Ahmed, M.. Lu, Y., Pan. S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Mangalam, K., Fan, H., Li, Y., Wu, C.-Y., Xiong, B., Feicht- enhofer, C., and Malik, J. Reversible vision transformers. In Proceedings of the IEEE/CVF Conference on Com- puter Vision and Pattern Recognition, pp. 10830-10840, 2022.
Sutton, S. Adapting bias by gradient descent: An incre- mental version of delta-bar-delta. In AAAI, volume 92, pp. 171-176. Citeseer, 1992.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Metz, L., Freeman, C. D., Harrison, J., Maheswaranathan, N., and Sohl-Dickstein, J. Practical tradeoffs between memory, compute, and performance in learned optimiz- ers. In Conference on Lifelong Learning Agents (CoL- LAs), 2022. URL http://tt//// thub. com/ google/ learned opt imization.
Werbos, P. J. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):1550- 1560, 1990.
Wichrowska, O., Maheswaranathan, N., Hoffmann, M. W., Colmenarejo, S. G., Denil, M., Freitas, N., and Sohl- Dickstein, J. Learned optimizers that scale and general- ize. In International conference machine learning, pp. 3751-3760. PMLR, 2017.
Oh, J., Hessel, M., Czarnecki, W. M., Xu, Z., van Hasselt, H. P., Singh, S., and Silver, D. Discovering reinforcement learning algorithms. Advances in Neural Information Processing Systems, 33:1060-1070, 2020.
Xu, Z., van Hasselt, H. P., and Silver, D. Meta-gradient reinforcement learning. Advances in neural information processing systems, 31, 2018.
OpenAI, T., Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint 2xX1v:2303.08774, 2023.
Zahavy, T., Xu, Z., Veeriah, V., Hessel, M., Oh, J., van Hasselt, H. P., Silver, D., and Singh, A self-tuning actor-critic algorithm. Advances in neural information processing systems, 33:20913-20924, 2020.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z.,, Lin, Z.,, Desmaison, A., Antiga, L.,, and Lerer, A. Automatic differentiation in pytorch. NIPS 2017 Autodiff Workshop, 2017.
Zhang, Y., Khanduri, P., Tsaknakis, I., Yao, Y., Hong, M., and Liu, S. An introduction to bilevel optimization: Foun- dations and applications in signal processing and machine learning. IEEE Signal Processing Magazine, 41(1):38- 59, 2024.
Pearlmutter, B. A. Fast exact multiplication by the hessian. Neural Computation, 611):177-160, 1994. doi: 10.1 neco. 1ec0.1994.0 6.4.6..... 1.177.

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
A. Appendix
A.I. Author Contributions
Iurüi Kemaev: MixFlow-MG concept. project leadership, algorithm and benchmarks design, implementationn, and analysis;
Dan A. Calian, Luisa M. Zintgraf, Gregory Farquhar: algorithm analysis, benchmarks design, testing implementation; Hado Hasselt: advising the project, algorithm refinement and analysis.
All authors contributed to paper writing.
A.2. Derivations for MixFlow-MG
To expose second-order derivatives in the update equations Equation (3), we propose reparameterising them have gradients VL AL(O,, 77 //000, separate argument in the combined update function T:
(0s-1,0++1 ((,,, Us 7/ 2:)) T(V Li, O2 v4, 1) 2;).
(13)
After applying the chain rule to the gradient of the validation loss with respect to and using the fact that the validation loss V in (3) does not depend on the last-step state tr, we obtain
d
dV
0v d0r
0v
dur
ov
o)
drj der
drj
00T
dn
Out
dn
00r
dr
(14)
0
dYr
30T
drj
To calculate this total derivative, let us unroll one step for for T.+1:
d
dVL;
d0,
0Y;., du
0Y
dn
LVI,
dn
00; dn
Ou; dn
On
01
[ d0;
LL 0m00
0T d0;
0r du,
0T
OVL;
003 dnj
00, dn
Ov: du
On
01
01
8²
d8,
0 dua
( 0LL
+ 01,11
+ @VL. 0n00,
00,
LVL;
002
dn
Ov; dr
On
Rewriting this in the block-matrix form results
d0c.1 dn dus+1 dn
d0;
18T.11 0r.,. 0L
0r
drj
0T
0L
0T
(15)
00;
ov 002
Ou;
dv:
8v La
0n000
On
dn
Or, alternatively,
(P+0)xO
d dn
0T.1
0r.-1 LL
0r
d';
aY
8L
81
(16)
00,
aV 002
Ou;
dn
0md0,
On
(P+O)xP
(P+0)xM
where P. O, and M are the sizes of 0, v. and y correspondingly.
12

---

MixFlow-MG: Scalable Mcta-Learning via Mixed-Mode Differentiation
30
Peak HBM by types (GPU), GiB
5000
algorithm
mixed-mode
Peak dynamic-to St-ttaii HBM ratio
algorithm
10-
platform
H100
default
Peak total HBM ratio
TPUv5p
6000
platform
H100
3000
TPUvSp
dynamic memory
2000
static memory
default
0-
10
15
10
15
10
15
Model size, billions
Model size, billions
Model size, billions
(a) Peak dynamic and static HBM
(b) Dynamic/ static HBM ratio
(c) Peak total HBM ratio
Figure Static and dynamic peak HBM w.r.t model size.
Assuming that the function approximator and loss function have continuous second-order derivatives, which is typically the case for neural networks, the following identities hold (c.f. Schwarz's 's theorem):
0²L" T 002
8 66e
0L
T
0LL
mma0,
00.0n
Combining them with Equation (16) we can rewrite the vector-by-hesee (VHP) and vector- vecto-by-mixeedddriittvess (VMP) products into their transposed versions, i.e. hessian-by-vector (HVP) and mixedd-erivaiiive-mari--yyyy byyr---b-bbaii-----ttrrrrr (MVP) products:
normal VJP T
0Y LL 84 L, 08?
0LL T
(1
01 2 LVL.
0L
L²L 08?
0L 06?
0Y OV
since
08?
06?
inefficient VHP
efficient HVP
normal VJP T
8Y I LVL 0m00,
8...
( JVvi /)
T
8
8
0Y oV
since
im00,
On00:
a0,am
00,0m
inefficient VMP
efficient MVP
A.3. Handling static device memory
In the terminology introduced in Section 4, static device memory is used for storing inputs and outputs, parameters 6, states v, and checkpointed gradient and allocated at the beginning of the on-device computation for its whole lifetime. So far, the reported performance metrics reflected only changes in dynamic memory usage because MixFlow-MG operates on per-inner step basis, i.c. does not change the static memory allocations.
Figure ((a) shows dynamic and static memory distribution for the chinchilla scaling experiments from Section 5.3. As be seen, MixFlow-MG reduces dynamic memory by 10-25x, thus turning static memory into the dominating part of the allocated device memory. This gets exacerbated by the fact that, models and their optimizers' states become larger, the overall dynamic-tor -assticc ratio shrinks from 50-100 to 5-10 for default implementation, as depicted in Figure 8(b). In total, this causes peak HBM memory gains to reduce from 10-25x (Figure 7) to 4-6x (Figure 8(c)).
Fortunately, the static memory factor can reduced by several folds with the following techniques or their combinations:
13

---

MixFlow-MG: Scalable Meta-Lcarning via Mixed-Mode Differentiation
For distributed setups with D interconnected devices, static tensors be efficiently (i.e. with minimal communication overhead) distributed using Fully-Sharded Data Parallelism (FSDP) (Rajbhandari al., 2020), thus redueing the static memory allocation per device by D times.
For momentum-bassd inner-level optimisers, such Adam (Kingma, 2014), one can use the technique proposed in Maclaurin et al. (2015) to invert per-step updates during backward pass instead of storing them in static device memory; moreover, combining it with per-inner update remat from Section 4 can allow avoid computational overheads; the same principle holds for arbitrary optimisers applied to the class of reversible networks (Behrmann et al., 2019; Kitaev et al., 2020; Mangalam et al., 2022).
The default per-inner update rematerialisation strategy can be improved using dynamic programming (Griewank, 1992), allowing to reduce static memory usage from linear to logarithmic by ' (the number of inner updates per each outer update).
All of these techniques are fully compatible with MixFlow-MG and allow to achieve the 10-25x gains shown in Figure 7 with affordable (if not zero) compute overhead. We leave the implementation details of these methods outside the scope of this paper, as they can be found in the corresponding original works.
A.4. Python code for mixed-mode mode bilevel gradients in JAX and Py PyToch
Code Listing 1. JAX implementation for fwdree grad in Algorithm 2
def wwwdww_wrrreeeeeeeeeeeeeeeeeeeee fn """Returns function implementing (inner__lo_ss
fn)
The returned function has a custom-defined VJP rule for implementing forward-over-reverso mode for Hessian-by-vector products that emerge in the meta gradient / bilevel optimisation scenario. This custom rule substantial amount memory and compute compared with default JAX autodi ff.
Args:
inner_ loss_ fn: function implementing inner loss calculation. It must accept 'params' the first argument.
Returns:
A function implementing 'grad (innorr_ooss___ea with a custom-defined VJ rule for forward-ovo--roverso Ressian-by-vector products. n**
8 eax.ccustm_vjjp
def twarev lateev__rra__dleeeeeeee (params, xinputs)
""*Pure implementation, ***
return jax.gra loss in) (parans, kinputs)
def
(parans, +inputs):
"**Forward pass implementation. *** return fwdrev_grad_fn (paraar (params, kinputs), (params, inputs)
def fwdrev_ gww._rrdd_ae. fn ba_kwaakdrd pass (residuals, ct):
"""Backward pass implementation. (params, inputs) grac_loss_fi jax. grad (inner_loss. fn, ragnums-range (len (inputs) 1)) hvp_ct = jax. jvp (lambda p: loss_ >inputs), (params,), ) ). return hvp_ct
residuals
wiweev_grrd_ff.deefre fwarev twarev_grad forward pass, fwdrev grad En backward pass
return fudrev_grad_fn
14

---

MixFlow-MG: Scalable Mcta-Lcarning via Mixed-Mode Differentiation
Code Listing PyTorch implementation for fiudrev_grad in Algorithm
def htt__fwdeevggrade (inner_loss_fn)
"""Returns à function implementing 'grad (inner_loss.
The returned function has a custom-defined VJP rule for impl emplementing forwardw pewaaroovvr-revers.. mode for Hessian-by-vectoor products that emerge in the meta gradient level optimisation scenario. This custom rule can save a substantial amount of memory \ compute compared with default implementation.
Args:
inner_loss_fn: function implementing inner loss calculation. It must accept 'params' the first argument.
Returns:
function implementing 'grad (inner_loss__eee with custon-defined VJP rule for reree Hessian-by-vactor products. иим
class
@staticmethed
def forward (context, params, kinputs): """Fowward pass implementation. ***
savexxt.sav__oooe booooaaacwarddddoooooooaaearadde (parans, *inputs)
return torch. func. grad (inner loss_ fn) (params, sinputs)
@staticmethed
def backward (context, ct):
"""Backward pass implementation. **.
params, iinutss context saved_tensors
grad_ grad_loss_fn torch. func. func.grad (less, argnums-taple (range (1en (inputs) + 11))
hvp_ct torch. func. (lambda grad_loss_fis fn (p, xinputs), (params,), (ct,) ))
return hvp_ct
return axdpevGrad.apply
A.5. Python snippet for per-inner step gradient checkpointing with saving inner gradients
Listing Python snippet for optimisations in Section
def inner_step(... Implements one inner step.
params grad_ (params, inputs) d_parrms jax.tree www.tree.eeee max.trre.map((
lambda x: ott
www..cccit..ecceeeeeeeeeeeeeeeee (x, (ineer_grads'), d params
def outer_steple step(...) $ Implements the outer step. inner_step = jax, checkpoint inner_step, jelicy-jax.cheekpoinee wwwckcoonntpolliciiiiiiicc swve_only_theee_eeeeeooo (inner_grads') new_params, jax. lax. scan (inner_step, ...)
15

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
A.6. Python code for the motivating example
Code Listing 4. Python implementation for the motivating example
def get__tty__ttt (seed, B, M, T, D, use_ use_loop_fusion use_mi xed_mode)
""*Returns a toy task and example a arguments.
Args:
seed: a random seed.
8: a batch size.
M: number of inner steps.
T: number of inner updates.
D: data and inner model size.
se__loop___uste whether to use loop fusion.
use_mixed_moce whether to use mixed mode.
Returns:
jitted function with arguments that correspond the toy task.
rngl, rgg2, rng3 jax. random.spll it (jex. random. PRNGKey (seed), 3)
params jax. normal (rngl, (D, 1)) xs, targets jax. random. n (rng2, (2, T, D))
val val targe jax. random. nermal (rng3, (2, B, D))
def toy_task (params, xs, targets, val_x, val_target);
def apply (params, x): y - -jpp.mamuul (x, params)
def fy, i): return
jnp.sin(y)) cop (y),
if use
for i in range (1, M + 1): Y, fly, i)
else:
y, jax. lax scan (f, y, jnp. (1, M 1))
return y
def loss (params, x, target):
return imp.mean, (apply (params, x) target) x 2)
def meta_less (params):
if use_mixeed_mor _ode:
grad_fn get_fwxrev_grad_fhh11 (1oss)
else:
grad in fax, grad
def inner step (params, and target): d_params frad_fn (params, +x_and_target)
params jax.tree. map (lambda F, dp: p 1e-3 dp, params, params) return params, ()
params, jax. lax. scan (inner_stepp params, (xs, targets)) return loss (params, val_x, val_target)
gead = jax. grad (meta (parans)
return meta_grad
return jax. jit (toy_task), (params, xs, targets, val_x, val val_target)
16

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
Figure HLO graph for the motivational example. Data nodes are depicted orange, compute operations (multiplications, trigonemetric functions etc.) in gray. It can be seen that the mixed-mode version contains far fewer data blocks. Also, this example demonstrates the complexity of the underlying low-level programs and the huge role of compiler in optimising raw computational graphs.
A.7. Compiled computational graphs for the motivating example
See Figure 9.
A.8. Detailed ablations on all used optimisations
See Figure 10 and Table 2 for ablations on 489M model and Table 3 for step time measurements on 44M model, which fits into single-core device memory.
Table Case study for 489M transformer.
Optimisations
489M transformer
Mixed mode
Block remat
Save grads
GPU
TPU
(G)
Time (s)
HBM (G)
Time (s)
371.2
N/A
273.9
N/A
363.7
N/A
176.6
N/A
180.1
N/A
123.7
N/A
182.4
N/A
130.8
N/A
286.0
N/A
168.1
N/A
289.2
N/A
176.8
N/A
174.8
N/A
43.8
5.13
54.8
5.45
46.9
4.12
A.9. Sweeps over data regimes for TPUs
See Figure 11.
A.10. Models and hyperparameters
For sweeps in Figure we used hyperparameters from Table 4. For benchmarks in Figure we used the models from Table 5. For scaling plots Figure and Figure we used models from Table with batch size 4 and inner steps per outer update.
17

---

MixFlow-MG: Scalable Mcta-Learning via Mixed-Mode Differentiation
Table 3. Case study for 44M transformer.
Optimisations
44M transformer
Mixed mode
Block remat
Save grads
GPU
TPU
HBM (G)
Time (8)
HBM (G)
Time (s)
94.2
N/A
70.2
0.75
76.6
N/A
45.8
0.70
54.2
1.33
32.8
1.05
54.5
1,30
34.1
1,03
76.4
N/A
45.1
0.88
76.6
N/A
45.5
0.70
45.2
1.51
12.7
1.17
16.4
1.19
12.9
0.94
HBM footprint in GiB (GPU), default
save inner grads
HBM footprint in GIB (TPU), default
save inner grads
blcck ramat
block remat
392
Talse
Palse
252
202
150
150
50
mixed mode
save Inner grads mixed mode
mixed mode
Save inner grads mixed mode
152
250
302
150
50
Figure All combinations of the used optimisation from Section 4 for 489M model. Note that GPU required saving inner gradients for peak memory gains, while TPU needs it only for reducing step time. Mixed-mode diffcrentiation and model blacks' rematerialisations are critical for both cases.
Peak aymiii HBM ratio
20
20
model
15.
106M
278M
587M
10-
10
2.68
5
5-
4.58
2000
4000
6000
8000
2
Sequence length
Batch
Number inner updates
Figure 11. Sweep over data regimes for chinchilla models for TPUs. For GPUs see Figure 5. The results are more noisy due to TPU- specific optimisations memory layout (e.g. memory padding).
18

---

MixFlow-MG: Scalable Meta-Learning via Mixed-Mode Differentiation
Table Sweep over data regimes in Figure 5: hyperparameters, values, and descriptions. When plotting each of three per-axis plots, we used the maximum values for the other axes (e.g. for the sequence length plot used batches of size 8 with 8 inner updates).
Parameter
Values
Description
Model size (x10%)
(106, 278, 587, 1018, 2639, 4516)
Parameters in inner transformer
# of inner updates (1)
2. 4, 6. 8)
Inner updates per outer update
Batch size
2, 6,4,6,5 8)
Inner model's 's batch size
Sequence length
{1024, 2048, 4096, 8192)
Context length
Table 5. Chinchilla models used in sweeps over each of the components, Section 5,3
Sweep over
d model
Ifw size
kv size
n heads
layers
d_model
128
2048
1024
16-256
16
ffw_size
384
512-8192
32
8
16
n_heads
768
1024
24-384
2-32
16
n.layers
256
1024
32
4-64
Table 6. Chinchilla models from Hoffmaan et al. (2022a) used in scaling benchmarks, Section 5.3.
Parameters (million)
d_model
ffw_size
n_heads kv_size
n_layers
44
512
2048
64
8
g
90
640
2560
64
10
13
140
768
3072
64
196
896
3584
64
14
16
278
1024
4096
64
16
18
489
1280
5120
128
10
587
5632
724
1536
6144
128
12
22
1,018
1792
7168
128
14
23
1,429
2048
8192
128
16
1,609
2176
2,007
2304
9216
128
18
28
2,639
2560
10240
128
20
30
3,802
2816
11264
128
22
4,516
3072
12288
6,796
3584
14336
128
28
40
9,293
4096
16384
128
32
42
11,452
4352
17408
128
12,295
4608
18432
128
12.569
4608
18432
128
32
47
13,735
4864
19456
128
32
47
16,183
5120
20480
128
40
47
19


Reviews:
[{"review_id": "1cMAKbOf7R", "reviewer": "Reviewer_ZDJP", "summary": "This paper proposes to smartly combine forward-mode and reverse-mode differentiations to efficiently perform JVP/HVP in bi-level optimization. This is motivated by the observation that different differentiation methods can show massively different performance in practice, despite their potentially same theoretical memory/compute complexity. The proposed MixFlow-MG achieves over 10x memory and up to 25% wall-clock time improvements over previous implementations of meta-learning algorithms in multiple benchmarks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the feedback of the Reviewer and address the raised concerns below.\n\n\n> 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. \n\nWe 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. \n\nIn 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).\n\nWe 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. \n\n> … it would be nice if meta-learning algorithm implemented with their method actually improves the final performance.\n\nWe 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.\n\n[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\n\n[2] _Table 2_ in Antoniou, Antreas, Harrison Edwards, and Amos Storkey. \"How to train your MAML.\" arXiv preprint arXiv:1810.09502 (2018)]\n\n[3] _Figure 2a_ in Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" Advances in neural information processing systems 32 (2019)\n\n> 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.\n\nWe 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. \n\nICML 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.\n\nSome examples of similar conference papers, including ones from the past ICML conferences:\n\n* 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\n\n* 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\n\n* 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\n\n* 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\n\n* 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", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The main focus of the paper is improving the implementation of meta-learning (or bi-level optimization) algorithms. I believe all proposed implementation tricks intuitively make sense and smartly designed.", "theoretical_claims": "N/A: The paper doesn't make any theoretical claims besides providing derivations of their MixFlow-MG algorithm.", "experimental_designs_or_analyses": "Given that the main focus of the paper is improving practical implementation efficiency of meta-learning algorithms, I appreciated that authors have provided extensive analyses on run-time and memory usage. 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. Traditionally, it's known that scaling iterative-differentiation (as opposed to implicit differentiation) meta learning algorithms to large models/datasets is bottlenecked by prohibitive compute/memory costs. Since this paper addresses this issue, it would be nice if meta-learning algorithm implemented with their method actually improves the final performance.", "relation_to_broader_scientific_literature": "Bi-level optimization with recent large models, despite having lots of useful applications like hyper-parameter optimization, is largely limited by its compute/memory complexity. This paper tackles this issue, potentially providing new research opportunities.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Using HVP or JVP to enable efficient bi-level optimization is a popular idea itself. However, the key contributions of this paper are practical implementation tricks for efficient HVP/JVP.", "other_comments_or_suggestions": "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.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "9XbiOsAYoV", "reviewer": "Reviewer_85Hz", "summary": "The paper introduces MixFlow-MG, an algorithm that uses mixed-mode differentiation for more efficient gradient-based bilevel optimization. It reduces memory consumption and computation time compared to standard methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for the thorough feedback and address the raised concerns below.\n\n\n> 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.\n…  However, the paper does not show the concrete problem formulation that can be solved by the proposed method. \n\n\nWe 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.\n\nWe 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.  \n\nTo 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.\n\n\n[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\n\n[2] _Table 2_ in Antoniou, Antreas, Harrison Edwards, and Amos Storkey. \"How to train your MAML.\" arXiv preprint arXiv:1810.09502 (2018)\n\n[3] _Figure 2a_ in Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" Advances in neural information processing systems 32 (2019)\n\n> The paper can be strengthened by including the experiment result on the application of BLO on LLM.\n\nTo 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. \n\nThat 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).\n\n[1] Hu, Nathan Zixia, et al. \"Meta-Learning Online Adaptation of Language Models.\" The 2023 Conference on Empirical Methods in Natural Language Processing.\n\n> In application of MAML, how many inner-level steps are applied? If it is one, what is the benefit of the proposed approach?\n\nWe 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).\n\nWe 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.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria are appropriate for addressing the problem of high computational costs in gradient-based bilevel optimization.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The authors validated their approach through:\n\nInefficiency analysis of standard Truncated-BPTT implementations.\nReparameterization to expose problem symmetry.\nDemonstration of performance gains with mixed-mode differentiation.\nAnalysis of device memory footprint.\nEvaluation of optimizations (block rematerialization, saving inner gradients).\nScaling analysis of dynamic memory requirements.", "relation_to_broader_scientific_literature": "The paper proposes an efficient gradient computation for BLO.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n\nThe paper provides a clear explanation of the problem and the proposed solution, and employs illustrative examples and visualizations to aid understanding.\n\nThe paper includes extensive experiments on the efficiency of the proposed method.\n\nWeaknesses\n\nWhile 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.", "other_comments_or_suggestions": "It is interesting that the proposed method enables the application of BLO to LLM. However, the paper does not show the concrete problem formulation that can be solved by the proposed method. The paper can be strengthened by including the experiment result on the application of BLO on LLM.", "questions_for_authors": "In application of MAML, how many inner-level steps are applied? If it is one, what is the benefit of the proposed approach?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YD8dqRjnZO", "reviewer": "Reviewer_zzTN", "summary": "The paper introduces MixFlow-MG, a method for scalable bilevel optimization in meta-learning using mixed-mode automatic differentiation. The core idea involves reparameterizing the inner-loop dynamics to replace inefficient reverse-over-reverse differentiation with forward-over-reverse mode. This reduces computational costs and memory usage. Experiments demonstrate up to 10× memory savings and 25% faster wall-clock time across tasks like hyperparameter learning, MAML, and loss weighting. The method is validated on large models (up to 16B parameters) and hardware (GPUs/TPUs), showing consistent gains in dynamic memory and compute efficiency.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for the actionable feedback. We add clarifications and respond to the questions below.\n\n> 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.\n\nJust 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. \n\nEssentially, 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. \n\nAlthough 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.\n\n>Only three bilevel optimization-based methods are evaluated using the proposed hypergradient computation strategy… Lack of comparison with the recent methods focusing on LLM… \n\nWe 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. \n\nIn 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.\n\nOur 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.\n\n> The implicit gradient-based algorithms are not discussed in the paper.\n\nImplicit 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.\n\n>Equation (5) is not rigorous, as it sets the second derivative term to zero. \n\nTypically, 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.\n\n> 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.\n\nOur 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. \n\nThat 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. \n\n>The main deep learning and meta-learning implementation trend is still PyTorch. However, this submission only demonstrates the algorithm on <...> JAX.\n\nWe 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.\n\n[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", "claims_and_evidence": "1. 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 in the submission.\n\n2. Equations (7) and (8) claim that the transport versions of VHP and VMP lead to more efficient computation; however, this is not clearly explained.", "methods_and_evaluation_criteria": "The papers evaluate and compare the memory overhead introduced by the proposed algorithm and various baseline methods across different meta-learning algorithms. However, the performance of the bilevel optimization-based methods is not reported.", "theoretical_claims": "There are no theorems proposed in the submission.", "experimental_designs_or_analyses": "Only three bilevel optimization-based methods are evaluated using the proposed hypergradient computation strategy.\n\nLack of comparison with the recent methods focusing on LLM for instance [1].\n\n[1] Choe, Sang, et al. \"Making scalable meta learning practical.\" Advances in neural information processing systems 36 (2023): 26271-26290.", "relation_to_broader_scientific_literature": "The submission focuses on efficient hypergradient computation in bilevel optimization. I believe there is no clear connection to broader scientific research beyond the scope of meta-learning.", "essential_references_not_discussed": "The implicit gradient-based algorithms are not discussed in the paper. \n\n[1] Lorraine, Jonathan, Paul Vicol, and David Duvenaud. \"Optimizing millions of hyperparameters by implicit differentiation.\" International conference on artificial intelligence and statistics. PMLR, 2020.\n\n[2] Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" Advances in neural information processing systems 32 (2019).\n\n[3] Gao, Boyan, et al. \"Loss function learning for domain generalization by implicit gradient.\" International Conference on Machine Learning. PMLR, 2022.", "other_strengths_and_weaknesses": "The main deep learning and meta-learning implementation trend is still PyTorch. However, this submission only demonstrates the algorithm on the VJP and JVP-friendly platform, JAX.", "other_comments_or_suggestions": "See the above sections.", "questions_for_authors": "See the above sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eE3xuikdvu", "reviewer": "Reviewer_zFod", "summary": "This paper develops a novel efficient bi-level optimization framework, called MixFlow-MG. Based on the fact that, when computing HVPs, forward-over-reverse mode is usually much more efficient than other modes, the authors first reparameterize the update function, and replaces the usual VHP and VMP computations in the meta-gradient into HVP and MVP computations, respectively, significantly improving the computational efficiency of meta-gradient computation. The analysis show that the proposed method, combined with other techniques such as block rematerialization and inner-gradient saving, significantly reduce the memory overhead, as well as wall-clock speedups on various learning scenarios, making meta-learning feasible even on extremely large models with billions of parameters.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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\n\nUsing 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.\n\n> Just in case, what if you add a simple experiment (in the appendix) showing that the proposed algorithm produces the exact meta-gradient?\n\nThe 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. \n\n**Q1**\n\nIndeed, 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).\n\n`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.\n\n**Q2**\n\nTo 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).\n \nIn 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.\n\nEssentially, 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!\n\n[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\n\n[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\n\n**Q3**\n\nYes, 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. \n\nWe 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.\n\n**Q4**\n\nThe 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.\n\nParameter $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.", "claims_and_evidence": "Overall, I think this paper is very strong. As a meta-learning researcher, it's really interesting and exciting to see that bi-level optimization now become scalable (up to some degree) even with LLMs. All the claims about its computational efficiency are strongly supported by clear evidence, everywhere. All the figures, from Figure 1 to Figure 7, clearly demonstrate that the proposed method dramatically reduce the memory burden and wall time, in a systematic way. I didn't find any issue regarding missing evidence for this paper.", "methods_and_evaluation_criteria": "The propose method is quite simple - just reparameterize the update function, and switch the order of computation to change VHP and VMP into more efficient HVP and MVP. \n\nHonestly, I'm not super expert in low-level implementation, so I was not aware that HVP can be more efficiently computed with the forward-over-reverse mode. While I still don't fully understand the technical details as to why this is the case, but anyway, assuming this is true, I believe the proposed method is very reasonable and is a big contribution to the meta-learning field.\n\nI don't find any specific issue regarding the evaluation criteria.", "theoretical_claims": "I carefully checked the derivation of meta-gradient in Section A.2, and it is correct.", "experimental_designs_or_analyses": "Again, while I'm not an expert in low-level implementation and relevant terminologies, I can clearly see that the analysis are systematic and well controlled. I don't find any relevant issues.", "relation_to_broader_scientific_literature": "This paper makes meta-learning scalable, so it can be a big contribution to the LLM and other relevant fields as well, i.e., all the area where meta-learning was infeasible due to the excessive memory issue.", "essential_references_not_discussed": "I didn't find any missing references.", "other_strengths_and_weaknesses": "I don't see any weaknesses of this paper.\n\nFor the strengths, see above.", "other_comments_or_suggestions": "Just in case, what if you add a simple experiment (in the appendix) showing that the proposed algorithm produces the exact meta-gradient? I don't think it's essential, but it might be helpful to make readers more strongly believe the correctness of the proposed algorithm.", "questions_for_authors": "I have a few simple questions. I'm open to further increase my score if those are resolved.\n\n1. I didn't fully understand how precisely the backprop is done with your algorithm. The one-step unrolling equation in Eq.(6) is for deriving the FMD algorithm, as far as I know (I might be wrong). But it seems that the proposed algorithm seems to be a kind of backprop algorithm. Could you elaborate on this? It's not entirely clear in Algorithm 1 and 2. Specifically, what is \"fwdrev_grad\"? Is it something you implemented, or an existing API? (I'm not familiar with JAX, so was not fully understand A.4)\n\n2. Regarding the proposed reparameterization, I didn't know that the publicly available auto-diff frameworks implements backprop in such a way (because there are extra terms in Eq.(6) compared to when we do not reparameterize). I thought they follow the original parameterization with $\\Phi$ (and in this case, the proposed trick cannot be applied as the blue and red terms in Eq.(6) disappear). Could you provide more details or references about this? \n\n3. It seems that most of the experiments assume rather short inner-learning trajectory, like $T \\in \\{2,4,8\\}$. Could you show the efficiency of your method when $T$ is very large, like $T=10^3$ or $10^4$? \n\n4. I'm a bit confused with the experiment in Section 3.2. Why is it a bi-level optimization problem? Could you elaborate more? What does it mean by controlling $M$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a method (MixFlow-MG) to reduce the memory consumption when computing exact meta-gradients of bi-level optimization problems whose inner loop uses a gradient-based optimizer.
This is achieved by re-parameterizing the inner loop which allows to replace double-backward differentiation with mixed-mode differentiation (forward-over-reverse) that scales more favourably in terms of memory.
Empirical results demonstrate memory reductions by up to 10x as well as small run time improvements on three different setups:

- The proposed method is simple and compatible with the parameterizations assumed by many libraries, allowing it to be broadly adopted.
  The submission contained a JAX implementation, and the authors will add a PyTorch implementation, further broadening applicability.
  Most reviewers appreciated this contribution and emphasized that bi-level optimization research is often limited by memory.
  The proposed method's achieved memory reductions allow existing works to scale by a factor of 10; this is a solid contribution.

- There were concerns regarding the motivation of improving the **exact** computation of meta-gradients, and the lacking comparison with approximate schemes.
  I understand this concern, but think that successfully scaling the exact computation in itself is a sustainable contribution that might lead to the development of new approximate schemes that inherit the proposed improvements. 
  The authors somewhat addressed this in their rebuttal and provided related work that suggests improved results by scaling the exact scheme, e.g. using longer inner loops.
  After reading the paper, I believe that the proposed approach can also benefit some existing approximate schemes, e.g. truncated backpropagation through time.
  Finally, the experiments presented in the paper focus on improving computational performance and do not show how MixFlow-MG improves performance down-stream.
  I think this would be good a addition for the paper, but secondary for its acceptance as this should follow from the demonstrated gains.

Overall, the paper is in good shape and the authors have offered concrete actions to further improve it; therefore I recommend acceptance.
Decision: Accept (poster)