SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {V a r} \left(\Delta \left(E _ {\text {r a n d}}\right)\right) = \frac {1}{r} \left(V _ {1} + V _ {2} - 2 V _ {3} - V _ {4} - V _ {5} + 2 V _ {6}\right) \tag {12}
$$
$$
V _ {1} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \Psi \left(\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + 2 \widehat {\xi} ^ {l}\right), \quad V _ {2} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \Psi \left(\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + 2 \tilde {\xi} ^ {l}\right)
$$
$$
V _ {3} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \Psi \left(\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + \widehat {\xi} ^ {l} + \tilde {\xi} ^ {l}\right) V _ {4} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \exp \left(\left(\xi^ {\mu_ {1}}\right) ^ {\top} \widehat {\xi} ^ {l}\right) \exp \left(\left(\xi^ {\mu_ {2}}\right) ^ {\top} \widehat {\xi} ^ {l}\right)
$$
$$
V _ {5} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \exp \left(\left(\xi^ {\mu_ {1}}\right) ^ {\top} \tilde {\xi} ^ {l}\right) \exp \left(\left(\xi^ {\mu_ {2}}\right) ^ {\top} \tilde {\xi} ^ {l}\right) V _ {6} = \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \exp \left(\left(\xi^ {\mu_ {1}}\right) ^ {\top} \widehat {\xi} ^ {l}\right) \exp \left(\left(\xi^ {\mu_ {2}}\right) ^ {\top} \tilde {\xi} ^ {l}\right) \tag {13}
$$
for  $\tilde{\xi}^l$  denoting  $\hat{\xi}^l$  with one of its dimensions flipped and:
$$
\Psi (\mathbf {x}) \stackrel {\text {d e f}} {=} D ^ {4} \exp (- 2 N) (1 + 8 \widehat {A}) ^ {- \frac {N}{2}} \exp \left(\frac {B ^ {2}}{2 (1 - 8 \widehat {A})} \| \mathbf {x} \| ^ {2}\right) \tag {14}
$$
Proof. Take a memory  $\xi^l \in \{-1, +1\}^N$  and an input  $\widehat{\xi}^l \in \mathcal{B}(\xi^l, \rho N)$ . Denote by  $\mathrm{neg}(\widehat{\xi}^l, i)$  a vector obtained from  $\widehat{\xi}^l$  by replacing  $\widehat{\xi}^l(i)$  with  $-\widehat{\xi}^l(i)$ . Let us study the change of the energy of the system as we flip the value of the  $i$ th dimension of the input  $\widehat{\xi}^l$  since the sign of this change solely determines the update that will be made. We have the following:
$$
\Delta \left(E _ {\text {r a n d}}\right) = E \left(\operatorname {n e g} \left(\widehat {\xi} ^ {l}, i\right); \xi^ {1}, \dots , \xi^ {M}\right) - E \left(\widehat {\xi} ^ {i}; \xi^ {1}, \dots , \xi^ {M}\right) = E _ {\text {s i g n a l}} + E _ {\text {n o i s e}}, \tag {15}
$$
$$
E _ {\text {s i g n a l}} = \frac {1}{r} \sum_ {k = 1} ^ {r} \left(W _ {k} ^ {l} - Z _ {k} ^ {l}\right), \tag {16}
$$
$$
E _ {\text {n o i s e}} = \frac {1}{r} \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \} \backslash \{l \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right), \tag {17}
$$
and furthermore:  $W_{k}^{i} = a_{k}^{i}b_{k}, Z_{k}^{i} = a_{k}^{i}c_{k}$  for:
$$
a _ {k} ^ {i} = D \exp (- \frac {N}{2}) \exp (B \omega_ {k} ^ {\top} \xi^ {i} - \widehat {A} \| \omega_ {k} \| _ {2} ^ {2}),
$$
$$
b _ {k} = D \exp \left(- \frac {N}{2}\right) \exp \left(B \omega_ {k} ^ {\top} \hat {\xi} ^ {l} - \hat {A} \| \omega_ {k} \| _ {2} ^ {2}\right), \tag {18}
$$
$$
c _ {k} = D \exp (- \frac {N}{2}) \exp (B \omega_ {k} ^ {\top} \mathrm {n e g} (\widehat {\xi} ^ {l}, i) - \widehat {A} \| \omega_ {k} \| _ {2} ^ {2}).
$$
If  $\omega_{1},\ldots ,\omega_{r}\sim \mathcal{N}(0,\mathbf{I}_{N})$  then, from the fact that  $E_{\mathrm{rand}}$  is the unbiased estimation of  $E_{\mathrm{reg}}$ , we get:
$$
\mathbb {E} \left[ X _ {k} \right] = \exp \left(\left(\xi^ {l}\right) ^ {\top} \widehat {\xi} ^ {l}\right),
$$
$$
\mathbb {E} \left[ Y _ {k} \right] = \exp \left(\left(\xi^ {l}\right) ^ {\top} \operatorname {n e g} \left(\widehat {\xi} ^ {l}, i\right)\right), \tag {19}
$$
$$
\mathbb {E} \left[ W _ {k} ^ {\mu} \right] = \exp \left(\left(\xi^ {\mu}\right) ^ {\top} \widehat {\xi} ^ {l}\right),
$$
$$
\mathbb {E} \left[ Z _ {k} ^ {\mu} \right] = \exp \left(\left(\xi^ {\mu}\right) ^ {\top} \operatorname {n e g} (\widehat {\xi} ^ {l}, i)\right),
$$
This is a direct consequence of the OPRF-mechanism introduced in [44]. Variables:  $X_{k}$ ,  $Y_{k}$ ,  $W_{k}^{\mu}$  and  $Z_{k}^{\mu}$  for  $\mu = 1, \dots, M$  are simply unbiased estimators of the softmax-kernel values obtained via applying OPRF-mechanism. Let us now compute the expected change of the energy of the system:
$$
\mathbb {E} [ \Delta (E _ {\text {r a n d}}) ] = \mathbb {E} [ E _ {\text {s i g n a l}} ] + \mathbb {E} [ E _ {\text {n o i s e}} ], \tag {20}
$$
$$
\mathbb {E} \left[ E _ {\text {s i g n a l}} \right] = \frac {1}{r} \sum_ {k = 1} ^ {r} \left(\mathbb {E} \left[ X _ {k} \right] - \mathbb {E} \left[ Y _ {k} \right]\right) = \frac {1}{r} \sum_ {k = 1} ^ {r} \left(\exp \left(\left(\xi^ {l}\right) ^ {\top} \hat {\xi} ^ {l}\right) - \exp \left(\left(\xi^ {l}\right) ^ {\top} \operatorname {n e g} \left(\hat {\xi} ^ {l}, i\right)\right)\right) \tag {21}
$$
$$
\mathbb {E} [ E _ {\text {n o i s e}} ] = \frac {1}{r} \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \} \backslash \{l \}} \left(\mathbb {E} \left[ W _ {k} ^ {\mu} \right] - \mathbb {E} \left[ Z _ {k} ^ {\mu} \right]\right) = \tag {22}
$$
$$
\frac {1}{r} \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \} \backslash \{l \}} \left(\exp ((\xi^ {\mu}) ^ {\top} \widehat {\xi} ^ {l}) - \exp ((\xi^ {\mu}) ^ {\top} \operatorname {n e g} (\widehat {\xi} ^ {l}, i))\right)
$$
We will first upper bound  $|\mathbb{E}[E_{\mathrm{noise}}]|$ . We have:
$$
\begin{array}{l} \left| \mathbb {E} \left[ E _ {\text {n o i s e}} \right] \right| \leq \frac {1}{r} \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \} \backslash \{l \}} \left(\exp \left(\left(\xi^ {\mu}\right) ^ {\top} \widehat {\xi} ^ {l}\right) + \exp \left(\left(\xi^ {\mu}\right) ^ {\top} \operatorname {n e g} \left(\widehat {\xi} ^ {l}, i\right)\right)\right) \\ \leq \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \} \backslash \{l \}} \left(\exp \left(N (1 - 2 (\tau - \rho))\right) + \exp \left(N (1 - 2 (\tau - \rho) + \frac {2}{N})\right)\right) \tag {23} \\ \leq 2 M \exp (N (1 - 2 (\tau - \rho) + \frac {2}{N})) \\ \end{array}
$$
We will now consider two cases:
Case 1:  $\widehat{\xi}^l (i) = \xi^l (i)$ :
In this setting, flipping the value of the  $i$ th dimension of the input vector increases its distance from the close neighbor. Therefore in this case we would like the energy change of the system to be positive (so that the flip does not occur). From the Equation 21, we obtain:
$$
\begin{array}{l} \mathbb {E} \left[ E _ {\text {s i g n a l}} \right] \geq \frac {1}{r} \sum_ {k = 1} ^ {r} \left(\exp (N (1 - 2 \rho)) - \exp (N (1 - 2 \rho) - 2))\right) = \tag {24} \\ \exp (N (1 - 2 \rho)) (1 - e ^ {- 2}) \\ \end{array}
$$
$$
\mathbb {E} \left[ \Delta \left(E _ {\text {r a n d}}\right) \right] \geq \exp (N (1 - 2 \rho)) \left(1 - e ^ {- 2}\right) - 2 M \exp \left(N \left(1 - 2 (\tau - \rho) + \frac {2}{N}\right)\right) \tag {25}
$$
Therefore, if the following holds:
$$
M \leq \exp (2 N (\tau - 2 \rho)) \frac {1 - e ^ {- 2}}{2 e ^ {2}}, \tag {26}
$$
then  $\mathbb{E}[\Delta (E_{\mathrm{rand}})] > 0$
Case 2:  $\widehat{\xi}^l (i) = -\xi^l (i)$ :
In this setting, flipping the value of the  $i$ th 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:
$$
\begin{array}{l} \mathbb {E} \left[ E _ {\text {s i g n a l}} \right] \leq \frac {1}{r} \sum_ {k = 1} ^ {r} \left(\exp (N (1 - 2 \rho)) - \exp (N (1 - 2 \rho) + 2))\right) = \tag {27} \\ \exp (N (1 - 2 \rho)) (1 - e ^ {2}) \\ \end{array}
$$
$$
\mathbb {E} \left[ \Delta \left(E _ {\text {r a n d}}\right) \right] \leq \exp \left(N \left(1 - 2 \rho\right)\right) \left(1 - e ^ {2}\right) + 2 M \exp \left(N \left(1 - 2 (\tau - \rho) + \frac {2}{N}\right)\right) \tag {28}
$$
$$
M \leq \exp (2 N (\tau - 2 \rho)) \frac {e ^ {2} - 1}{2 e ^ {2}}, \tag {29}
$$
then  $\mathbb{E}[\Delta (E_{\mathrm{rand}})] < 0$ . Note that the bound from Inequality 26 is stronger than the one from Inequality 29. That completes the proof of the first part of the theorem.
Now we will compute the variance of  $\Delta(E_{\mathrm{rand}})$ . Denote:
$$
Z _ {k} = \sum_ {\mu \in \{1, \dots , M \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right) \tag {30}
$$
Note that if  $\omega_{1},\ldots,\omega_{r}$  are chosen independently then  $Z_{k}$  for  $k = 1,\ldots,r$  are independent. The following is true:
$$
\begin{array}{l} \operatorname {V a r} \left(\Delta \left(E _ {\text {r a n d}}\right)\right) = \operatorname {V a r} \left(E _ {\text {s i g n a l}} + E _ {\text {n o i s e}}\right) = \operatorname {V a r} \left(\frac {1}{r} \sum_ {k = 1} ^ {r} \sum_ {\mu \in \{1, \dots , M \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right)\right) \\ = \operatorname {V a r} \left(\frac {1}{r} \sum_ {k = 1} ^ {r} Z _ {k}\right) = \frac {1}{r ^ {2}} \sum_ {k = 1} ^ {r} \operatorname {V a r} \left(Z _ {k}\right) = \frac {1}{r ^ {2}} \sum_ {k = 1} ^ {r} \operatorname {V a r} \left(\sum_ {\mu \in \{1, \dots , M \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right)\right) \tag {31} \\ = \frac {1}{r ^ {2}} \sum_ {k = 1} ^ {r} \left(\mathbb {E} \left[ \left(\sum_ {\mu \in \{1, \dots , M \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right)\right) ^ {2} \right] - \left(\mathbb {E} \left[ \sum_ {\mu \in \{1, \dots , M \}} \left(W _ {k} ^ {\mu} - Z _ {k} ^ {\mu}\right) \right]\right) ^ {2}\right) \\ \end{array}
$$
Therefore we have:
$$
\begin{array}{l} \operatorname {V a r} \left(\Delta \left(E _ {\text {r a n d}}\right)\right) = \frac {1}{r ^ {2}} \sum_ {k = 1} ^ {r} \left(\sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} \left[ W _ {k} ^ {\mu_ {1}} W _ {k} ^ {\mu_ {2}} \right] + \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} \left[ Z _ {k} ^ {\mu_ {1}} Z _ {k} ^ {\mu_ {2}} \right]\right) \\ - \frac {2}{r ^ {2}} \sum_ {k = 1} ^ {r} \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} \left[ W _ {k} ^ {\mu_ {1}} Z _ {k} ^ {\mu_ {2}} \right] \tag {32} \\ - \frac {1}{r ^ {2}} \sum_ {k = 1} ^ {r} \left(\sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} \left[ W _ {k} ^ {\mu_ {1}} \right] \mathbb {E} \left[ W _ {k} ^ {\mu_ {2}} \right] + \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} \left[ Z _ {k} ^ {\mu_ {1}} \right] \mathbb {E} \left[ Z _ {k} ^ {\mu_ {2}} \right]\right) \\ - \frac {2}{r ^ {2}} \sum_ {k = 1} ^ {r} \sum_ {\mu_ {1}, \mu_ {2} \in \{1, \dots , M \}} \mathbb {E} [ W _ {k} ^ {\mu_ {1}} ] \mathbb {E} [ Z _ {k} ^ {\mu_ {2}} ] \\ \end{array}
$$
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 convenience):
$$
\mathbb {E} \left[ W _ {k} ^ {\mu} \right] = \exp \left(\left(\xi^ {\mu}\right) ^ {\top} \widehat {\xi} ^ {\iota}\right), \tag {33}
$$
Let us now define:
$$
\Psi (\mathbf {x}) = D ^ {4} \exp (- 2 N) \exp (B \omega^ {\top} \mathbf {x} - 4 \widehat {A} \| \omega \| _ {2} ^ {2}). \tag {34}
$$
Note that the following is true:
$$
\mathbb {E} \left[ W _ {k} ^ {\mu_ {1}} W _ {k} ^ {\mu_ {2}} \right] = \Psi \left(\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + 2 \widehat {\xi} ^ {l}\right)
$$
$$
\mathbb {E} \left[ Z _ {k} ^ {\mu_ {1}} Z _ {k} ^ {\mu_ {2}} \right] = \Psi \left(\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + 2 \mathrm {n e g} \left(\widehat {\xi} ^ {l}, i\right)\right) \tag {35}
$$
$$
\mathbb {E} \left[ W _ {k} ^ {\mu_ {1}} Z _ {k} ^ {\mu_ {2}} \right] = \Psi \big (\xi^ {\mu_ {1}} + \xi^ {\mu_ {2}} + \widehat {\xi} ^ {l} + \mathrm {n e g} (\widehat {\xi} ^ {l}, i) \big)
$$
Thus it remains to find closed-form formula for  $\Psi (\mathbf{x})$  for any given  $\mathbf{x}\in \mathbb{R}^N$ .
From the proof of Theorem 3.1 in [44], we get for  $A < 0$ :
$$
\mathbb {E} \left[ \exp \left(A \| \omega \| ^ {2} + B \omega^ {\top} \mathbf {x}\right) \right] = (1 - 2 A) ^ {- \frac {N}{2}} \exp \left(\frac {B ^ {2}}{2 (1 - 2 A)} \| \mathbf {x} \| ^ {2}\right) \tag {36}
$$
$$
\Psi (\mathbf {x}) = D ^ {4} \exp (- 2 N) (1 + 8 \widehat {A}) ^ {- \frac {N}{2}} \exp \left(\frac {B ^ {2}}{2 (1 - 8 \widehat {A})} \| \mathbf {x} \| ^ {2}\right) \tag {37}
$$
Plugging to Equation 32 formulae from Equation 33 and Equation 35 and utilizing Equation 37 for  $\Psi$ , we obtain the formula for the variance from the statement of the Theorem.
B.1 Warm-up for Mnemosyne and other optimizers: additional results
Preliminaries: At each timestep  $t$ , gradient  $\nabla f(\mathbf{x}_t)$  is input to the optimizer. The gradient is preprocessed as proposed in [2]. Coordinate-wise Mnemosyne's using two temporal encoders is applied. The Mnemosyne's memory cell interfaces with the rest of the system similarly to any RNN-cell. Each cell uses exponential discount factor  $\tau = 0.1$ ,  $r = 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 optimizeee.
Meta-training: We refer to training optimizer's parameters  $\theta$  as meta-training to distinguish from the optimizee NN training. Mnemosyne's optimizer is meta-trained on MNIST classification task with 3
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: (1 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 3 small ViT models. The optimizee MLPs are sampled from this hyperparameter distribution:  $l \in [1,2]$  hidden layers of size in range [20, 40] and sigmoid or relu activation function. The optimizee ViTs have  $l \in [1,3]$  layers,  $h \in [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 to 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-trained by gradient descent using Adam optimizer with learning rate  $\eta = 3e^{-4}$  to minimize a combination of two loss functions. The first is the task loss given by the sum of optimizee losses in a 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 trajectory, often observed in behaviour cloning. Our imitation loss acts as a regularizer which prevents Mnemosyne's optimizer from over-fitting on the optimizee task that it is trained on. We emphasize that expert's learning rate  $\eta_{\mathrm{exp}} = 3e^{-2}$  was not obtained via any tuning process.
Our optimizer model has minimal input feature engineering and our meta-training setup is significantly simpler than those considered in the literature [49, 13, 47, 71]. Even so, we can successfully apply Mnemosyne's optimizer to a variety of tasks due to its efficient memory mechanism. Furthermore, Mnemosyne'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-training. Fig. 10 shows that Mnemosyne can optimize MLPs with different NN architectures 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 [2] struggles to minimize classification loss beyond a few hundred steps.
Transformers: 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. 10, here Mnemosyne's optimizer is faster than standard analytical optimizers and much more stable than LSTM optimizer. Fig. 11 shows the benefit of using expert imitation-loss for long-horizon stability of the Mnemosyne's optimizer.
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. It 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-training). Furthermore, for larger 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 if meta-trained on much simpler tasks as we show in Fig. 3.
Figure 11: Impact of training the optimizer with combined meta loss and imitation loss can be seen in generalization to a long horizon rollout. All variants were trained only on length 100 rollouts.
Figure 12: Ablation Studies. Left: Comparison of the Mnemosyne's linear CAM with regular attention memory blocks with different history cache lengths  $(h)$ . 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.
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, we could not implement a regular attention based optimizer with cache length more than 100. On the other hand, Performer's memory cell defining CAM can attend to theoretically unbounded history and out-performs regular attention variants with fixed memory requirement.
B.3 Different RF-mechanisms: detailed look
Fig. 12 (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 use 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. 12 (right)) and thus it became our default choice.
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.
Figure 13: Comparison of the coordinate-wise Mnemosyne with ReLU kernel, square kernel and FAVOR+ approximation of the softmax kernel. We present the meta-training loss of the three variants. FAVOR+ mechanism outperforms other variants.
B.5 Ablations over different discount factors and number of RFs in CAM mechanism
In Fig. 14, we present detailed ablation studies over discount factors  $\tau$  as well as the number of random features applied by our default CAM mechanism leveraging hyperbolic cosine random features.
B.6 Benchmarking different depths of the temporal module
Finally, we run ablations over different number of temporal encoders in Mnemosyne's temporal block. We noticed that modest increase of the number of encoders improves loss in meta-training and meta-training very deep variants is particularly challenging (as requiring much more data). Since in this paper we decided to use simple meta-training strategies and furthermore increasing the number of temporal encoders did not lead to substantial gains, we decided to choose shallow temporal encoders' architectures. The results are presented in Fig. 15.
B.7 Compute Resources Used
All Mnemosyne optimizer variants were trained and tested on a 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.
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.
Figure 15: Comparison of the meta-training loss for Mnemosyne variants applying different number of temporal encoders  $k$ . Since several variants on the left figure performs similarly, on the right figure we highlight top two. The meta training is conducted on the MLP optimization tasks and MNIST data.
Figure 16: Coordinate-wise Mnemosyne across different ViT architectures and datasets, as described in Sec. B.8.1. Mnemosyne matches or outperforms optimal Adam variants without any hyperparameter tuning.
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: 2. We used the coordinate-wise variant of the Mnemosyne. We run tests on the following datasets: imagenet2012, places365 and caltech-birds-2011. We were optimizing the last layer of the ViT-architecture and used Adam expert with learning rate  $\eta = 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 of 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  $\eta = e^{-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 a hand-designed optimizer is. The results are presented in Fig. 16 and Fig. 17. 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.
B.8.2 ViT multi-layer fine-tuning
Here, we used a 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  $= 1e^{-3}$ ). For comparison, the same baseline Adam variant is to fine-tune the complete model.
Table 2: Hyperparameters for the different ViT models used in this paper