SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 (we explicitly put the gains coming from the two variants of Mnemosyne as compared to the expert variant). This shows that Mnemosyne does not learn to imitate the expert.
Model	Heads	Layers	Hidden Dim.	MLP Dim.	Params	Patch Size
ViT-Base	12	12	768	3072	86M	16
ViT-Large (16)	24	16	1024	4096	307M	16
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 andcls layer of ViT-H (see Tab: 2 for hyperparameter) using tensorwise  $(\sim 1M$  params), while the head was trained using Adam. The rest of the transformer parameters are fixed to the pre-trained value for all methods. The batch size was set at 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  $\sim 50M$  parameters with a 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 BERT-pretraining 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: 4. We trained all parameters from scratch for all methods, with a batch size of 512. For the Mnemosyne results shown in Fig: 7, we trained all parameters except the token embedding using Tensorwise Mnemosyne ( $\sim$ 86M parameters). The token embedding was trained using Adam with learning rate  $1e - 4$ . 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
We use coordinate-wise Mnemosyne to prompt-tune [38] T5XXL model [55] (see Table 5 for hyper-parameters) on SuperGLUE 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.
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 Layers	Decoder Layers	Heads	Head Dim.	Embedding Dim.	MLP Dim.	Params	Compute
T5XXL	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 of using  $g_{\theta}$  is more general than the one presented in Eq. 1 and is of the form given below for  $g_{\theta} = (g_{\theta_1}^{\mathrm{init}}, g_{\theta_2}^{\mathrm{hist}})$ ,  $\theta = [\theta_1; \theta_2]$  and a context-vector  $\mathbf{c}$ :
$$
\left\{ \begin{array}{l} \mathbf {x} _ {0} = g _ {\theta_ {1}} ^ {\text {i n i t}} (\mathbf {c}), \\ \mathbf {x} _ {t + 1} = g _ {\theta_ {2}} ^ {\text {h i s t}} (f, \mathbf {x} _ {0}, \dots , \mathbf {x} _ {t}) \quad \text {i f} t > 0 \end{array} \right. \tag {38}
$$
The optimizer  $g$  is now explicitly split into two-parts: (1)  $g_{\theta_1}^{\mathrm{init}}$  that learns initial optimization point from the context-vector, e.g. the image encoding the scene (as it is the case in learnable-MPC setting [72]), and (2)  $g_{\theta_2}^{\mathrm{hist}}$  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  $g_{\theta_1}^{\mathrm{init}}$  applies spatial bi-directional attention while  $g_{\theta_2}^{\mathrm{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
Preliminaries: 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 as iterative Linear Quadratic Regulator (iLQR) [41] is one way to implement MPC. However, in challenging layouts with narrow corridors and doors and with conservative safety and collision avoidance constraints, iLQR struggles to converge to the optimal trajectory. Sequential Quadratic Programming (SQP) [62] is often a more robust optimizer that can handle difficult non-linear constraints. However, SQP is significantly, sometimes  $\sim 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 [30]. We learn Mnemosyne's optimizer to imitate SQP
Figure 18: Navigation in a photo-realistic environment with MPC policy using Mnemosyne optimizer for initializing SQP solver.
behaviour and initialize it with an approximately optimal trajectory as a starting point from which SQP refines to an optimized and feasible trajectory.
Training details: Mnemosyne's optimizer,  $g_{\theta_1}^{init}$ , receives current robot pose  $p_r$ , a goal pose  $p_g$  and a visual occupancy grid as the context, c. The occupancy grid is processed by an image encoder module, a ViT where the attention mechanism is approximated by bidirectional Mnemosyne memory. As in a ViT, the occupancy grid is first pre-processed by a convolution layer and then flattened to a sequence. Each element (token) of the sequence corresponds to a different  $5 \times 5$  patch of the original frame which is then enriched with positional encoding. The pre-processed input is then fed to 3 Mnemosyne attention and MLP layers of hidden dimension 64. The final embedding of one of the tokens is chosen as a latent representation of the occupancy grid,  $l_{oc}$ .  $p_r$ ,  $p_g$  and  $l_{oc}$  are concatenated and processed by an MLP which outputs the predicted action trajectory,  $\mathbf{x}_0$ .
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 a pair of input context  $\mathbf{c}$  and the final optimal trajectory  $\mathbf{x}_T$  was recorded. A total of  $\sim 500$ , 000 training examples were collected. Mnemosyne's optimizer was trained with supervised learning on the SQP dataset by minimizing mean squared error between the SQP optimal trajectories and the predicted trajectories.
Results: After training, the predicted trajectory from Mnemosyne was used to initialize SQP optimization. Without Mnemosyne initialization, SQP optimization was capped at maximum 10 iterations. It took on average 4.78 iterations and 0.12sec for the SQP solution to complete. With Mnemosyne initialization, SQP is only run for 1 iteration to reach the optimal trajectory. SQP generates a trajectory that satisfies kinematic, dynamic and safety constraints for the robot which transformer alone e.g. [7] cannot natively guarantee. This reduces the optimization time by more than half to 0.048sec on average which is under real-time constraint. It includes 0.011sec for Mnemosyne inference and the rest for SQP iteration. A sequence of snapshots during navigation with Mnemosyne-SQP optimizer in a sample environment is shown in Fig. 18.