SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Deepali Jain *
Google DeepMind,
jaindeepali@google.com
Krzysztof Choromanski *
Google DeepMind and Columbia University,
kchoro@google.com
Avinava Dubey *
Google Research,
avinavadubey@google.com
Sumeet Singh
ssumeet@google.com
Vikas Sindhwani
sindhwani@google.com
Tingnan Zh
tingnan@google.com
Jie Tan
jietan@google.com
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.
Learning-to-learn (L2L) systems [66, 63, 4, 5, 52, 28, 60, 74, 2, 12] used to train machine learning (ML) optimizers can be thought of as a natural lifting (to the optimizer-space) of the idea that has revolutionized ML: replacing hand-engineered features with the learned ones.
The idea to train ML optimizers is tempting, yet the lift to the optimizer-space comes at a price: the instances used to train such systems are the optimization problems themselves. Generalization in such a setting means: the ability to transfer knowledge to "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 a subject on its own. Indeed, the literature on meta-learning is voluminous [73, 6, 53, 76, 1, 39, 24] and of critical importance in many disciplines such as Robotics, where transfer knowledge from simulator to hardware is a notoriously difficult problem [42, 32, 57].
A standard approach to learning optimizers is to cast it as a sequential decision problem, where a function  $f$  called an optimizee is optimized via another function  $g_{\theta}$  (an optimizer) with learnable parameters  $\theta$ . Function  $f$  can take as input  $\mathbf{x}$  the parameters of a neural network (NN) and output its corresponding test loss on a given task. The optimizer  $g_{\theta}$  updates  $\mathbf{x}$  as:
$$
\left\{ \begin{array}{l} \mathbf {x} _ {0} \leftarrow \text {o p t i m i z a t i o n i n i t i a l p o i n t}, \\ \mathbf {x} _ {t + 1} = g _ {\theta} (f, \mathbf {x} _ {0}, \dots , \mathbf {x} _ {t}) \quad \text {i f} t > 0 \end{array} \right. \tag {1}
$$
Function  $g_{\theta}$  is either trained by minimizing the meta-loss objective  $\mathcal{L}_{\theta}$ , which is usually a sum of  $f$ -losses over certain time horizons, with supervised/reinforcement 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  $(\mathbf{x}_0, \dots, \mathbf{x}_t)$ , but its processed-version, e.g. the sequence of the corresponding gradients  $(\nabla f(\mathbf{x}_0), \dots, \nabla f(\mathbf{x}_t))$  if  $f$  is differentiable (which does not need to be the case [50]). In practice, the processed-versions often have much richer structure [49]: "...various rolling statistics including...loss 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  $g_{\theta}$  (which is a subject on its own) and different training strategies of  $g_{\theta}$ . Our interest is in the core design of  $g_{\theta}$  since it acts as a memory-based system. Indeed, most models of  $g_{\theta}$  leverage recurrent NN cells, such as LSTMs [2, 58, 57], that keep the history of the optimization-rollout in the form of a compact learnable latent state. In addition, due to the fact that inputs  $\mathbf{x}$  to  $f$  are usually high-dimensional (e.g. neural networks' parameter-vectors),  $g_{\theta}$  is often factorized to process independently different dimensions of  $\mathbf{x}$  [2].
The Transformer-revolution in ML set the stage for a very different way to model sequential data and in general: to model memory - the attention mechanism [67, 19, 54, 8, 22, 11]. It is natural to ask whether attention architectures can be used to replace LSTM memory cells in L2L 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.
Contributions. 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 leveraging minimal computational resources, (c) matches accuracy-wise SOTA hand-designed optimizers with carefully tuned hyper-parameters (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 is of critical practical importance. Furthermore, Mnemosyne provides space complexity comparable to that of its standard 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 also conduct a thorough theoretical analysis of the compact associative memory used by Mnemosyne, to the best of our knowledge never done before.
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-temporal system, (c) hierarchical spatial attention mechanism to further reduce memory footprint of the spatial attention encoders that can be thought of as a novel attention pooling mechanism (see: 3.2.1). Obtained L2L mechanism effectively acts as a 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.2). 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.
The research on L2L systems involves a plethora of techniques: curriculum learning (e.g. incrementally 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 [71] and more [9, 51, 15, 68].
Regular Transformers are recently considered in this setting, in particular to tune hyperparameters [16], to learn BFGS-type optimization for motion reconstruction [26] or as memory-systems 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 energy-models [36] (while used in the causal attention setting). Regular Transformers can be thought of as differentiable dictionaries applying powerful associative memory mechanisms [37, 56], i.e. modern Hopfield networks [29] with exponential memory. Linear low-rank attention mechanisms are their compact variants [61] 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 as 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, A.1) providing particularly low variance of the softmax-kernel estimation, a key ingredient of those modern associative memory models.
3 Learning to learn with spatio-temporal attention
In this section, we give the description of Mnemosyne. Since the system consists of several components, 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 optimize  $f: \mathcal{D} \to \mathbb{R}$ . In the simplest case, we can take:  $\mathcal{D} \subseteq \mathbb{R}^d$  for  $d \in \mathbb{N}_+$ . However it will be convenient to impose a tree-like structure on the elements  $\mathbf{x} \in \mathcal{D}$ . We will assume that there exists a tree  $\mathcal{T}_f$  such that for every  $\mathbf{x} = (x_1, \dots, x_d)^\top \in \mathcal{D}$ , the set  $\{x_1, \dots, x_d\}$  is partitioned into non-empty subsets  $S_1, \dots, S_l$  corresponding to different leaves of  $\mathcal{T}_f$ , where  $l$  stands for their number. Not only does  $\mathcal{T}_f$  define partitioning of  $\{x_1, \dots, x_d\}$  through the subsets residing in different leaves, but it also imposes a natural hierarchy. The construction might seem artificial at first glance, but we have a good reason to introduce it early on - it is a predominant description of the NN structure (our main interest in this paper is to optimize NNs with learnable optimizers thus we identify elements of  $\mathcal{D}$  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_i$  are just their flattened representations) and tree-induced hierarchy describes how the NN layers/modules emerge from those individual tensors<sup>2</sup>. This more general framework covers also the setting  $\mathcal{D} \in \mathbb{R}^d$ , where  $\mathcal{T}_f$  is a star-tree with different leaves corresponding to different variables  $x_i$ .
3.2 Topological encoder
Standard L2L systems [2] act independently on the individual parameters to be optimized. We refer to this strategy as a 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  $\mathcal{T}_f$  being a star-tree (see: Sec. 3.1). However it comes at a 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-wise framework (see: Sec. 5.2), but supports an arbitrary  $\mathcal{T}_f$ , 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  $\mathbf{T}$  from each leaf of a given input  $\mathcal{T}_f$ . Tensor  $\mathbf{T}$  is first flattened/vectorized and a representation vector  $\mathbf{r}$  is associated with each parameter of  $\mathbf{T}$ . Since
2An example of the instantiation of this mechanism is a pytree structure that is a default representation of the neural network parameters used in JAX [25] - a popular machine learning framework used to transform numerical functions and train neural networks.
in this paper we minimize feature engineering for the L2L systems, our default choice for  $\mathbf{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  $S_{\mathbf{T}}$  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  $\mathbf{T}$ .
3.2.1 Compactifying topological encoder with the hierarchical pooling
Using bi-directional linear attention from Performers is justified by the fact that sequence  $S_{\mathbf{T}}$  can be in practice very long. In fact it can easily surpass  $1M$  tokens (for instance if  $\mathbf{T}$  is a square weight-tensor corresponding to two consecutive layers of size  $> 1K$  each). In that setting, even linear attention might not suffice. To address it, we introduce additional hierarchical pooling mechanism, leading to the hierarchical pooling encoder (HPE). Sequence  $S_{\mathbf{T}}$  is first split into chunks of a fixed length  $L \in \mathbb{N}_{+}$ :  $S_{\mathbf{T}}^{1}, S_{\mathbf{T}}^{2}, \ldots$  (the last chunk might be shorter). Topological encoding is applied in parallel to each chunk. The procedure is repeated for the resulting sequence of the topological encodings, which is already shorter by the multiplicative factor of  $L$  (potentially with a different  $L$  even though here we assume that  $L$  is the same). The total number of repetitions is in practice a small constant  $h_{\mathrm{pool}} \in \mathbb{N}$  and leads to the final sequence of length  $l = \frac{\mathrm{len}(S_{\mathbf{T}})}{L^{h_{\mathrm{pool}}}}$ , where  $\mathrm{len}(S_{\mathbf{T}})$  is the original length and  $l$  as a small constant. If  $\mathrm{len}(S_{\mathbf{T}})$  is small enough, hierarchical pooling is not applied. The resulting  $l$ -length sequence of latent  $d$ -dimensional encodings is output as a topological encoding and fed to the temporal encoder defined below. We refer to different tokens of that sequence as meta-tokens,  $\mathcal{M}$ .
3.3 Temporal encoder: Compact Associative Memory (CAM) model
Preliminaries. The temporal module consists of one or more temporal encoders stacked together that process the meta-tokens,  $\mathcal{M}$ . It treats the length  $l$  of the sequence  $\mathcal{M}$  as a batch size  $b$ . For a given meta-token, denote by  $\{\xi^{\mu}\}_{\mu = 1}^{M} \subseteq \mathbb{R}^{d}$  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  $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{N \times d}, \mathbf{W}_V \in \mathbb{R}^{d \times d}$  as follows:
$$
\mathbf {q} ^ {\mu} = \mathbf {W} _ {Q} \xi^ {\mu}, \quad \mathbf {k} ^ {\mu} = \mathbf {W} _ {K} \xi^ {\mu}, \quad \mathbf {v} ^ {\mu} = \mathbf {W} _ {V} \xi^ {\mu} \tag {2}
$$
Take a kernel  $\mathrm{K}:\mathbb{R}^N\times \mathbb{R}^N\to \mathbb{R}$ , and its linearization:  $\mathrm{K}(\mathbf{x},\mathbf{y}) = \mathbb{E}[\phi (\mathbf{x})^\top \phi (\mathbf{y})]$  for some (randomized)  $\phi :\mathbb{R}^N\to \mathbb{R}^r$ . We refer to  $\phi (\mathbf{x}),\phi (\mathbf{y})$  as random feature (RF) vectors and define a hidden state encapsulating memory of the system of first  $t$  patterns as:
$$
\left\{ \begin{array}{l} \mathbf {N} _ {t} = \sum_ {\mu = 1} ^ {t} \lambda_ {t} (\mu) \phi \left(\mathbf {k} ^ {\mu}\right) \left(\mathbf {v} ^ {\mu}\right) ^ {\top} \in \mathbb {R} ^ {r \times d}, \\ \Psi_ {t} = \sum_ {\mu = 1} ^ {t} \lambda_ {t} (\mu) \phi \left(\mathbf {k} ^ {\mu}\right) \in \mathbb {R} ^ {r} \end{array} \right. \tag {3}
$$
A discount-function  $\lambda_t: \mathbb{R} \to \mathbb{R}$  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  $\mathbf{h}_{\mathrm{Mne}}(t) = (\mathbf{N}_t,\Psi_t)$  is of size independent from the number of its implicitly stored patterns  $t$ . When patterns are added,  $\mathbf{h}_{\mathrm{Mne}}(t)$  needs to be efficiently updated on-the-fly. It is easy to see that this can be done for the exponential discount strategy:  $\lambda_{t}(\mu) = \exp (-\tau (t - \mu))$  with  $\tau \geq 0$  ( $\tau = 0$  turns off discounting). We have the following:
$$
\mathbf {N} _ {t + 1} = \exp (- \tau) \cdot \mathbf {N} _ {t} + \phi (\mathbf {k} ^ {t + 1}) (\mathbf {v} ^ {t + 1}) ^ {\top}, \tag {4}
$$
$$
\boldsymbol {\Psi} _ {t + 1} = \exp (- \tau) \cdot \boldsymbol {\Psi} _ {t} + \phi (\mathbf {k} ^ {t + 1})
$$
With the definition of the temporal encoder's hidden state, we can now explain how it acts on the input vectors  $\xi \in \mathbb{R}^d$ . New vector  $\xi' = \xi + \Delta \xi$  is obtained as follows, where  $\mathbf{q} = \mathbf{W}_Q \xi$ :
$$
\Delta \xi = \frac {\mathbf {N} _ {t} ^ {\top} \phi (\mathbf {q})}{\phi (\mathbf {q}) ^ {\top} \Psi_ {t}} = \sum_ {\mu = 1} ^ {t} \frac {\lambda_ {t} (\mu) \phi (\mathbf {q}) ^ {\top} \phi \left(\mathbf {k} ^ {\mu}\right)}{\sum_ {\beta = 1} ^ {t} \lambda_ {t} (i) \phi (\mathbf {q}) ^ {\top} \phi \left(\mathbf {k} ^ {i}\right)} \mathbf {v} ^ {\mu} \tag {5}
$$
Vector  $\xi^{\prime}$  can be computed in time  $O_{M}(1)$  and is given as a convex combination of value vectors  $\mathbf{v}^{\mu}$ , with coefficients proportional to approximated kernel values, but modulated by the discount-function.
For  $\mathbf{W}_Q = \mathbf{W}_K = \mathbf{W}_V = \mathbf{I}_d$ ,  $\lambda_t \equiv 1$  and with exact kernel values  $\mathrm{K}(\mathbf{q}, \mathbf{k}^i)$  in Eq. 5 (rather than their approximated versions  $\phi(\mathbf{q})^\top \phi(\mathbf{k}^i)$ ), dynamical systems defined by Eq. 5 become effectively Hopfield networks and, as energy-based models with energies given as  $E(\xi; \{\xi^\mu\}_{\mu=1}^t) = -\sum_{\mu=1}^t \mathrm{K}(\xi, \xi^\mu)$ , retrieve memory-vectors upon energy-minimization-driven convergence [56]. The retrieval quality depends on the kernel, with the softmax-kernel  $\mathrm{K}(\mathbf{x}, \mathbf{y}) \stackrel{\mathrm{def}}{=} \exp(\mathbf{x}^\top \mathbf{y})$  providing particularly strong theoretical results. For arbitrary  $\mathbf{W}_Q$ ,  $\mathbf{W}_K$ ,  $\mathbf{W}_V$ , but still with  $\lambda_t \equiv 1$  and exact kernel values, Eq. 5 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  $\lambda_{t}$ ) compact (since it provides efficient hidden state and input update-rules via linearized kernels from Performers) associative memory model (CAM) (as opposed to the regular associate memory model from Hopfield networks).
In Mnemosyne, we choose softmax-kernel as K, since it is a default choice for Transformers. We observed that the FAVOR++ mechanism defining  $\phi$  from [44], and denoted by us as  $\phi_{F++}$ , provides the most robust performance, but since it is harder to implement in the temporal encoder setting (see: discussion in Sec. A.1), in practice we apply the so-called hyperbolic cosine random features from [18], performing very similarly, as  $\phi_{F++}$ . More details including in particular exact definitions and ablation studies over different random feature mechanisms can be found in Sec. A.1, B.3.
3.4 Putting it all together
Details of the hierarchical pooling encoder (HPE) and the compact associative memory (CAM) modules are depicted in Fig. 1. 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  $\mathbf{r}_i$  (see: Sec. 3.2) of every parameter  $x_i$  is processed separately in parallel by a CAM module followed by an MLP layer to produce the update  $\Delta x_i$ . 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 a 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  $\mathcal{S}_{\mathrm{T}}$  in the parameter-tree is processed as a whole by an HPE to produce meta-tokens  $\mathcal{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  $\mathcal{L}$  has the same shape as  $\mathcal{M}$ . It is transformed into a fix-sized encoding  $\mathbf{e}$  by a spatial attention encoder (SPE). This encoding is broadcast to all the input tokens of the tensor via concatenations with vectors  $\mathbf{r}$ . Each resulting vector  $\mathbf{r}' = \mathbf{r} \odot \mathbf{e}$  is processed by a single MLP-layer to get the update tensor  $\Delta \mathbf{T}$ .
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 \times 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 as follows. Each tensor is vectorized and chunked into sub-sequences that are spatially encoded by the bi-directional Performers. Presented pooling mechanism consists of two layers. This results in the input tensors  $\xi^0, \xi^1, \xi^2$  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  $\mathbf{W}$  matrices to keys (shown above) and queries and then non-linearly transformed via the  $\phi$ -mapping. The transformed variants are leveraged by the associative memory.
Figure 2: Two modes of Mnemosyne application: (a) coordinate-wise and (b) tensor-wise.
4 The theory of Mnemosyne's compact associative memory