Mnemosyne: Learning to Train Transformers with Transformers
Authors: Deepali Jain
Krzysztof Marcin Choromanski
Kumar Avinava Dubey
Sumeet Singh
Vikas Sindhwani
Tingnan Zhang
Jie Tan
Venue: NeurIPS 2023 poster (2023)
Keywords: learnable optimizers
Transformers
efficient attention
spatio-temporal attention
TL;DR: This paper introduces a new class of learnable optimizers, based on the Transformer-architectures, that can be used for general purpose optimization, in particular to train other Transformer-models.

Abstract: In this work, we propose a new class of learnable optimizers, called Mnemosyne. It is based on the novel spatio-temporal low-rank implicit attention Transformers that can learn to train entire neural network architectures, including other Transformers, without any task-specific optimizer tuning. We show that Mnemosyne: (a) outperforms popular LSTM optimizers (also with new feature engineering to mitigate catastrophic forgetting of LSTMs), (b) can successfully train Transformers while using simple meta-training strategies that require minimal computational resources, (c) matches accuracy-wise SOTA hand-designed optimizers with carefully tuned hyper-parameters (often producing top performing models). Furthermore, Mnemosyne provides space complexity comparable to that of its hand-designed first-order counterparts, which allows it to scale to training larger sets of parameters. We conduct an extensive empirical evaluation of Mnemosyne on: (a) fine-tuning a wide range of Vision Transformers (ViTs) from medium-size architectures to massive ViT-Hs (36 layers, 16 heads), (b) pre-training BERT models and (c) soft prompt-tuning large 11B+ T5XXL models. We complement our results with a comprehensive theoretical analysis of the compact associative memory used by Mnemosyne which we believe was never done before.

Mnemosyne: Learning to Train Transformers with Transformers
Deepali Jain Google DeepMind, jai @indeepali@google.coo com
Krzysztof Choromanski Google DeepMind and Columbia University, kchoro@google, com
Avinava Dubey Google Research, avinavadubey@goglee.com
Sumeet Singh Google DeepMind,
Vikas Sindhwani Google DeepMind, s sindhwani @gwowla...gggogee..oooo com
esumeet@google.co com
Tingman Zhang
Jie Tan Google DeepMind, jietan@google.com com
Google DeepMind, cingnan@google. com
Abstract
In this work, we propose a new class of learnable optimizers, called Mnemosyne. It is based on the novel spatio-tempores low-rank implicit attention Transformers that can learn to train entire neural network architectures, including other Trans- formers, without any task-specific optimizer tuning. We show that Mnemosyne: (a) outperforms popular LSTT optimizers (also with new feature engineering to mitigate catastrophic forgetting of LSTMs), (b) can successfully train Transformers while using simple meta-training strategies that require minimal computational re- sources, (c) matches accuracy- wise SOTA hand-designed optimizers with carefully tuned hyper-parameters (often producing top performing models). Furthermore, Mnemosyne provides space complexity comparable that of its hand- desindd first-order counterparts, which allows to scale to training larger sets of parameters. We conduct extensive empirical evaluation of Mnemosyne on: (a) fine-tuning a wide range Vision Transformers (ViTs) from medium size architectures to massive ViT-Hs (36 layers, 16 heads), (b) pre-training BERT models and (c) soft prompt-tuning large 11B+ T5XXL models. We complement our results with a comprehensive theoretical analysis of the compact associative memory used by Mnemosyne which we believe was never done before.
Introduction
12 used to train machine learning (ML) optimizers can be thought as a natural lifting (to the optimizer-space) of the idea that has revolutionized ML: replacing hand-engineered features with the leamned ones.
Learning-to-lean (I.2L. systems 66 63
28 M 74
The idea train ML optimizers tempting, yet the lift the optimizer-space comes at price: the instances used to train such systems are the optimization problems themselves. Generalization in such a setting means: the ability to transfer knowledge "similar" optimization tasks not seen in training. Rigorous mathematical analysis of the properties of L2L systems, that involves defining distributions over optimization problems, becomes challenging and is subject on its own. Indeed, 39 24 and of critical importance in many disciplines such as Robotics, where transfer knowledge from simulator hardware is
the literature on meta-learning is voluminous 73 53
76
notoriously difficult problem 42 32 57
"qqual contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

A standard approach to learning optimizers is to cast it as a sequential decision problem, where a function f called an optimizee optimized via another function go (an optimizer) with learnable parameters 0. Function I can take as input x the parameters of a neural network (NN) and output its corresponding test loss on given task. The optimizer go updates x as:
Xg optimization initial point,
(I)
Xt+1
ge Xo. Xo...... xt)
if > 0
Function go is either trained by minimizing the meta-loss objective Le. which is usually a sum of f-losses over certain time horizons, with superviseddreinforeeeees learning [40] or in a completely supervised way, where it learns to imitate expert-optimizers [13. It usually does not act directly on the sequence (Xo. .... Xt), but its processed-version. e.g. the sequence of the corresponding gradients (x0), .... f(x, )) if f is differentiable (which does not need to be the case (50). In practice, the pnocessed-versiong often have much richer structure 49]: "..vousious rolling statistics including... data structures, the rolling momentum / rms terms, as well as the gradient clipping state..."
In this paper, we abstract from the low-level design of the sequential inputs to go (which is a subject on its own) and different training strategies of ge- Our interest is in the core design of go since it acts as a memory-
system. Indeed, most models of go leverage recurrent NN cells, such as LSTMs in the form of a compact learnable latent state. In addition, due to the fact that inputs x to I are usually high-dimensional (e.g. neural networks' parameter-vectors). ge is often factorized to process independently different dimensions of x (2].
2 158 577. that keep the history of the optimization-rollont
The Transformer-revolutio in ML set the stage for a very different way to model sequential data and in general: to model memory- the attention mechanism [67
It is natural to ask whether attention architectures can be used to replace LSTM memory cells in L.L systems. Applying Transformers here is compelling modeling long-range relationships over time by avoiding catastrophic forgetting (which is what LSTMs struggle with, but Transformers are particularly good at) is especially relevant for optimizing highly non-convex objectives in deep learning. Yet these benefits come at the cost of quadratic (in the sequence-length) space and time complexity.
19 54 g 22
Contributions. We propose new class of learnable optimizers, called Mnemosyne. is based on the novel spatio neatiottempora low-rank implicit attention Transformers that can learn to train entire neural network architectures, including other Transformers, without any task-specific optimizer tuning. We show that Mnemosyne: (a) outperforms popular LSTM optimizers (also with new feature engincering to mitigate catastrophic forgetting of LSTMs). (b) can successfully train Transformers while using simple meta-training strategies leveraging minimal computational resources, (c) matches accuracy- optimizers with carefully tuned
(often producing top performing models). As we show in Sec. 5 SOTA hand-designed optimizers are very sensitive to the hyperparameter choice: a hyperparameter selection that is optimal for one dataset might perform very poorly on another one. Thus the ability of the optimizer to automatically "implicitly learn" optimal hyperparameters of critical practical importance. Furthermore, Mnemosyne provides space complexity comparable to that of its standard hand-designed first-order counterparts, which allows scale to training larger sets of parameters. We conduct an extensive empirical evaluation of Mnemosyne on: (a) fine-tuning a wide range of Vision Transformers (ViT's) from medium-size architectures to massive ViT-Hs (36 layers, 16 heads), (b) pre-training BERT models and (c) soft prompt-tune tuniit large 11B+ T5XXL models. We also conduct a thorough theoretical analysis of the compact associative memory used by Mnemosyne, to the best of our knowledge never done before.
wise SOTA
Mnemosyne leverages several algorithmic techniques: (a) efficient Transformers, called Performers 18 applying implicit low-rank attention and guaranteeing linear (in the history length for causal and parameter-tensor size for the spatial attention) time and space complexity, (b) bi-directional and uni-directional attention combined in the unified spatio peatio-temporal system, (c) hierarchical spatial attention mechanism to further reduce memory footprint of the spatial attention encoders that can be thought of as novel attention pooling mechanism (see: [.2.TD. Obtained L2L mechanism effectively acts as compact associative memory (CAM) (see: Sec. [4] fed with latent representations summarizing groups of parameters induced by the natural structure of the target model to be trained (the so-called topological encodings, see: Sec 3.22) It thus provides the best of both worlds: efficiency due to the compact fixed-size hidden state (as in LSTMs) and expressiveness since this hidden state approximates regular Transformer's attention via the CAM-mechanism. We also believe that this paper lays the groundwork for the research on general-purpose attention-based learnable optimizers and foundational models for L2L systems.
2

---

2
Related work
The research on L2L systems involves a plethora of techniques: curriculum learning (e.g. incremently increasing optimizer's unroll length [13), randomly scaled optimizees in training with relative scaling of input gradients [47), hierarchical RNN-architectures with lower memory and compute overhead
that are capable of capturing inter-parameter dependencies 77] and more 9 15 68].
Regular Transformers are recently considered in this setting, in particular to tune hyperparameters [1]]. to learn BFGS-type optimization for motion reconstruction [26] as memory-system in class-
incremental learning [31]. Scalable Transformers [65 64] were designed to address computational
limitations of their regular counterparts. Several efficient attention mechanisms were proposed, based
on hashing [35), clustering [59], dimensionality reduction [69] or sparsity [75].
In this paper, we apply in particular methods approximating attention via low-rank decomposition of the attention matrix [18
44
46
17
43 20). due to their intrinsic connection with associative
memory and mmeryy-moed can be thought of as differentiable dictionaries applying powerful associative memory mechanisms 37 56), i.c modern Hopfield networks [29] with exponential memory. Linear low-rank attention mechanisms are their compact variants [6] with intriguing theoretical properties and as such are perfect candidates for scalable memory-systems with respect to the temporal axis (Mnemosyne addresses also the problem of the efficient spatial encodings of the trainable parameters, see Sec. (3.2). That interpretation has profound practical consequences giving guidance on the optimal version of the low-rank attention mechanism. Mnemosyne uses the so-called hyperbolic cosine random features (see: Sec. 3.3.1 I providing particularly low variance of the softmax-kernel estimation, a key ingredient of those modern associative memory models.
[36] (while used in the causal attention setting). Regular Transformers
3 Learning to learn with spatio-temporal attention
In this section, we give the description of Mnemosyne. Since the system consists of several compo- nents, we start by providing a high-level overview. We then discuss individual components in more depth: topological (spatial) encoder in Sec. 3.2 and temporal (causal) encoder in Sec. 3.3
3.1 Preliminaries: tree-like optimization domains
Consider an optimizee f: + R. In the simplest case, we can take: S Rd for d E N+. However it will be convenient to impose a tree-like structure on the elements x E D. We will assume that there exists a tree T such that for every x .... (a)) E DD the set { {Ii,...; ...: Ia} is partitioned into non-empty subsets S1. Si corresponding to different leaves of Ty. where I stands for their number. Not only does Ty define partitioning of {.1, ..., xd} through the subsets residing different leaves, but it also imposes natural hierarchy. The construction might seem artificial at first glance, but we have a good reason to introduce it early on it predominant description of the NN structure (our main interest in this paper is to optimize NNs with learnable optimizers thus we identify elements of DD with different instantiations of the particular NN model). In this context, different leaves correspond to individual tensors of parameters (e.g. weight-matrices or bias-vectors; here subsets S, are just their flattened representations) and tree-induced hicrarchy describes how the NN layers/modules emerge from those individual tensors 2 This more general framework covers also the setting DE E RRd where Ty is a star-tree with different leaves corresponding to different variables Tę.
3.2 Topological encoder
Standard L.LL systems (2) act independently on the individual parameters to be optimized. We refer to this strategy coordinate-wise approach. It has the advantage of enabling lots of data for meta-training (since the optimization trajectory of each scalar parameter is a viable training point) and corresponds to Ty being star-tree (see: Sec. B.Ip. However it comes at price of the memory footprint since space complexity becomes linear in the total number of trainable parameters. Mnemosyne can be successfully applied in the coordinate-wist framework (see: Sec. (5.2). but supports an arbitrary T,, in particular the natural variant, where leaves correspond to the individual tensors of the NN to be optimized and that we refer to as the tensor-wise approach.
Mnemosyne acts independently on each tensor T from each leaf of a given input Ty. Tensor T is first flattened/vectorrized and representation vector is associated with each parameter of T. Since
2 An example of the instantiation of this mechanism is a pytree structure that is default representation of the neural network parameters used in JAX (25) popular machine learning framework used to transform numerical functions and train neural networks.
3

---

in this paper we minimize feature engineering for the L2L systems, our default choice for r is the 2-dimensional vector of the absolute value of the gradient dimension and its sign (that was used on a regular basis in several papers on the subject), but we emphasize that Mnemosyne is agnostic to the particular representation choice. The resulting sequence of representations ST is then transformed by the bi-directional attention of the Performer model [18]. A latent encoding of a fixed token from that sequence is then output as a topological encoding of T.
3.2.1
Compactifying topoloiiaal encoder with the hierarchical pooling
Using bi-directional linear attention from Performers is justified by the fact that sequence ST can be in practice very long. In fact it easily surpass 1M tokens (for instance if T is square weight-tensor corresponding to two consecutive layers of size > 1K each). In that setting, even linear attention might not suffice. To address it, introduce additional hierarchical pooling mechanism, leading to the hierarchical pooling encoder (HPE). Sequence ST is first split into chunks of a fixed length L E N+: St. Sk. ... (the last chunk might be shorter). Topological encoding is applied in parallel to each chunk. 'he procedure is repeated for the resulting sequence of the topological encodings, which already shorter by the multiplicative factor of L (potentially with different L even though here we assume hat L is the same). The total number of repetitions is in practice a small constant Aooo E N
and leads to the final sequence of length as a small constant. If (Sr) is small enough, hierarchical pooling is not applied. The resulting 1-length sequence of latent d-dimensional encodings is output as topological encoding and fed to the temporal encoder defined below. We refer to different tokens of that sequence as meta-tokens, M.
len(ST)
7."peal where len(Sr) is the original length and I
3.3 Temporal encoder: Compact Associative Memory (CAM) model
Preliminaries. The temporal module consists of one more temporal encoders stacked together that process the meta-tokens, .M. It treats the length 1 of the sequence .M as a batch size b. For given meta-token, denote by {E" }==1 C Rd its corresponding latent encodings obtained over time. We will refer to them here as memory-vectors (patterns). We obtain their latent embeddings: queries, keys and values via learnable linear transformations Wq. WK K
R Nxd
Wy E Rddd as follows:
qdt Wot", k" ="=Wk WKS", v Wvg"
(2)
Take a kernel K R x R + R. and its linearization: K(x,y) T [(y]] for some (randomized) Φ RN + Rr, We refer to Φ(x), d(y) as random feature (RF) vectors and define hidden state encapsulating memory of the system of first t patterns as:
N, = Σ. - (VV)T E ERrxd V Σ Ep=1 {e(H))((")) ER
(3)
A discount-function X : IR + IR is applied to deprioritize older patterns.
3.3.1 Updating hidden states and the outputs of the temporal module
Note that temporal encoder's hidden state hMme (4) (N,, Vz) ) is of size independent from the number of its implicitly stored patterns l. When patterns added, h! needs be efficiently updated on-the-fly. It is easy to see that this can be done for the exponential discount strategy: de(H) = exp(-T(t- 4)) with 2 (7 0 turns off discounting). We have the following:
N = exp( (T) N + ((k*+)(((++)))))
(4)
L 1 exp( (T) V. + o(k'+1)
With the definition of the temporal encoder's s hidden state, we can now explain how it acts on the input vectors & E Rdd. New vector g' E + Ag is obtained as follows, where q = WoC:
N ((q) A = ( V
Σ
T o(k))
E MO)0(( v
(5)
Vector ç' can be computed in time Owr(1) and given as a convex combination of value vectors va, with coefficients proportional to approximated kernel values, but modulated by the discount-function

---

For We WK = L and with exact kernel values Kíq. k') in Eq. 5 (rather
than their approximated versions Φ(q) o(k)), dynamical systems defined by Eq. 5 become ef- fectively Hopfield networks and, as energy-based models with energies given as E(: 1 m
- 1 (t) retrieve memory-vectors upon energy sanaaammmimiizzti-drieeeeee convergence [56] The
retrieval quality depends on the kernel, with the softmax-kemel K(x. y) def exp(x y) providing
particularly strong theoretical results. For arbitrary Wq. W.. Wv, but still with X = and exact kernel values, Eq. turns into regular Transformers' attention. Upon replacement of the exact kernel values with the approximate ones, Performer model is recovered.
We think about Eq. 5 as a generalized (since it uses Xς) compact (since provides efficient hidden state and input update-rules via linearized kemels from Performers) associative memory model (CAM) (as opposed to the regular associate memory model from Hopfield networks).
In Mnemosyne, we choose roftmax-kermn as K, since is a default choice for Transformers. We observed that the FAVOR+ mechanism defining & from 44 and denoted by us or-+, provides the most robust performance, but since is harder to implement in the temporal encoder setting (see: discussion in Sec. (.II) in practice apply the so-called hyperbolic cosine random features from [8], performing very similarly, as O More details including in particular exact definitions and ablation studies over different random feature mechanisms can be found in Sec. A. B.3
3.4 Putting it all together
Details of the hierachical pooling encoder (HPE) and the compact associative memory (CAM) modules are depicted in Fig. [ With all the components of Mnemosyne described, we present the complete design of Mnemosyne for two application modes: coordinate-wise and tensor-wise. shown in Fig. 2 In the coordinate-wise application, the enriched gradient input r, (see: Sec. 3.2 of every parameter 2; processed separately in parallel by a CAM module followed by an MLP layer to produce the update Ar, This makes the optimizer design simple and allows it to learn optimization from small-scale problems since each parameter input serves as a separate training example. In this setting, CAM stores fix-sized memory state for each parameter thereby making the optimizer memory state scale linearly with the number of parameters. In the tensor-wise application, every tensor ST in the parameter-tre tree is processed as whole by an HPE to produce meta-tokens M for CAM. Now the CAM memory state becomes very compact because it stores a state for each of the (small number of) meta-tokens rather than every token in the original input sequence. CAM output L has the same shape as M. It is transformed into a fix-sized encoding e by a spatial attention encoder (SPE). This encoding is broadcast to all the input tokens of the tensor via concatenations with vectors r. Each resulting vector r C e is processed by a single MLP-layer to get the update tensor AT.
Mediectooll
Bid scciona Perferar
N
143
W K..
12
(1)
N
pollieg
Necorbiica pool sncldde ((PG)
COMPIN MoSCoaite memoyyy (210)
Figure 1: Pictorial description of the hierarchical pooling encoding (HPE) and compact associative memory (CAM) in Mnemosyne on the example of modifying a single weight-tensor of a given NN. Consider training a single 3 x 4 weight tensor of a toy feedforward fully connected NN ((a)). Three snapshots of this tensor ((b)) represent its consecutive instantiations in the optimization process. HPE ((c)) acts follows. Each tensor is vectorized and chunked into sub- sebb-sequnnes that are spatially encoded by the bi-directional Performers. Presented pooling mechanism consists of two layers. This results the input tensors 5", 5² to CAM. Here the first dimension (the final number of meta-tokens) serves as a batch one (b- = 2) and N = 4 (see: notation from Sec. 3.3). Those tensors are first linearly mapped via W matrices to keys (shown above) and queries and then non-linearly transformed via the d-mapping. The transformed variants are leveraged by the associative memory.
5

---

CAM
MLP
Aut
(a)
HPE
M
CAM
e
SPE
ST
MLP
AT
(b)
Figure 2: Two modes of Mnemosyne application: (a) coordinate- and (b) tensor-wise.
4 The theory of Mnemosyne's compact associative memory
In this section, we analyze Mnemosyne's compact associative memory (CAM) from the theoretical point of view. We show that, its regular non-compact counterpart, is capable of storing patterns, but (as opposed to the former) in the implicit manner. We start by providing a concise introduction to associative memory models "analytical" (rather than combinatorial) energy-based nearest-neighbor search systems. We then present our main result (Theorem 4.3) stating that "on average" CAMs restore exponentially many (in space dimensionality) number of patterns, provided that they spread "well enough". To the best of our knowledge, this the first result of this type.
4.1 Regular exponential associative memory
As in several other papers providing a theoretical analysis of the associative memory, we consider feature vectors taken from the set [-1,+1] +1,%, We denote by {C*) the set of all the memory- vectors to be stored. For the given input ç, the regular exponential associative memory model, the energy of the system is defined as:
M
Eveg(E: (E. = Σ
exp(5 T t*).
(6)
The dynamical system defining the interactions with the associative memory and whose goal is to retrieve the relevant memory for the given input vector E {-1,+11A (its nearest neighbor in (E),M)) has the following form: o - Ti - (Ti, (a)) + for the initial point a, where ii, i2, are chosen independently at random and Tj (-1, +1J + {-1,+1] only updates the jth entry its input as follows (for Elj;; x| denoting vector ç, but with its jth dimension equal to r):
(S(L) sgn (Freg (Eli; - -1::5' 1): Freg(LLi (Elj; 1):5",,..
(7)
Thus at every step, random dimension the input vector is chosen and its value is flipped if that operation decreases the energy of the system.
Definition 4.1 (the capacity of the associative models). We say that the model described above stores there exists E (o, }) such that (f") = Et for any andd? g" taken from the Hamming ball B(g", pN) centered in g and of Hamming radius p.N.
memories {5" (M Jum1
It was proven in (21 that if the memories are chosen uniformly and independently random from [-1,+1}%, then with probability approaching as N - oo the model stores all the memories as long as the memory-set is not too large. Most importantly, the upper bound the memorics size exponential in N.
Remark 4.2. Despite the exponential capacity of the model, all its memories need to be explicitly stored (to compute the energy mnergy-function) and compute time is proportional to their number. Thus for large number of memories, the space and time complexity makes the retrieval infeasible in practice.
4.2 Mnemosyne's Compact Associative Memory (CAM)
We denote by W1, w2, ...up samples chosen independently from the multivariate Gaussian distribution (0, IN) and define the energy the system as follows (see: Sec. (.3. A.I
Exand 5\ (.. = = MM('.
where M is given as: M(S ..., g....MM MM oF++ (E+). We refer to the number of random projection vectors
the number of random features (RFs).
Equipped with this new energy function, we define the corresponding dynamical system in the same way as for the regular associative memory model. Calculating energy Erand can be now

---

done efficiently in time (rr), once vector M is computed (thus independently from the number of implicitly stored memories). In the online/streaming setting, when the memories come one by one, updating vector M can be done in time O(rr) per query.
4.3 The capacity of Mnemosyne's memory
We are ready to present our main theoretical result. We assume the setting from Sec. 4.2 The extended version of this result, providing in addition concentration results, is given in Sec. A.2
+1], the memory-vectors. Assume that the Hamming distance between any two memory-vectors is at least TN for some > O. Take some 0 <. p 2. Then the following is true for any memory-vector E for 2e - the expected change of the energy of the compact associative memory system A (Erand) associated with flipping the value of the dimension of 2 is positive if that operation increases the distance from its close neighbor g and is negative otherwise.
Theorem 4.3 (storage of compact associative memories). Denote by
l = 1,... P and any input 2 E Big', ) as long as M S exp(2N(T - 2p))
Remark 4.4. Theorem 4.3 says that the expected value of the change of the energy of the system has the correct sign even the number of stored patterns exponential in their dimensionality, provided that the patterns are well separated. By the analogous argument the from the proof of Theorem 4.3 in Sec. A.2 a similar statement for the method applying other RF-mechanism for softmax -ernel estimation can be derived. However + provides the smallest variance among all competitors as the most accurate currently known mechanism for the unbiased softmax-kerun approximation.
5 Experiments
This section is organized as follows. Sec. 5.1 is a warm-up, where we provide initial comparison of Mnemosyne with several hand-designed optimizers and an LSTM-based learned optimizer baseline, on the tasks of training smaller NN architectures. Our results show that Mnemosyne consistently outperforms other variants and that popular Adam optimizer [4 is the second best. Note also that Adam is an optimizer of choice for training large Transformer models. Therefore in the following sections, we focus on the detailed comparison of Mnemosyne with Adam for larger architectures.
In Sec. 5.2, we test the coordinate-wise Mnemosyne for ViT fine-tuning and soft prompt-tuning 11B+ T5XXL models. Optimizer's 's memory scales linearly with the NN size for the coordinate- wise variant and Adam. However, depending on the size and number of temporal attention layers in the coordinate-wist Mnemosyne, the linear multiplicative constant can be prohibitively large. This restriction is alleviated by the tensor-wise variant. In Sec. 5.3 we present the results for the tensor- wise Mnemosyne on BERT MI. pre-training and VIT one-tuning. The tensor-wise Mnemosyne's memory state scales with the number of tensors instead of the number of model parameters. However, this variant requires meta-training with larger tensors for best generalization results since now each tensor serves as a single training example instead of each parameter which was the case for the coordinate-wise variant. To take advantage of the efficiency of the tensor-wise and the efficacy of the coordinate-wise variant, we present Super-Mnemosyne in Sec. 5.4, that merges them.
All Mnemosyne variants are meta-trained on small scale MLP and VIT training tasks for short horizon of 100 steps. More detailed description of the meta- meta-training set-up is given in the Appendix (Sec: B) along with the details of the experiments in this section. Additional results including: (a) comparison of the CAM mechanism with regular attention (Sec. B.2), (b) studies over different RF-mechanisms (Sec. B.3) ablations over: (c) different discount factors and different number of random features for the CAM mechanism (Sec. B.5) and (d) different depths of the temporal modules
of Mnemosyne (Sec. B.6) are also given in the Appendix.
We have finalized the design choices for Mnemosyne optimizer used throughout this section based on the extensive theoretical analysis in previous sections and various ablation studies detailed in the Appendix. In particular:
1. The choice of the random feature (RF) mechanism is motivated by strong theoretical guarantees for the capacity of the corresponding associative memory. Other RF mechanisms that do not admit our theoretical analysis clearly underperform in the ablation results.
2. Our default hyperbolic cosine RF mechanism motivated by the provably reduced estima- tion variance as explained previously.
7

---

DGC
EG0
Rlsdcing WULD Original (1.TMM)
Figure 3: First three plots: Test loss curves for training VITs with Mnemosyne on: MNIST. CIFARIO and CIFAR100. Minimal feature engineering is applied. Last plot: Test loss curves for training an MLP on MNIST. Both Mnemosyne and LSTM are used within VeLO architecture using sophisticated feature enginecring.
3. We choose hyperbolic cosine R as they are shown more robust than regular positive RFs in ablation.
4. Our use of exponential discount strategy (EDS) to smoothly discount the past is theoretically motivated previously. We show empirical verification of EDS in the appendix and choose the discount factor based on ablation results.
5.1 Warm-up: Mnemosyne vs other optimizers for smaller NN architectures
We start by comparing coordinate-wise Mnemosyne with standard optimizers: Adam 34 RMSProp 78. SGD [3] as well as popular learnable optimizers using LSTMs 22. All optimizers were tested on the task of training Vision Transformer (VIT) 23 architectures not seen during meta-training. Considered VIT has 3 attention and MLP layers of size 16 and 2 heads. Loss curves for image classification with ViTs on different datasets (MNIST, CIFARIO, CIFAR 100) are shown as first three plots in Fig. 3 Minimal feature engineering techniques were applied. In the last plot of Fig. 3 we inserted Mnemosyne into VeLO optimizer [48) which originally used LSTMs. This learnable optimizer applies sophisticated feature engineering to mitigate the problem of catastrophic forgetting of the LSTM-cells. Note that we did not replicate their large-scale meta- tett-trainnng set-up, only the optimizer architecture and features. Mnemosyne outperforms its counterparts in all these scenarios. In particular, it successfully trains creention-baede architectures unseen tasks.
Now we proceed with the comparison of Mnemosyne and Adam with different learning rates on larger Transformers. We emphasize that we do not conduct any hyperparameter tuning for Mnemosyne.
5.2 Coordinate-wise Mnemosyne results
VIT last-layer fine-tuning: We tested coordinate-wise Mnemosyne with 2 temporal encoders on ViT fine-tuning
23 Due to memory constraints, we only trained the last layer of ViT by freezing
other parameters of the model. The results three different datasets: imagenet2012, places365 and
caltech-birds- 2011 for ViT-H(32) (36 layers, 16 heads, 32 x 32 patch shape) are presented in Fig. 4 Additional results, including ablations with varying architecture sizes and other datasets, shown Appendix Sec: B.8 We conclude that Mnemosyne (without any hyperparameter tuning) matches the top-performing Adam variant. Note that Adam is very sensitive to the choice of the learning rate Ir.
In Table D we compare coordinate-wise Mnemosyne with two other classes of learnable optimizers: a) Those using $4 memory units [27] based on state space machines as alternative to LSTMs and Transformers b) The Lion family learnable optimizers [4] using a different approach of symbolic programming. In the case of VLT-L architecture, we skip S4 due poor performance the previous tasks. Mnemosyne is clearly the best among all three variants. In contrast to Lion, Mnemosyne does not require any hyperparameter tuning.
14
Imagenet2012 VIT-I(122
Acam --200
AGa +300
11
VIH+
placea 365
bind 2011 V114422)
120
Figure 4: Fine-tuning the last layer of ViT-H(32) on different datasets with coordinate-wise Mnemosyne.
8

---

Method
VIT-H (ps=14)
VIT-H
ViT-L
(ps=8)
(ps=32)
Mnemosyne
72.3
81.4
74.0
S4
7.5
35.0
Lion Ir=le-1
67.2
76.4
66.8
Lion Ir=le-2
68.3
80.0
70.0
Lion Ir=2e-3
70.0
80.0
71.4
Lion Ir=le-3
72.3
81.2
71.4
Lion 1r=1e-4
70.0
81.0
71.4
Table 1:
Comparison of coordinate-wise Mnemosyne with Lion and S4-based learned optimizer. Accuracy after fine sine-tunig the last layer of different VIT architectures and different patch sizes (ps) Imagenet2012 dataset for 500 iterations.
VIT multi-layer fine-tuning: The light Mnemosyne variant with single temporal encoder can scale up to fine-tune 2 Transformer layers along with the last layer. We show the result for fine-tuning ViT-B(16) on Cifar100 in Fig, [ The non-M -on-Mnnmosyne layers are trained using an untuned Adam and the Adam variants for comparison fine-tune the full model. We observe that full-model finetuning is a hard task where all but one Adam variant fail to reach the optimal accuracy. Here, even by training part of the network with Mnemosyne, achieve the best accuracy. Also note that the best Adam learning rate in this experiment (le-)) was performing poorly for the previous one. We see this in subsequent results as well.
Figure 5: Fine-tuning ViT-B(16).
Soft prompt-tuning T5XXL: We tested coordinate-wise Mnemosyne for soft prompt-tuning IIB+ TSXXL Transform- ers (55) on the SuperGLUE task (38). This method was intro- duced as scalable alternative to fine-tuning pre-trained models for several downstream tasks, with learnable prompts injected into Transformer layers and modulating its behaviour. The trainable soft-prompt contains 12288 parameters. Mnemosyne outperforms all Adam variants (see: Fig. 6)
800
700
600-
500
150
20.
2.9
120
Panntioon
Figure 6: Soft prompt-tuning T5XXL.
5.3 wissor-wise Mnemosyne results
In this section, we evaluate the memory-efficient tensor-wise Mnemosyne.
BERT NLP Transformer pre-training: We showcase the ability of tensor-wise Mnemosyne to pre-train BERT-base text Transformer 22 on the standard masked language modeling (MLM) task.
Coordinate wise Mnemosyne was not applicable here due to the memory footprint and thus it became a practical test for the memory efficient tensor wise variant. With that variant, we were able to train 86M parameters of Bert-base model thereby showcasing the scalability of tensor-wise Mnemosyne. We
Auen Acom lod Niamoooye
compare Mnemosyne with different variants of Adam in Fig 7 We see that Mnemosyne matches the performance of the best Adam variant though it was never exposed to the MLM task during meta-training. Several Adam variants get stuck in a local optima, worse than that found using Mnemosyne.
Figure 7: MLM task with BERT.
VIT multi-layer fine-tuning: Now we show the performance of tensor-wise Mnemosyne on fine-tuning ViTs. Although tensor-wise Mnemosyne can scale up to the full ViT model, here we freeze large tensors in the Transformer layers and fine-tune the rest. Training large tensors with tensor-wise Mnemosyne requires data and compute intensive large scale meta-training (to ensure that spatial encoders learn how to encode long sequences well). This will be the focus of the future work. For this experiment, consider the ViT-H(32) model and three datasets:

---

Cifar100, places365, imagenet in Fig. 8 Mnemosyne shows equivalent or better performance as compared the optimal Adam.
cffar00 VITH
place366 VIHHOC
Adu 3-0-2 110.0 scoat
imageeet201
1400
VIH+102)
Figure Fine-tuning multi-layer ViT-H(32) on different datasets with tensor-wise Mnemosyne.
Accuece
12
clfa1100 VIPH(32)
places385 WIFH632)
Imagenec2012
VIT+132)
20
100
1254
2000
1350
3000
1718
Becrtton
Figure 9: Fine-tuning full ViT-B(16) on different datasets with Super-Mnemosyne
5.4 Super- Super-Mnemosynee combining coordinate- and tensor-wise strategies
In this last set of VIT-experiments, combine tensor- wise and coordinate- sorddinate-wisel light Mnemosyne. Since applying the former on large tensors requires more intense meta-training which is out of the scope of this paper, decided optimize the largest tensors with the coordinate-we with the tensor-wise Mnemosyne (to keep meta-training simple). This combination turned out to minimize total memory usage. As
and others
in Fig. Mnemosyne outperforms optimal Adam variants.
Broader impacts & limitations
We believe that Mnemosyne opens research on attention-based optimizers for general- purposs optimization. Our system should be used responsibly due to the potential for misuse, significant societal impact, and carbon footprint of Transformers 10 70 the future we want analyze in more depth the impact on more complex meta-training strategies on the quality of learned optimizers.
7 Conclusion
We proposed new class learnable optimizers applying efficient spatio-tempont attention, called Mnemosyne. We show that they outperform their LSTTM-based counterparts and can be successfully used to fine/soft prompt-tune and pre-train large Transformer models, matching optimal hard-coded variants without any hyper-parameter tuning, often producing top performing models.
10

---

References
[1] Ferran Alet, Tomás Lozano-Pérez, and Leslie Pack Kaelbling. Modular meta-learning. 2nd Annual Conference Robot Learning, CoRL. 2018, Zürich, Switzerland, 29-31 October 20/8, Proceedings, volume 87 of Proceedings of Machine Learning Research, pages 556-868. PMLR, 2018.
[2] Marcin Andrychowicz, Misha Denil, Sergio Gomez Colmenarejo, Matthew W. Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning learn by gradient descent by gradient descent. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors. Advances in Neural Information Processing Systems 29; Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 3981-3989, 2016.
[3] Emily M. Bender, Timnit Gebru, Angelina McMillan-Majon and Shmargaret Shmitchell. On the dangers of stochastic parrots: Can language models be too big?.. In Proceedings of the 2021 ACM Conference Fairness, Accountability, and Transparency, FAccT '21, page 610-623, New York, NY, USA, 2021. Association for Computing Machinery.
[4] Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Geesei. On the optimization of synaptic learning rule. In Conference on Optimality in Biological and Artificial Networks, 1992.
[5] Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Geesei. On the search for new learning rules for anns. In Neural Processing Letters, pages 26-30, 1995.
[6] Glen Berseth, Zhiwei Zhang, Grace Zhang, Chelsea Finn, and Sergey Levine. Comps: Continual meta policy search. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[7] Authony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06811 2022.
[8] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan. Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[9] Yue Cao, Tianlong Chen, Zhangyang Wang, and Yang Shen. Learning optimize in swarms. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 15018-15028, 2019.
[10] Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Vore Kather- ine Lee, Adam Roberts, Tom B. Brown, Dawn Song, Úlfar Erlingsson, Alina Oprea, and Colin Raffel. Extracting training data from large language models. CORR, abs/2012.07805,. 2020.
[11] Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George F. Foster, Llion Jones, Mike Schuster, Noam Shazeer, Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Zhifeng Chen, Yonghui Wu, and Macduff Hughes. The best of both worlds: Combining recent advances in neural machine translation. In Iryna Gurevych and Yusuke Miyao, editors, Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume 1: Long Papers, pages 76-86. Association for Computational Linguistics, 2018
[12] Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang. and Wotao Yin. Learning to optimize: A primer and benchmark. CoRR, abs/2103.1 12828. 2021.
11

---

[13] Tianlong Chen, Weiyi Zhang. Jingyang Zhou, Shiyu Chang, Sijia Liu, Lisa Amini, and Zhangyang Wang. Training stronger baselines for learning optimize. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Ties Lin, editors, Ad- vances in Neural Information Processing Systems 33 Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[14] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong. Cho-Jui Hsieh, et al. Symbolic discovery of optimization algo- rithms. arXiv preprint ar arXiv:2332.006675 2023.
[15] Yutian Chen, Matthew W. Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Timothy P. Lillierap, and Nando de Freitas. Learning learn for global optimization of black box functions. CoRR, abs/1 61 1.08224, 2016.
[16] Yutian Chen, Xingyou Song, Chansoo Lee, Zi Wang. Qiuyi Zhang, David Dohan, Kazuya Kawakami, Greg Kochanski, Arnaud Doucet, Mare' Aurelio Ranzato, Sagi Perel, and Nando de Freitas. Towards learning universal hyperparameter optimizers with transformers. CoRR, Jbs/2205.13320, 2022.
[17] Krzysztof Choromanski, Han Lin, Haoxian Chen, Tianyi Zhang, Arijit Schanobish, Valerii Likhosherstov, Jack Parker-Holder, Tamás Sarlós, Adrian Weller, and Thomas Weingarten. From block-Toepline matrices to differential equations on graphs: towards general theory for scalable masked transformers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 1962-3983. PMLR, 2022.
[18] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[19] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellster Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. CoRR, abs/2204.02311 2022.
[20] Sankalan Pal Chowdhury, Adamos Solomou, Avinava Dubey, and Mrinmaya Sachan. On learning the transformer kernel. Transactions of Machine Learning Research, 2022.
[21] Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On a model of associative memory with huge storage capacity. Journal of Statistical Physics, 168(2):288- 299, may 2017.
[22] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pages
4171 4416. Association for Computational Linguistics, 2019.
12

---

[23] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[24] Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. Generalization of model-agnostic meta-learning algorithms: Recurring and unseen tasks. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, viriual, pages 5469-5480, 2021.
[25] Roy Frostig, Matthew Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. 2018.
[26] Erik Gärtner, Luke Metz, Mykhaylo Andriluka, C Daniel Freeman, and Cristian Sminchisescu. Transformer-based learned optimization. arXiv preprini arXiv:2212.01005, 2022.
[27] Albert Gu. Karan Goel, and Christopher Ré, Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.
[28] Sepp Hochreiter, A. Steven Younger, and Peter R. Conwell. Learning to learn using gradient descent. In Georg Dorffner, Horst Bischof, and Kurt Homik, editors, Artificial Neural Networks ICANN 2001, International Conference Vienna, Austria, August 21-25, 2001 Proceedings, volume 2130 of Lecture Notes in Computer Science, pages 87-94. Springer, 2001.
[29] John J. Hopfield. Hopfield network. Scholarpedia, 2(5):1977, 2007.
[30] Jeffrey Ichnowski, Yahav Avigal, Vishal Satish, and Ken Goldberg. Deep learning can accelerate grasp-optimized motion planning. Science Robotics, 5(48))eabd7710, 2020.
[31] Ahmet Iscen, Thomas Bird, Mathilde Caron, Alireza Fathi, and Cordelia Schmid. A memory transformer network for incremental learning. CoRR, Abs/2210.044. 0.04485, 2022.
[32] Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data-efficien robotic grasping via randomized- to to-canonical adaptation networks. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 12627-12637. Computer Vision Foundation / IEEE, 2019.
[33] Chi Jin, Praneeth Netrapalli, Rong Ge, Sham M. Kakade, and Michael I. Jordan. Stochastic gradient descent escapes saddle points efficiently. CoRR, 1b2/1992.04811, 2019
[34] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[35] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[36] Dmitry Krotov and John J. Hopfield. Dense associative memory for pattern recognition. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29; Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 1172- 1180, 2016.
[37] Dmitry Krotov and John J. Hopfield. Large associative memory problem in neurobiology and machine learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 1021.
13

---

[38] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for -fficieteere prompt tuning. In Farie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Viruall Event Punta Cana, Dominican Republic, 7-11 November, 2021, pages 3045-3059. Association for Computational Linguistics, 2021.
[39] Jeffrey Li, Mikhail Khodak, Sebastian Caldas, and Ameet Talwalkar, Differentially private meta-lcarning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[40] Ke Li and Jitendra Malik. Learning to optimize. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[41] Weiwei Li and Emanuel Todorov. Iterative linear quadratic regulator design for nonlinear biological movement systems. In ICINCO (1), pages 222-229. Citeseer, 2004.
[42] Jacky Liang, Saumya Saxena, and Oliver Kroemer. Learning active task-oriented exploration policies for bridging the sim-to-real gap. In Mare Toussaint, Antonio Bicchi, and Tucker Hermans, editors, Robotics: Science and Systems XVI, Virtual Event / Corvalis, Oregon, USA, July 12-16, 2020, 2020.
[43] Valerii Likhosherstov, Krzysztof Choromanski, Jared Davis, Xingyou Song, and Adrian Weller. Sub-linear memory: How to make performers slim. CoRR, abs/2012. 11344,,,..11 2020.
[44] Valerii Likhosherstov. Krzysztof Choromanski, Avinava Dubey, Frederick Liu, Tamás Sarlós, and Adrian Weller. Chefs' random tables: Non-trigonometric random features. to appear at AAAI 2023, abs/2205. 15317, 2 2022.
[45] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 5747-5763. Association for Computational Linguistics, 2020.
[46] Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie- Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. CoRR, abs/2106. 12566, 2021.
[47] Kaifeng Lv, Shunhua Jiang, and Jian Li. Learning gradient descent: Better generalization and longer horizons. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 2247-2255. PMLR, 2017.
[48] Luke Metz, James Harrison, C. Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, and Jascha Sohl-Dickstein. Velo: Training versatile learned optimizers by scaling up. CoRR, abs/2211 1.09760, 2022.
[49] Luke Metz, Niru Maheswaranathan, C. Daniel Freeman, Ben Poole, and Jascha Sahl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. CoRR, abs/2009. 1243, 2020.
[50] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, C. Daniel Freeman, and Jascha Sohl- Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning. ICML 2019, 9-15 June 2019, Long Beach. California, USA, volume 97 of Proceedings of Machine Learning Research, pages 4556-4565. PMLR, 2019.
[51] Luke Metz, Niru Maheswaranathan, Jonathon Shlens, Jascha Sohl-Dickstein, and Ekin D. Cubuk. Using learned optimizers make models robust to input noise. CoRR, wttp//996.0333. 2019.
14

---

[52] Dattika K. Naik and Richard J. Mammone. Meta-neural networks that learn by learning. [Proceedings 1992] IJCNN International Joint Conference on Neural Networks, 1:437-442 vol.1, 1992.
[53] Vitchyr H. Pong, Ashvin V. Nair, Laura M. Smith, Catherine Huang, and Sergey Levine. Offline meta-reinffrreemet learning with online self- superviiio.... In Kamalika Chaudhuri. Stefanic Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 17811-17829. PMLR, 2022.
[54] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
[55] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 5551, 2020.
[56] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Milena Pavlovic, Geir Kjetil Sandve, Victor Greiff, David P. Kreil, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. CoRR, has/2008.02211,, 2020.
[57] Sachin Ravi and Hugo Larochelle. Optimization as model for few-shot learning. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[58] Anirban Roy and Sinisa Todorovic. Learning to learn second-order back- soccppropagation for enns using Istms. In 24th International Conference on Pattern Recognition, ICPR 2018, Beijing, China, August 20-24, 2018, pages 97-102. IEEE Computer Society, 2018.
[59] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68, 2021.
[60] Adam Santoro, Sergey Bartunov, Matthew M. Botvinick, Daan Wierstra, and Timothy P. neural networks. In Maria-Floring Balcan and Kilian Q Weinberger, editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 1842-1850. JMLR. .org, 2016
Lillicrap. Meta-learning with
[61] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, PMLR, 2021.
volume 139 of Proceedings of Machine Learning Research, pages
[62] Sumeet Singh, Jean- Slotine, and Vikas Sindhwani. Optimizing trajectories with closed- loop dynamic sqp. arXiv preprint wrXiv:2109.0708/ 2021.
[63] Richard S. Sutton. Adapting bias by gradient descent: An incremental version of delta-bar-delta. In William R. Swartout, editor, Proceedings of the 10th National Conference on Artificial Intelligence, San Jose, CA, USA, July 12-16, 1992, pages 171-176. AAAI Press / The MIT Press, 1992.
[64] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[65] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler, Efficient transformers: A survey, CoRR. htts/2009.0673.... 2020.
[66] Sebastian Thrun and Lorien Y. Pratt, editors. Learning to Learn. Springer, 1998.
15

---

[67] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information P Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.
[68] Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In Marina Mcila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 00563-10563. PMLR, 2021.
[69] Sinong Wang, Belinda Z. Li, Madian Khabsa. Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, (bs/2006.04766. 2020.
[70] Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, et al. Ethical and social risks of harm from language models. arXiv preprint arXiv:2 2ttp::://2204433 1..04399 2021.
[71] Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando de Freitas, and Jascha Sohl-Dicksteen. Learned optimizers that scale and generalize. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6- 11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 3751-3760. PMLR, 2017.
[72] Xuesu Xiao, Tingnan Zhang, Krzysztof Marcin Choromanski, Tsang- Wei Edward Lee, An- thony G. Francis, Jake Varley, Stephen Tu, Sumeet Singh, Peng Xu, Fei Xia, Sven Mikael Persson, Dmitry Kalashnikov, Leila Takayama, Roy Frostig, Jie Tan, Carolina Parada, and Vikas Sindhwani. Learning model predictive controllers with real-time attention for woal-world naviga- tion. In Karen Liu, Dana Kulic, and Jeffrey Ichnowski, editors, Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pages 1721. PMLR, 2022.
[73] Huaxiu Yao, Linjun Zhang, and Chelsea Finn. Meta-learning with fewer tasks through task interpolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[74] Arthur Steven Younger, Sepp Hochreiter, and Peter R. Conwell. Meta-learning with backprop- agation. IJCNN' 01. International Joint Conference on Neural Networks. Proceedings (Cat. No.0ICH37222), 3:2001-2006 v 2001.
[75] Manzil Zaheer. Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33:1 13:17283-17297, 2020.
[76] Zihao Zhao, Anusha Nagabandi, Kate Rakelly, Chelsea Finn, and Sergey Levine. MELD: meta-reinforcement learning from images via latent state models. In Jens Kober, Fabio Ramos, and Claire J. Tomlin, editors, 4th Conference on Robot Learning, CoRL 2020, 16-18 November 2020, Virtual Event / Cambridge, MA, USA, volume 155 of Proceedings of Machine Learning Research, pages 1266-1261. PMLR, 2020.
[77] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In IEEE international conference on computer vision,
pages 19-27, 2015.
[78] Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 11127-11 135.. Computer Vision Foundation / IEFE, 2019.
16

---

A Proofs
A.1 Kernels & their linearizations for temporal encoders in Mnemosyne
We tested different transformations é and discovered that those leading to most accurate approxima- tion of the seltmax-kerned lead to most effective memory mechanisms for Mnemosyne's 's temporal encoders (see: Sec. 3.3 Our starting variant is the so-called FAVOR+ mechansim from [18), given
def
[x]l2 and Wir... Wep _N(0,
as follows for [(z, r)
exp(
Φr+ (z) [(z, (exp(wi 2), exp 2))
Random vectors W1, Wp form a block-orthogonal ensemble (see: [8]
We applied also its im-
provement relying on the so-called hyperbolic cosine random features, where I[ is the concatenation operator:
Φ1F+(Z) = r(z, r) (exp(() 2), exp(-w) 2))
(10)
i=
Both randomized transformations provide unbiased estimation of the softmax-kernel yet the latter one (that can be cast as modified 6Ff via the antithetic Monte Carlo trick) has provably lower approximation variance.
A.1.1 The curious case of linearization with bounded features
The last variant for the efficient estimation of the softmax-kernel we applied, is very recent mechanism FAVOR++ from 44 given as:
D VF II
OF (z)
expl A Bu: CT111))
(1-4A)) A and E (0, 1) is free parameter. As opposed to the previous variants, mechanism o (z) provides estimation via bounded random variables (since Ā > 0), leading to stronger concentration results (beyond second moment) and still unbiased approximation.
where we have: - V -
The optimal choice of p depends the kernel inputs. The formula for p optimizing the variance of the kernel matrix estimation K= [K(q', k)),s_=i....,M induced by the softmax-kernel K (in the bi-directional case) is not tractable. However choosing by optimizing certain derivative of the variance-objective was showed to work well in several applications 44
1(27 + N) + 8N
27 - N
p²
(11)
4y
for 7 1 M ΣM lq + k312. Since ? be rewritten as: Y N3= (EM
EE | + k) for q EM1 and k setting can be clearly done in time linear in M one-time procedure. Then the computation of hMme (M) follows. Compute-time per memory-vector remains (1s
Σド=
M
- k,, computing p² in the bi-directional
However Mnemosyne's temporal encoder applied uni mireeirctional attention. In the uni-directional case, we have: %/- = 1 q/ + k/ l| (1l + Ej=m ||k] 112 + 2(q²) (k()), where k(t) = Ej k for t = L,..., M. Instead of one 7, we now have M values Yt since not all memories are known at once. We can still achieve O, (M)) compute-time per 7r repeating the trick from the bi-directional case, but that would need be followed by the re-computation of @(k)) (with new p-parameter) for 1,... which of course is not possible since vectors [k]_ are not explicitly stored (and for a good reason computational benefits), see: Eq. 3
Thickening Mnemosyne's memory: To obtain efficient uni-directional Mnemosyne's memory
cell also for the 0F44-mechanise 3 replacing hMmo(!) (t) = (N,, V2) with H = ((Ng) (NN))eR pen, (V Σ,, Ar), where we have:
propose "thicken" in that setting the hidden state from Eq.
E = Ej=k' k', Ar- Ei=1 ||k| and furthermore: Nř, ve correspond to versions of N₂ and V
17

---

respectively, using parameter to define mapping ø. The set 52 is obtained by discretizing interval (0, 1) into a fixed number of chunks (and effectively quantizes p E (0. 1)). The strategy is now clear: when the new pattern comes, we first update the entire thickened state, and then compute p². We finalize by finding E closest to p" to transform an input and using for that the "slice" of the hidden state corresponding to p. We see that all these operations can be made efficiently with only e-multiplicative term (independent from the number of patterns M) in space and time complexity.
FAVOR+ mechanism, as FAVOR+, can also be adapted to its hyperbolic cosine variant. In practice FAVOR+ mechanism worked similarly to FAVOR+ yet the proper adaptation of the latter one was important, since (see: Sec. 4) this variant provides strongest theoretical guarantees for the capacity of the entire compact associative memory model.
A.2 The proof of the extended version of Theorem 4.3
We start by providing an extended version of Theorem enriched with the exact formula of the variance of ((Eand)) We prove it below. We borrow the notation from Sec. A.1
Theorem A.1 (storage of compact associative memories). Denote by E' ....."M EM f-1.+11 the memory-vectors, Assume that the Hamming distance between any two memory-vectors is at least TN for some >> O. Take some < p 3. Then the following is true for any memory-vector E for 1-e 2.= the expected I 1= 1,..; μ and any impul T E B(g', as long M S (xx(22N(T - 2p))
change of the energy of the compact associative memory system A ((rand) associated with flipping the value of the dimension of ê is positive if that operation increases the distance from its close neighbor g and is negative otherwise. Furthermore, the variance of A A(Erand) ) is of the form:
1 (Vi + V/2 2VV V42 - V/3 2V))
(12)
7
where:
Vi
Σ
V(E"! + 5 +2²,, V/2 =
Σ V + g +
Pn E
V/2 Σ
Vi + guz + + ) V4- Σ
exp((5"1) T e
(1...
#ire M
V/3 Σ
exp((E")] 2 ) Vs
Σ
" Exp((E"1)) (xx((()))))) (e))
(...2().....M)
#1.22E(),....
(13)
for E denoting E with one of its dimensions flipped and:
B2
(x) D (Pexp(-2NN) + +8A 7 exp
|x|1
(14)
2(1 - (A1)
Proof. Take a memory gl E +1N and an input @ E B(5', pN). Denote by neg(E" .)) a vector obtained from ? by replacing (i) with (1). Let study the change of the energy of the system flip the value of the ith dimension of the input ? since the sign of this change solely determines the update that will be made. We have the following:
((Erand) E(neg((..))) E....h))
(M E E.. = Esigual + Enoise:
(15)
where:
Exigmal Σ
( Z%b,,
(16)
Enoise
Σ
1
Σ
(Wh - 2%1,
(17)
r
k= pE(I
18

---

and furthermore: W = bb z = d for:
N exp( A Allwall (1wx)1))
a exp(
N eexp(Buje A A 1b2),
by D Dexp
(18)
N
Ck D Dexp(
expl Buy neg(E i) A 11x))
[w,,... Wy ~ (N IN) then, from the fact that Erand is the unbiased estimation of Ereg. we get:
E[Xa] exp((E)) Te), E| (ex((()) neg(E) @ i)),
(19)
B[W = (5*) )
E 2 = cxp(((") nog((") (es(E",))) 1)),
This is a direct consequence of the OPRF-mechanism introduced in 44). Variables: Xg. Ys. W and Z for ...,, M are simply unbiased estimators of the applying OPRF-mechanism. Let us now compute the expected change of the energy of the system:
values obtained
(Bran)) = [signall + E. B[Eniis)],
(20)
where:
E Esignall
Σ
[ E(Yx))
Σ R=1
(exp((e))TT) - - Tnog(= i1)
(21)
k=1
and
E Enoise] = -
-
E
(E[W?) - E[
(22)
Σ
(exp((E))TT) - neg( i))
MAL
We will first upper bound
E Eooisell
We have:
E Enoisell
1| < Σ
Σ
(exp((E") + T (
"neg(e". 1))
km
Σ
(exp(N 2(T p))) + exp(N(( 2(T p) 2
N })}
(23)
pe (I.....MM
< 2M exp(NN1 - 2((T - p) + N)
We will now consider two cases:
Case 1: (1) = <
In this setting, flipping the value of the ith dimension the input vector increases its distance from the close neighbor. Therefore in this case would like the energy change of the system be positive (so that the flip does not occur). From the Equation 21 we obtain:
E Esignail N
(cxp(N(1- 2p)) - cxp(NN( - 2p) - 2))) =
(24)
1=1
exp(N((1 21)) 1
19

---

Thus we obtain:
p) ))
E E (Brand)) 2 (1)(N(112))((111 -²- -
22 eM1xp(N(1 - 2
(25)
Therefore, if the following holds:
1 1
M S exp(2N(T - -2p))
(26)
2ca
then > O.
Case 2: (i) = -
In this setting, flipping the value of the ith dimension of the input vector decreases its distance from the close neighbor. Therefore in this case we would like the energy change of the system to be negative (so that the flip does not occur). From the Equation 21 we obtain:
E [esinnll] v
(exp( N(1 - 2p)) - exp(NV1 - 2p) + 2)))
(27)
exp(N(( 2p)) - -e²)
Thus obtain:
E| ([ran)))) S exp((((1-2p))) 2p)) 2) + 2M exp(N(1 2(r - p) + 2)
(28)
Therefore, if the following holds:
-
M S exp(2N(T 2p))
(29)
2c
then E Erand) )] O. Note that the bound from Inequality 26 is stronger than the from Inequality 29 That completes the proof of the first part of the theorem.
Now we will compute the variance of A Erand)) Denote:
Z+
Σ
(W %
(30)
(L.......M)
Note that if wi,..., wy are chosen independently then Z₁ for it I,..., r are independent. The following is true:
Var A (Eran)) Var( (assgna + Enoise) Var
Σ
(W 2
[E(L.....M
Var(
1
Σ
7k)
2
1
Var Zzk)
Var
Σ
(Wh -
(31)
k=1
k=1
k=1
(E(L
(W - 2 1)
E
Σ
-w
E
[E M}
(2....... M)
20

---

Therefore we have:
Var(A(Bip (Erand)
B +
E
B|2K 2k"1 2
(L..e((
2
Σ
ElW x x
k= A1-22Ef1.....
(32)
E
B|W [E +
E
(I...E(
P1PP2 E (}....M)
2
Σ
Σ
k=
Note that from the fact that our random feature map based estimators are unbiased, we get (as we already noted before in Equation 19 and put here again for Reader's 's convenience):
(33)
E B[Z = exp((E")) neg( i)),
Let us now define:
[x() = D' exp(-2N) cxp(Bw Tx - 4A|lw112).
(34)
Note that the following is true:
B[ W (g 1 + 542 + 2
1 - = ( 2 + e²² + 2neg E i)) E W _ (5 + t02 + - + neg ( i))
(35)
Thus it remains find closed-form formula for (x) for any given R.V.
From the proof of Theorem 3.1
444, we get for A <
B2
Elexp(A] Bw x)] = (1 2A) $ exp
2(1 - 2A) |xx||2
(36)
Thus obtain:
B²
(x() = D exp((22)(( + -A) exp
|
(37)
2(1 8A)
Plugging to Equation 32 formulae from Equation 33 and Equation 35 and utilizing Equation 37 for V, we obtain the formula for the variance from the statement of the Theorem.
B Experiment details
B.1 Warm-up for Mnemosyne and other optimizers: additional results
Preliminaries: At each timestep t, gradient f)((x: is input to the optimizer. The gradient is pre- processed as proposed in 2 Coordinate- wise Mnemosyne's using two temporal encoders is applied. The Mnemosyne's 's memory cell interfaces with the rest the system similarly to any RNN-cell. Each cell uses exponential discount factor 0.1, 16 random projections, 16 hidden dimensions and 1 attention head. The memory cell output is fed to a fully connected layer, returning the update to be applied to the NN parameters of the optimizee.
Meta-training: refer to training optimizer's parameters 0 as meta-training to distinguish from the optimizee NN training. Mnemosyne's optimizer is meta-trained on MNIST classification task with 3
21

---

2 struggles to minimize classification loss Figure 10: Validation loss curves when training MLP with Mnemosyne compared to other methods for MNIST image classification. Optimization curves for 4 different MLP architectures in this order: layer, 20 hidden dim, sigmoid activation), (2 layers, 20 hidden dim, sigmoid activation), (1 layer, 40 hidden dim, sigmoid activation), (1 layer, 20 hidden dim, relu activation) are shown.
small MLP and small VIT models. The optimizee MLPs are sampled from this hyperparameter distribution: E [1,2] hidden layers of size in range [20, 40) and sigmoid or relu activation function. The optimizee ViTs have 7 E [1. 3] layers, h E [1,3] heads, with hidden dimension in range [16, 64), mlp dimension in range [16, 64) and head dimension in range [8, 16). The optimizee task is train
the model for 100 steps on batches of 64 image-class examples.
Hybrid loss function to improve generalization: To promote generalization, we use the random- scaling trick proposed by [47). Mnemosyne's optimizer is meta- tett-trained by gradient descent using Adam optimizer with learning rate 1) = 3e to minimize a combination of two loss functions. The first is the task loss given by the sum of optimizee losses in truncated roll-out of 5 MNIST training steps. The other one is an imitation loss given by the mean squared error between Mnemosyne's updates and expert-optimizer (Adam) updates for same inputs. Importantly, this imitation loss is different from the one proposed in 13 which uses off-policy expert roll-outs for imitation. In our case, we provide expert supervision for the on-policy updates. This mitigates the problem of divergence from expert's 's trajectory, often observed in behaviour cloning. Our imitation loss acts as a regularizer which prevents Mnemosyne's optimizer from fvtr-fitting the optimizee task that it is trained on. We emphasize that expert's learning rate leexp = 3e -2 was not obtained via any tuning process.
Our optimizer model has minimal input feature engincering and our meta- veeaattainnng setup is significantly simpler than those considered in the literature [9]
1 47 71 Even so, successfully apply Mnemosyne's 's optimizer to a variety of tasks due to its efficient memory mechanism. Furthermore, Mnemosyne's 's memory cells can be easily combined with any of the existing L2L methods that use LSTMs for memory-encoding.
Results: After meta-training, Mnemosyne's optimizer was tested on NN training tasks with different NN architectures and datasets. Recall that Mnemosyne only saw one ML. task of MNIST classifier training for 100 steps during meta-traing veeaattaiinnng Fig. I0 shows that Mnemosyne can optimize MLPs with diffcrent NN archtitectures and activation functions on MNIST image classifier training. Note that, Mnemosyne converges significantly faster than popular analytical optimizers, RMSprop and Adam while retaining similar asymptotic performance. Mnemosyne can train NNs for long horizons of thousands of steps while baseline LSTM optimizer
beyond a few hundred steps.
The results were already presented in the main body of the paper (see: Sec. 5.1). We want to add that, as for experiments from Fig. TO. here Mnemosyne's optimizer faster than standard analytical optimizers and much more stable than LSTM optimizer. Fig. IT shows the benefit of using expert imitation-loss for long-horing horizon stability of the Mnemosyne's 's optimizer.
Transformers:
Our results on training Transformers with Mnemosyne naturally lead to the question of the role that Transformer-based optimizers can play in training Transformers architectures. is well known that Transformer training requires nontrivial optimization techniques [45] e.g. learning rate schedulers (for that reason SGD was replaced with Adam in Transformer-trainigg architectures training is slow, often prohibitively (unless the model is trimmed down, for instance by replacing long-range attention modeling with local attention of the controllable attention radius). Attention- based optimizers can potentially address this problem, since they improve convergence (and thus effectively reduce training time) even meta-trained on much simpler tasks as we show in
Furthermore, for larger
Fig. 3
22

---

10a
LSTM (trained with meta loss)
LSTM (trained with combined loss)
6 100
Mnemosyne (trained with meta loss)
Mnemosyne (trained with combined loss)
Loss
10
3 10
2 100
1000
2000
3000
4000
5000
Itterations
Figure 11: Impact of training the optimizer with combined meta and imitation loss can be seen in generalization to a long horizon rollout, All variants were trained only on length 100 rollouts.
Rog All (1m3) Rog
Rog (h=5c)
Piemosone (AVOR Mnemmsyne (AVORR1)
1001 Minemaggic
Rog Rog
Teaatoos:
Ablation Studies. Left: Comparison the Mnemosyne's 's linear CAM with regular attention memory blocks with different history cache lengths (f). Middle: Meta-training curves of Mnemosyne optimizer with FAVOR+ and FAVOR+ mechanism for CAM. Right: Meta-training curves of Mnemosyne optimizer with different kernel transformation functions for CAM.
Figure 12:
B.2 Mnemosyne's CAM mechanism vs regular attention
We have tried to use regular Transformer blocks to encode associative memory for Mnemosyne's temporal module. For applying regular attention to online optimizer autoregressively, a limited-length cache of historical gradients has to be maintained. A self-attention map over the history sequence is generated and used to encode memory. Fig. 12 (left) shows the meta-training curves for regular attention optimizers with different history cache lengths. As we increase the cache length, the performance improves and the memory requirement scales quadratically. Due to this limitation, could not implement a regular attention based optimizer with cache length more than 100. On the other hand, Performer's 's memory cell defining CAM attend to theoretically unbounded history and out-performs regular attention variants with fixed memory requirement.
B.3 Different RF-mechanisms: detailed look
Fig. (middle) compares the performance of Mnemosyne's optimizer applying FAVOR+ and FAVOR+ mechanisms in CAM. FAVOR+ mechanism provides strongest theoretical guarantees for the capacity of the associative memory model. It also leads initially to faster convergence, but asymptotically performs similarly as the FAVOR+ variant. Due to the simpler implementation of FAVOR+, we it for all experiments with Mnemosyne's optimizer.
Optimizers with both regular positive and hyperbolic random features kernel learn similarly, but the latter has much lower variance (see: Fig. D (right)) and thus it became our default choice.
23

---

B.4 Ablations over different kernel functions in CAM mechanism
In Fig. 13 we compare different kernel functions for linear attention in CAM. Our ablation study shows that FAVOR+ method outperforms other functions. This informs the choice of the kernel in all our experiments.
Square ReLU
FAVOR+
Meta Loss
E2086
herations
Figure 13: Comparison of the coordinate-wisg Mnemosyne with ReLU kernel, square kernel and FAVOR+ approximation of the softmax kernel. We present the meta-training loss the three variants. FAVOR+ mechanism outperforms other variants.
B.5 Ablations over different discount factors and number of RFs in CAM mechanism
In Fig. I4 we present detailed ablation studies over discount factors т as well as the number of random features applied by our default CAM mechanism leveraging hyperbolic cosine random features.
B.6 Benchmarking different depths the temporal module
Finally, ablations different number of temporal encoders in Mnemosyne's temporal block. We noticed that modest increase of the numher of encoders improves loss in meta-training and meta-training very deep variants is particularly challenging (as requiring much data). Since in this paper we decided to use simple meta-training strategies and furthermore increasing the number of temporal encoders did not lead substantial gains, we decided choose shallow temporal encoders' architectures. The results are presented in Fig. 15
B.7 Compute Resources Used
All Mnemosyne optimizer variants trained and tested TPU pod containing 4 TPU v3 chips with JAX Hundreds of rounds of training and inference were needed to compare different variations, tasks and meta-losses.
Minemsay (1-18)
Mosemooyne (r=bi)
2.75
Misamasyne
2.00
50
1.00
0.75
250
500
750
1000
1250
1500
1750
2000
3000
1250
1500
1750
2000
Figure 14: Left: The comparison of Mnemosyne applying different discount factors with Adam optimizer in meta-training (MLP optimization). Right: The comparison of Mnemosyne applying different number of random features in the hyperbolic cosine random feature mechanism used in CAM.

---

10
Anvencsynne: knd
Mieensoyne: k-d
25
25-
2.0
Mnamosyne: k=2
Mnamosyne:
Mnemosyne: k-6
15
Mnemasyne:
15
Mnemasyna:
10
05
500
3000
1500 1750
2000
250
750
1250
1750
Berations
Barations
Figure 15: Comparison of the meta- mocta-training loss for Mnemosyne variants applying different number temporal
encoders k. Since several variants on the left figure performs similarly, on the right figure we highlight top two.
The matlin Irainino ic pondluctad /ums the MT contimiiatin vatioun tnelre and MNIST slato
Man
Saae Lo4
ITageMx2012
Imageraden2 (TTLIZ)
(11 111
WEHS50
20
22
Figure 16: Coordinate- wise Mnemosyne across different ViT architectures and datasets, as described in Sec. Mnemosyne matches or outperforms optimal Adam variants without any hyperparameter tuning.
B.S.]
B.8
Coordinate-wise Mnemosyne versus hard-coded optimizers for larger ViTs
B.8.1
VIT last-layer fine-tuning
In this study, we benchmarked Mnemosyne on different sizes of ViT architectures: ViT-Base, ViT- Large and ViT-Huge (ViT-B(x), ViT-L(x) and ViT-H(x) respectively, where x defines the patch size), see: Tab: We used the coordinate-wise variant of the Mnemosyne. We run tests on the following datases: im0genet2012, places365 and caltech-birds-201 the ViT-architecture and used Adam expert with learning rate = 3e -2 as a regularizer (see: our discussion above on meta-training). The learning rate was not tuned in any way. In fact (as we show below) Adam optimizer applying this learning rate is characterized by the sub-optimal performance. We tried two versions Mnemosyne: (a) a variant that solely optimizes the last layer of ViT (reported in the main body) and (b) the hybrid variant, where Mnemosyne is used to optimize the weight-matrix of the last layer and Adam with learning rate 7) = -3 to optimize the bias vector. That learning rate was also not tuned in any particular way and, as before, if applied purely within Adam, produces sub-optimal results. The purpose of that last experiment was to assess how efficient the strategy of optimizing jointly with Mnemosyne and hand-designed optimizer is. The results are presented in Fig. 16 and Fig. I7 We see that: (a) Mnemosyne without any hyperparameter tuning matches or outperforms optimal Adam variants, (b) it also substantially outperforms Adam variant used as an expert in meta-training. This is valid for both: regular Mnemosyne as well as the hybrid version.
We were optimizing the last layer of
B.8.2
VIT multi-layer fine-tuning
Here, we used light version of coordinate- wise Mnemosyne using a single temporal encoder layer with hidden dimension 8. This reduced the memory requirement of the Mnemosyne optimizer state. We fine-tune ViT-B model on CIFAR-100 dataset with batch size 128. We were able to fine-tune last 2 transformer layers along with the embedding, cls and head layers with Mnemosyne. Rest of the model was fine-tuned with Adam (learning rate 3).. For comparison, the same baseline Adam variant is to fine-tune the complete model.
25

---

Irageors0112
21
IT 20
NbONA 12.5%
Mastocee:: 15.7%
place
Lerd
0110011
VITH0221
30
Figure 17: The results from Fig. 16 but narrowed down to the comparison between two Mnemosyne variants and Adam optimizer applying learning used in meta-training of these variants of Mnemosyne. The expert substantially underperforms in all the cases Mnemosyne as compared to the expert variant). This shows that Mnemosyne does not learn imitate expert.
explicitly put the gains coming from the variants of
Table 2: Hyperparameters for the different ViT models used in this paper
Model
Heads
Layers
Hidden Dim.
Dim.
Params
Patch Size
ViT-Base
12
768
3072
86M
ViT-Large (16)
24
16
1024
4096
307M
ViT-Large (32)
24
16
1024
4096
307M
32
ViT-Huge
32
16
1280
5120
632M
32
B.9 Tensor-wise Mnemosyne versus hard-coded optimizers for VIT-H
We finetuned the embedding and els layer of ViT-H (see Tab: for hyperparameter) using tensorwise (~ 1M params), while the head was trained using Adam. The rest of the transformer parameters are fixed to the pre-trained value for methods. The batch size was set 128 for all methods.
B.10 Super-Mnemosyne: combining coordinate- and tensor-wise strategies for ViTs
We finetuned the top-8 layers of the ViT-Base model (see Tab: 2 along with the head, cls and embedding layer before we ran out of memory ie ~ 50M parameters with batch size of 256. Large tensor such as: a) the MLP block withing each layer, b) the head layer was finetuned using lite version of coordinate-wise. Rest of the tensors were finetuned using tensorwise. The bottom 4 layers of the model were kept fixed for Mnemosyne. For Adam baselines we finetuned all layers.
B.11 BERTTpretraiiiin NLP Transformers with Mnemosyne
We trained the Bert base model, whose Hyperparameters are shown in Tab: 3
The details of the training dataset used is shown in Tab:44 We trained all parameters from scratch for all methods, with batch size of 512. For the Mnemosyne results shown Fig: trained all parameters except the token embedding using Tensorwise Mnemosyne (~ 86M parameters). The token embedding was trained using Adam with learning rate le For Adam baseline we trained all parameters.
Table 3: Hyperparameters for the Bert base model
Model
Heads
Layers
Hidden Dim.
MLP Dim.
Params
Compute
Loss
Bert-Base
12
12
768
3072
110M
4x2 TPUv3
MLM
B.12 Soft prompt-tuning massive T5XXL Transformers with Mnemosyne
use coordinate-wine Mnemosyne to prompt-tune T5XXL model 55
(see Table for
hyper-parameters) on SuperGLUEE benchmark. Batch size 32 was used. The length of the soft-prompt sequence was 30 and each soft-prompt vector was of size 4096, making the total number of trainable parameters 122880.
26

---

Table 4: Dataset used for pre training.
Dataset
# tokens
Avg. doc len.
Books 77
1.0B
37K
Wikipedia
3.1B
592
Table 5: Hyperparameters for the T5XXL model
Model
Encoder
Decoder
Heads
Head Dim.
Embedding
MLP Dim.
Params
Compute
Layers
I Layers
Dim.
T5XXI.
24
24
64
64
4096
10240
11B
2x2x4
TPUv3
C
Mnemosyne for training initial optimization conditions
Mnemosyne can be also applied to learn initial conditions for the optimization.
In this section, our considered model using ge is more general than the one presented in Eq. and is of the form given below for 9e
(g :902 ). 0 [01; 02] and context-vector c:
:o.و (c); 90. (J, x.,
(38)
X++1
ift> 0
The optimizer is now explicitly split into two-parts: (1) gbgit that learns initial optimization point from the context-vector, e.g. the image encoding the scene (as the case in learnable-MPC setting (72]), and (2) ghist that processes the history of the optimization steps, as we described before.
Depending on the application, either one or the other optimizer (or both) are turned on. Critically, both are encoded as light scalable Transformers. Optimizer 98, applies spatial bi-directional attention while hist uses the spatio-temporal variant, described in the main body of the paper.
Below we show how this paradigm can be used in Robotics to learn initial points for the MPC optimization.
C.1 Spatial Mnemosyne for initializing MPC optimizers
We applied Mnemosyne with bidirectional spatial attention for learning trajectory optimizers to be used for wheeled robot navigation in complex, photo-realistic simulated environments. The robot uses vision sensors for observing an occupancy grid of the environment and navigates using linear and angular velocity control. Navigation in challenging environments requires efficient high-speed robot control. Model Predictive Control (MPC) presents an efficient approach to the navigation problem, provided that the motion-planning with the environment model can be carried out within computational real-time limits (72). Motion planning with efficient trajectory optimizers such iterative Linear Quadratic Regulator (iLQR) 41 is way to implement MPC. However, in challenging layouts with narrow corridors and doors and with conservative safety and collision avoidance constraints, iLQR struggles converge the optimal trajectory. Sequential Quadratic Programming (SQP) [62] is often a more robust optimizer that can handle difficult non-lineat constraints. However, SQP significantly, sometimes 10 times, slower than iLQR and hence cannot be deployed in real-time settings. Deep Learning models have been shown to accelerate SQP by warm starting the optimization process
Preliminaries:
30
We learn Mnemosyne's 's optimizer to imitate SQP
Figure 18: Navigation in photo-realistic environment with MPC policy using Mnemosyne optimizer for initializing SQP solver.
27

---

behaviour and initialize it with an approximately optimal trajectory starting point from which SQP refines to an optimized and feasible trajectory.
Mnemosyne's optimizer, gimat, receives current robot pose Pr, a goal pose Pg and visual occupancy grid as the context, c. 'hee occupancy grid is processed by image encoder module, a ViT where the attention mechanism is approximated by bidirectional Mnemosyne memory. As in ViT, the occupancy grid is first pre-processed by convolution layer and then flattened to sequence. Each element (token) of the sequence corresponds to a different 5 x 5 patch of the original frame which is then enriched with positional encoding. The pre-procee pre-processed input then fed Mnemosyne attention and MLP layers of hidden dimension 64. The final embedding of one of the tokens chosen as latent representation of the occupancy grid, Loe- Pr. Pg and loc are concatenated and processed by an MLP which outputs the predicted action trajectory.
Training details:
An offline dataset of SQP optimization examples is collected by running MPC navigation agent in 2787 different environments. Each navigation run has 180 steps on average. For each MPC step, one instance of trajectory optimization with SQP was run and pair of input context and the final optimal trajectory XT was recorded. A total of ~ 500, 000 training examples were collected. Mnemosyne's 's optimizer was trained with supervised learning the SQP dataset by minimizing mean squared error between the SQP optimal trajectories and the predicted trajectories.
used to initialize SQP optimization. Without Mnemosyne initialization, SQP optimization was capped at maximum 10 iterations. It took average 4.78 iterations and 0.12sec for the SQP solution complete. With Mnemosyne initialization, SQP only run for 1 iteration to reach the optimal trajectory. SQP generates trajectory that satisfies kinematic, dynamic and safety constraints for the robot which transformer alone e.g. [77] cannot natively guarantee. This reduces the optimization time by more than half to 0.048sec on average which under real-time constraint. Mnemosyne inference and the rest for SQP iteration. A sequence of snapshots during navigation with Mnemosyne-SOP optimizer sample environment is shown in Fig. 18
Results: After training, the predicted trajectory from Mnemosyne
includes 0.01lsec for
28


Reviews:
[{"review_id": "eJsuUBHqaS", "reviewer": "Reviewer_V1vR", "summary": "The paper proposes to use transformers to learn to optimise. The proposed architecture relies on the performer architecture and leverage attention to encode long-term dependencies. The authors discuss two variants, a coordinate-wise and tensor-wise approach, to account for the dependencies between the parameters to optimise. They also propose to use a hierarchal pooling strategy to bypass the linear attention of performers. The authors compare their method to classical optimizers, such as ADAM, and LSTM-based optimisers. The paper also provides a theoretical analysis of the compact associative memory used by the method.", "questions": "I would like to ask the authors the following questions:\n\n1/ Can you provide a summary of the meta-training procedure and a discussion of the choices made? How robust is Mnemosyne to the specific procedure used and the set-up?\n \n2/ Can you explain why long-term dependencies are important when learning an optimizers? Can you provide evidence that the proposed architecture is able to capture these and leverage them?\n\n3/ Can the authors articulate the importance of the theoretical analysis?\n\n4/ Why did you consider the VeLo architecture in the experiments? Did you consider other options?\n\n5/ Can you provide details of the configuration of Mnemosyne in the experiments? How many of these did you try out and how did the results vary?", "limitations": "I would have liked the authors to be more vocally self-critical about the pros and cons of their proposed method. The paper reads like the approach is a sliver bullet, which obviously is not the case. I would have liked to see a discussion of the strengths and weaknesses of the method, an analysis of the robustness, an empirical validation of the claims (like importance of long-term dependencies), etc.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The authors propose a transformer-based optimiser by combining a topological and a temporal encorder. The resulting transformer is outperforming LSTM-based counterparts on the experiments conducted by they authors. The method is also competitive with widely used optimisers such as ADAM. The overall results are encouraging and the reported results cover both NLP and CV applications.", "weaknesses": "While the authors propose a transformer-based architecture to learn an optimiser, their choices are mainly driven by practical considerations. The authors do not motivate their choices beyond these and do not discuss trade-offs and/or other possible choices. The methodological contributions of this work are modest beyond combining existing components and the use of a hierarchical pooling strategy. While I have no doubt that making these components work well together is non trivial, the challenges to make them work together are not discussed. Hence, it appears that the paper does not do justice to work that has been done in the background. More importantly, the lack of details would make it difficult for other researchers to reproduce the results.\n\nThe authors provide a theoretical analysis of the compact associative memory used, but the importance and implications of this result is not discussed.\n\nWhile the experiments cover the NLP and CV domain, reported results were underwhelming. The authors conducted first small scale experiments where they compared to a wide range of off-the shelve optimizers. However, in the NLP and CV domain, they only optimised a subset of the parameters with the proposed approach, cast doubts on the broad applicability, its scaling properties and its robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">*choices driven mainly by practical considerations…implications of the theoretical analysis*\n\nSee \"Implications of theoretical analysis on design choices\" in the global rebuttal.\n\n> *discussion of other choices*\n\nWe chose Adam for non-learnable and LSTM for learnable optimizer baselines due to their popularity. See \"New experiments: Comparison with other learnable optimizers\" in the global rebuttal for experiments with more baselines.\n\n> *discussion of tradeoffs, self-criticality*\n\nSee \"Discussion of limitations\" in the global rebuttal.\n\n> *modest methodological contribution…combining existing components…the challenges to make them work together are not discussed…*\n\nWe justify that our contributions go beyond *“combining existing components”* as follows:\n1. Implementation of linear attention in auto-regressive (rather than bi-directional) setting has not been done before. It requires algorithmic modifications like exponential discount strategy (l. 169). For auto-regressive FAVOR++, additional algorithmic ideas were needed (see: l.679-692 and \"Implications of theoretical analysis on design choices\" in the global comment).\n2. Our hierarchical pooling encoder (HPE) mechanism in the tensor-wise Mnemosyne is *new*. HPE is different than other pooling techniques in Transformers (e.g. [Funnel-Transformers](https://arxiv.org/abs/2006.03236)), and relevant beyond LOs.\n\nCombining all these components was non-trivial (as the Reviewer noted) and we will address those challenges properly in the final version of the paper.\n\n> *the lack of details…a summary of the meta-training procedure*\n\nWe plan to publish the code upon the acceptance of the paper. We did our best to include *all the details* in the paper and Appendix. (See \"Summary of details in the Appendix\" in the global comment). Mnemosyne is robust since our meta-training is simple by design. More extensive meta-training can improve the generalization of Mnemosyne but that study is out of scope for our paper.\n\n1. Mnemosyne has been meta-trained on simple tasks (training toy-size MLP+ViT models), often not related to the downstream problem (full pre-training of base BERT 86M-parameter model or soft prompt-tuning of 11B+ T5XXL model etc.).\n2. Different meta-loss functions were tried (with/without regularizers).\n3. For the loss functions with regularizers, the hyperparameters of the regularization terms *are not tuned* at all.\n\nWe kindly ask the Reviewer if they would like to see any other details. We are happy to provide them during the discussion phase.\n\n> *reported results were underwhelming…they only optimised a subset of the parameters*\n\nIt is not true that we only optimized a fraction of parameters in all experiments. Our experiments include pre-training **full base BERT model with 86M parameters** (l.333-345), which was not done before with LOs (due to space complexity challenges) as well as several (rather than one) layers of ViT models (l.345-357). The results in Sec. 5.4 are for training **top 8 (out of 12) layers (~50M parameters)** of the ViT-Base model with the Super-Mnemosyne variant. We do think that showing that LOs can be successfully applied on such a scale is a significant contribution that was not achieved before.\n\nFurthermore, fine-tuning a subset of layers or soft prompt-tuning massive models (in the paper we successfully soft prompt-tune 11B+ T5XXL models, see: l.323-330) are very popular techniques of customizing large models. This is the case since pre-training from scratch or fine-tuning all the parameters is not possible without extensive computational resources. That is why we do think that our work is directly relevant to researchers leveraging modest computational resources which is well within the scope of the ML democratization efforts. The research on training a subset or a more compact set of parameters is actually one of the most rapidly growing areas of the present efforts on Transformers (as often providing competitive results) and led to several important methods proposed recently (such as [LORA](https://arxiv.org/abs/2106.09685) or [adapters]( https://arxiv.org/abs/1902.00751)). Even more recently (after our paper was written), the mechanism of [SubTuning](https://arxiv.org/pdf/2302.06354.pdf) was introduced. The authors explicitly state (in the context of fine-tuning) that :\n\n*\"... not all layers are created equal: different layers across the network contribute variably to the overall performance, and the optimal choice of layers is contingent upon the downstream task and the underlying data distribution…\"*\n\nFinally, our methods contribute to the field of stable long-range optimization with LOs (see: our response to the comment *\"long-term dependencies\"* below).\n\n> *long-term dependencies*\n\nSee \"Mnemosyne can perform stable long-horizon optimization\" in the global rebuttal.\n\n> *VeLO*\n\nWe chose to consider the VeLO architecture in our experiments for several compelling reasons:\n1. VeLO is the SOTA in the LO literature, making it a pertinent benchmark.\n2. VeLO’s feature engineering tricks and hierarchical model further strengthens it as a benchmark. These tricks improve the performance of the vanilla LSTM-based Los.\n\nOur objective was to show that we could improve VeLO’s performance by substituting the LSTM module with our compact associative memory (CAM) module. Note that for our experiments, we used a small-scale meta-training approach. This enabled us to clearly showcase one of the crucial differences between the LSTM and Mnemosyne approach to optimization: the former class of methods often struggle to generalize effectively to long-horizon evaluation tasks, while Mnemosyne mitigates this issue (see also: our discussion regarding long-term dependencies). By demonstrating the improvements achieved by simply replacing one memory mechanism with the other one, we highlighted the potential of the Mnemosyne’s memory model as an efficient mechanism for optimizing for long optimization horizons."}, {"review_id": "wCuiLTfBRq", "reviewer": "Reviewer_xFuP", "summary": "The paper introduces Mnemosyne, a new learnable optimizer that can train neural networks without task-specific tuning. The authors claim that it employs minimal resources and handles larger parameter sets, outperforming LSTM optimizers and matching state-of-the-art ones. Its efficacy is demonstrated in the fine-tuning of Vision Transformers, BERT models, and large T5XXL models. In addition, it also provides the theoretical analysis of the compact associative memory used by the proposed Mnemosyne.", "questions": "See the weakness.", "limitations": "The authors have discussed the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "（1）The proposed Mnemosyne introduces a novel approach by utilizing spatio-temporal low-rank implicit attention Transformers to create learnable optimizers. This innovative idea goes beyond traditional optimizers by eliminating the need for task-specific tuning in training neural networks. The idea seems interesting.\n\n（2）Mnemosyne has demonstrated superior performance in various use cases, such as fine-tuning Vision Transformers, pre-training BERT models, and soft prompt-tuning large T5XXL models. This shows its broad applicability and robustness", "weaknesses": "(1) The method's effectiveness on broader model training remains unknown as it only trains or fine-tunes a subset of Transformer layers. Moreover, given the Transformer's huge data requirements, its performance on larger datasets is also unclear.\n\n(2) As the method only trains or fine-tunes some layers of certain Transformer architectures, its title's emphasis on training Transformers seems somewhat overstated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the review and all the comments. We address them in detail below.\n\n> *The method's effectiveness on broader model training remains unknown as it only trains or fine-tunes a subset of Transformer layers*\n\nWe would like to clarify that our experiments involve in particular pre-training *full base BERT model with 86M parameters* rather than thee subset of model’s parameters (l.333-345), which to the best of our knowledge was never done before with learnable optimizers (due to space complexity challenges) as well as several (rather than one) layers of ViT models (l.345-357). The results in Sec. 5.4 are for training *top 8 (out of 12) layers* (*~50M parameters*) of the ViT-Base model with the Super-Mnemosyne variant. Those examples best illustrate the effectiveness of the spatial encode (and consequently, tensor-wise Mnemosyne)r which is one of the critical contributions of this work and can be used to overcome memory issues of the learnable optimizers relying on the coordinate-wise strategy. To summarize, we do think that showing that learnable optimizers can be successfully applied on such a scale is a significant contribution that was not achieved before.\n\nWe do not claim that our approach using spatial encoding resolves all memory issues of learnable optimizers. But we do think that it is an important step towards scaling up learnable optimizers to massive Transformer-models, as allowing to train architectures of sizes that were beyond the reach of the proposed before learnable optimizers.\n\nFurthermore, fine-tuning a subset of layers or soft prompt-tuning massive models (in the paper we successfully soft prompt-tune 11B+ T5XXL models, see: l.323-330) are very popular techniques of customizing large models. This is the case since pre-training from scratch or fine-tuning all the parameters is not possible without extensive computational resources. That is why we do think that our work is directly relevant to researchers leveraging modest computational resources which is well within the scope of the ML democratization efforts. We do not think this research is as important as training larger and larger models. The research on training a subset or a more compact set of parameters is actually one of the most rapidly growing areas of the present efforts on Transformers and led to several important methods proposed recently (such as [LORA](https://arxiv.org/abs/2106.09685) or [adapters](https://arxiv.org/abs/1902.00751)). Those methods have been proven to be competitive with full model fine-tuning for many downstream tasks, yet they clearly require much smaller computational resources. Even more recently (after our paper was written), the mechanism of [SubTuning](https://arxiv.org/pdf/2302.06354.pdf) was introduced. The authors explicitly state (in the context of fine-tuning) that :\n\n*\"...We observe that not all layers are created equal: different layers across the network contribute variably to the overall performance, and the optimal choice of layers is contingent upon the downstream task and the underlying data distribution…”*\n\nWe do hope that our response addresses all the concerns raised by the Reviewer regarding the broader impact of our method.\n\n>*…given the Transformer's huge data requirements, its performance on larger datasets is also unclear*\n\nOur temporal modules consists of at most two temporal encoders (in most applications just one), since our ablation studies showed that deeper architectures provider negligible gains (see: Fig. 14) given small set of training tasks (which we wanted to keep to make the entire meta-training procedure simple). Thus the underlying Mnemosyne’s encoding entire memory mechanism is compact. We ended up applying very simple meta-training on a couple of toy-size MLP and ViT tasks (see: meta-training paragraph in Sec. B.1). This was the case also when the target downstream task was not related to the meta-training tasks, e.g. for base BERT pre-training or T5XXL soft prompt-tuning. In those cases Mnemosyne was operating on the 86M and 11B+ size models respectively. We did not observe any need to scale up meta-training to a larger set of tasks to efficiently optimize larger models. That was also the case when we increased the number of fine-tuned ViT layers (see: l.345-357 in Sec. 5). The only modification of meta-training was introducing larger weight-matrices in meta-training (so that the spatial encoder learns how to encode them), yet the sheer number of tasks and their difficulty remained the same.The entire meta-training consisted of only 100 iterations. All these observations clearly show that Mnemosyne is not data-hungry in meta-training and meta-training does not become a bottleneck when downstream optimization tasks become bigger. We will clarify it in the final version of the paper, by providing a separate paragraph on the scalability of meta-training in the main part of the paper.\n"}, {"review_id": "TehybJT9uN", "reviewer": "Reviewer_JvEY", "summary": "This paper proposes to use transformers to learn optimization policies for training deep learning models. Here, the sequence corresponds to the state sequence (representing the internals of the model) in the trianing process. The intuition is that when we collect many sequences from many training processes, we can extract knowledge from these sequences via using transformers. \n", "questions": "\n- When we do this kind of meta learning, does that mean that the model states need to be stored for every batch? That seems very expensive. \n- Can you comment environment we need to run the proposed solution, e.g., a few GPUs & a few days or 1000 A100s for a month? \n", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It appears to be a quite challenging project. Merely to make this kind of meta-learning works seem to be already very compute-heavy. Also, the vanilla version of the transformer has quadratic space and time complexity so it will be slow. \n\nThis work developed a quite large collections of careful hacks to make the idea work. For example, performance that uses low-rank matrix to approximate attention, and additional pooling is performed to further trim down computations. There are also other low level tricks like using exponential moving average. \n\nThere are also some theoretical results related to associate memory. \n\nThe sheer amount of work seems sufficient for a neurips acceptance. \n", "weaknesses": "I honestly cannot quite interpret the associate memory part. It seems to suggest the proposed architecture has sufficient expressive power. I thought the more important question is to understand why \n\nAlso, the lack of source code seems a minus. In general, I feel the quality of source code plays an increasingly important role in final decision, e.g., you can clearly see some system effort is more serious than others. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the review and all the comments.We address them in detail below. We are happy to hear that the Reviewer appreciates the amount of work conducted in this work and evaluates it as sufficient for NeurIPS acceptance.\n\n> *cannot quite interpret the associate memory part*\n\nOur theoretical results show that the applied memory mechanism is capable of storing (on average) optimization history of the exponential length (as it acts as the large capacity associative memory model).\nWe have run additional tests comparing the attention linearization mechanism with the popular Performer-ReLU and Performer-square variants. We have not put them in the original manuscript due to space constraints, but refer to them in the rebuttal (see the pdf attached to our global comment). Our results show that sub-exponential (on average) storage capacity of the associative memory model related to the ReLU and square variants does not admit the theoretical analysis from Sec. A.2 hurts performance. As summarized above, our theoretical results from Sec. A.2 show that the provided memory mechanism is capable of storing (on average) exponential number of points from the optimization history, effectively enabling the model to take into account much longer optimization history and consequently, improving overall optimizer’s performance.\n\nWe want to emphasize that our theoretical analysis critically relies on the fact that considered random feature map mechanisms approximate the softmax kernel. Therefore it clearly distinguishes those methods from other random feature map mechanisms (also applied in Performers) that linearize non-softmax kernels. As we explain above, this distinction is confirmed by the additional conducted experiments that we include in the rebuttal. Thus the importance of our theoretical analysis lies in the fact that it leads to the choice of the right linearized associative memory mechanism without any additional ablation studies.\n\nIn the final version of the paper, we will make those theoretical insights much more explicit in the main body of the paper and include aforementioned results on comparing ReLU- and square- variants with the hyperbolic cosine one.\n\n\nIn the rebuttal, we have added the comparison of our associative memory model with the optimizers replacing regular LSTM memory unit with the powerful S4 memory unit based on [state space machines](https://arxiv.org/pdf/2111.00396.pdf), as an alternative to LSTMs (and Transformers). Our results are presented below:\n\nExperiment 1: imagenet-ViT-H(patch size = 14), last layer fine-tuning (500 iterations):\n\n| Methods | Accuracy(%) |\n|--------------|:-----------:|\n| Mnemosyne | **72.3** |\n| S4 | 7.5 |\n\nExperiment 2: imagenet-ViT-H(patch size = 8), last layer fine-tuning (500 iterations):\n\n| Methods | Accuracy(%) |\n|--------------|:-----------:|\n| Mnemosyne | **81.4** |\n| S4 | 35.0 |\n\nOur experiments confirm our previous findings. We show in particular that S4-based architectures fail to fine-tune vision Transformer models under consideration, as opposed to Mnemosyne. Thus RNN-based approaches seem to underperform as compared to attention-based methods for the learnable optimization.\n\n> *the lack of source code*\n\nThank you for the comment. We plan to publish the code upon the acceptance of the paper. We did our best to make sure that *all the details* are included in the paper and put many of them in the Appendix, in particular:\n\n1.  detailed discussion of the meta-training phase: we believe we have provided all the details regarding meta-training in the meta-training paragraph of Sec. B.1: (a) the process for creating mlp and vit tasks for meta-training (l. 760-764), (b) number of iterations used in meta-training (l. 764), (c) the particular form of the hybrid-loss function, including Adam regularizer with the explicitly given non-tuned learning rate (l. 765-776);\n2.  ablations studies over varying tau-parameters for the exponential discount strategy (Fig. 13 - left), leading to our default choice of tau;\n3.  ablation studies over different number of applied random features r, leading to our default choice of r (Fig. 13 - right);\n4.  ablation studies over Mnemosyne architectures with varying number of temporal encoders (Fig. 14), leading to the conclusion that shallow encoder modules suffice for Mnemosyne;\n5.  used computational resources (Sec. B.6);\n6.  hyperparameters for different models that Mnemosyne was tested on in the paper (Table 1);\n7.  hyperparameters for the base BERT models that Mnemosyne was tested on in the paper (Table 2);\n8.  hyperparameters for the T5XXL models that Mnemosyne was tested on in the paper (Table 4). \n\n> *When we do this kind of meta learning, does that mean that the model states need to be stored for every batch? That seems very expensive. Can you comment on the environment we need to run the proposed solution, e.g., a few GPUs & a few days or 1000 A100s for a month?*\n\nThank you for the comment. All Mnemosyne variants were trained and tested on a TPU pod containing 4 TPU v3 chips with JAX. Mnemosyne is meta-trained on an MNIST classification task with 3 small MLP and 3 small ViT models. This makes the meta-training very efficient. Saving optimizer states for these small models has a very small memory and computational footprint. A meta-training run takes less than an hour on average. More details are given in the Appendix."}]

Author rebuttal:
# New experiments: Comparison with other learnable optimizers
In the rebuttal we provide the comparison with two other classes of learnable optimizers:
1. Those using powerful [S4 memory units](https://arxiv.org/pdf/2111.00396.pdf) based on state space machines as an alternative to LSTMs (and Transformers).
2. The [Lion](https://arxiv.org/abs/2302.06675) family of LOs using a different approach of symbolic programming.

Our results are presented in the attached pdf. In the last experiment (Table 3), we skip S4 due to poor performance on the previous tasks. Mnemosyne is clearly the best among all three variants (in contrast to Lion, does not require hyperparameter tuning; also it applies simplistic meta-training: 4TPU v3 chips, see: Sec. B.1).

# Implications of theoretical analysis on design choices
Our paper provides much deeper theoretical insight than most others on learnable optimizers (LOs) (that focus on the empirical results). We have provided extensive theoretical and practical motivations for our design choices in the paper. We will make it more explicit in the final version:
1. The choice of the random feature (RF) mechanism is motivated by strong theoretical guarantees for the capacity of the corresponding associative memory (AM) (entire Sec. 4 and Sec. A.2). We are not aware of similar results for LSTM-based LOs. Other RF mechanisms that do not admit our theoretical analysis clearly underperform, as we show in the additional experiments in the rebuttal (see the attached pdf file). We compare there with Performer-ReLU and Performer-square RF mechanisms. Thus the importance of our theoretical analysis lies in the fact that it leads to the choice of the right linearized AM mechanism without any additional ablation studies.
2. Our default hyperbolic cosine RF mechanism (see: ll.186-191 and Sec. A.1) is motivated by the provably reduced estimation variance. We explicitly say it in l. 654-656.
3. We propose a novel autoregressive FAVOR++ (AFAVOR++) variant (see: Sec. A.1.1, l.679-692).
4. In Fig. 12 we show that: (a) hyperbolic cosine RFs are more robust than regular positive RFs (last plot), (b) AFAVOR++ provides faster convergence (middle plot).
5. We introduce novel exponential discount strategy (EDS) (l.169) to smoothly discount the past, theoretically motivated and compatible with Performers (see: Fig. 13 (left) for the detailed empirical verification of EDS).

# Summary of details in the Appendix
1. detailed discussion of the meta-training phase (meta-training paragraph of Sec. B.1), including:
	* the process for creating MLP and ViT tasks for meta-training (l. 760-764),
	* number of iterations used in meta-training (l. 764),
	* the particular form of the hybrid-loss function, including Adam regularizer with the explicitly given non-tuned learning rate (l. 765-776);
2. ablations studies over varying tau-parameters for the exponential discount strategy (Fig. 13 - left), leading to our default choice of tau;
3. ablation studies over different number of applied random features r, leading to our default choice of r (Fig. 13 - right);
4. ablation studies over Mnemosyne architectures with varying number of temporal encoders (Fig. 14), leading to the conclusion that shallow encoder modules suffice for Mnemosyne;
5. used computational resources (Sec. B.6);
6. hyperparameters for different ViT models that Mnemosyne was tested on in the paper (Table 1);
7. hyperparameters for the base BERT models that Mnemosyne was tested on in the paper (Table 2);
8. hyperparameters for the T5XXL models that Mnemosyne was tested on in thee paper (Table 4).

# Mnemosyne can perform stable long-horizon optimization
In our experiments, we made a noteworthy observation that Mnemosyne optimizers outperform their LSTM-based counterparts in terms of achieving a stable optimization phase over extended optimization horizons (see: Fig. 10 and Fig. 11). This results bears crucial importance since our goal is to achieve strong generalization from *small-horizon meta-training* to *long-horizon evaluation*, a notoriously difficult problem in the field of LOs. Prior attempts to address this challenge involved increasing the meta-training budget by including long-horizon tasks and made meta-training more complicated. Our approach takes a different direction that we do think about as novel. By incorporating a more expressive memory mechanism, we demonstrate that it is possible to learn highly generalizable optimization algorithms without the need of more complicated meta-training. We do believe that this finding opens up new possibilities for efficient and effective optimization in long-range scenarios and is an important advancement in the field. It also contributed to the efforts of democratizing ML.

# Discussion of limitations
1. Tensor-wise Mnemosyne is motivated by the limitations of coordinate-wise LOs (including coordinate-wise Mnemosyne). In Sec 3.2.1, we say:
	> *In that setting, even linear attention might not suffice. To address it, we introduce an additional hierarchical pooling mechanism, leading to the hierarchical pooling encoder (HPE).*
2. In the experiments (Sec. 5.3), we *explicitly* mention the limitations of the coordinate-wise Mnemosyne:
	>"Coordinate-wise Mnemosyne was not applicable here due to the memory footprint and thus it became a practical test for the memory efficient tensor-wise variant."
3. The tensor-wise variant is one way to address this problem, but it is not perfect. We say in Sec. 5.3:
	>"Although tensor-wise Mnemosyne can scale up to the full ViT model, here we freeze large tensors in the Transformer layers and fine-tune the rest. Training large tensors with tensor-wise Mnemosyne requires data and compute intensive large scale meta-training (to ensure that spatial encoders learn how to encode long sequences well). This will be the focus of the future work."
4. In the Broader impacts & limitations section we discuss how to further develop our meta-training strategies.

Meta-review: This work proposes a transformer-based optimiser, which was shown to outperform LSTM-based counterparts. Experiments indicate that the learned optimizer is competitive with ADAM. Authors considered problems in both the NLP and the CV domain. The authors adequately addressed the concerns raised by the reviewers during the rebuttal and I would encourage the authors to incorporate the feedback they provided. The paper will be stronger by removing the ambiguities and imprecisions in the original manuscript.
Decision: Accept (poster)