Mnemosyne: Learning to Train Transformers with Transformers

In this work, we propose a new class of learnable optimizers, called \textit{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.

Paper

References (85)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer JvEY6/10 · confidence 4/52023-07-07

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.

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. This 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. There are also some theoretical results related to associate memory. The sheer amount of work seems sufficient for a neurips acceptance.

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 Also, 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.

Questions

- 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 environment we need to run the proposed solution, e.g., a few GPUs & a few days or 1000 A100s for a month?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA

Reviewer xFuP6/10 · confidence 2/52023-07-14

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.

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. (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. (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.

Questions

See the weakness.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have discussed the limitations.

Reviewer V1vR5/10 · confidence 4/52023-07-31

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.

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. The authors provide a theoretical analysis of the compact associative memory used, but the importance and implications of this result is not discussed. While 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.

Questions

I would like to ask the authors the following questions: 1/ 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? 2/ 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? 3/ Can the authors articulate the importance of the theoretical analysis? 4/ Why did you consider the VeLo architecture in the experiments? Did you consider other options? 5/ 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?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

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.

Authorsrebuttal2023-08-18

A note to Reviewer V1vR

We would like to once more thank the Reviewer V1vR for the comments. We have addressed in detail all of them. Furthermore, in our rebuttal we have added additional empirical results showing strong performance of our method. We would like to sincerely ask the Reviewer to read our response and consider changing the score. Thank you very much. The Authors

Reviewer xFuP2023-08-16

Response to the authors

I am not an expert in the relevant field, but after reading the author's response and considering the opinions of other reviewers, I decided to keep the original score.

Reviewer V1vR2023-08-21

re: Rebuttal by Authors

I thank the authors for their detailed response. A couple of follow-up comments. 1/ I do not dispute the contributions mentioned above. However, these contributions seem to have been mainly driven by practical (to make sure the method could scale to at least modest size architectures), rather than methodological insights. if this is not the case I would love to hear the authors elaborate on these insights. 2/ It is great to hear the authors plan to make the code available. Arguably, I did not check the more detailed description of the implementation in the Appendix. I would encourage the authors to refer to the relevant Appendix in the main body of the paper. Regarding the robustness of the meta-training: (i) can the authors elaborate on what they mean by "Mnemosyne is robust since our meta-training is simple by design" and (ii) how can they claim this procedure is robust if they did not investigate the impact of the hyperparameters on it? 3/ I had missed the fact that authors had used entire Bert-base in the NLP experiments. My apologies. This being said, my point was that the method does not seem to be applicable to more recent neural nets (even though the authors put efforts and introduced novelties to scale it better). This suggest there are inherent practical limitations to the proposed method. Also, I do not dispute the fact that recent work focusses on parameter-efficient fine-tuning. However, i believe the proposed method is orthogonal to these works.

Reviewer V1vR2023-08-21

re: Rebuttal by Authors (continued)

(continuation from previous post; comments here follow the general response to authors.) 4/ Re Implications of theoretical analysis on design choices: Thanks for the additional clarifications. i would encourage the authors to ensure these details are included in the main body of the paper. What remains unclear is how the theoretical analysis motivates FAVOR++. Can you please elaborate? 5/ Re Mnemosyne can perform stable long-horizon optimisation: While appreciate the additional discussion provided by the authors, it does not answer my original questions which were (i) can you explain why long-term dependencies could be important when learning an optimiser and (ii) can you provide direct evidence that Mnemosyne is able to capture? Fig 10 and 11 show strong optimisation performance (which is great!), but it does not tell me that this strong performance is due to an ability to capture long-term dependencies.

Authorsrebuttal2023-08-21

Thank you very much for the additional comments. We address them in detail below. 1. The whole field of learnable optimizers is definitely practically-motivated, since it aims to address the practical limitations of the brute-force counterparts. In that sense, our work is grounded in the practical considerations (e.g. scalability of the LOs; critical for the field). We do not consider this to be a weakness of the paper. However, as explained in detail in the rebuttal, we have developed several theoretical techniques that are interesting on their own, in particular: * the autoregressive linear attention mechanism with smoothly decaying attention weights, * scalable spatio-temporal Transformer model with a novel pooling algorithm, * the autoregressive variant of the FAVOR++ mechanism (that prior to that paper was given only in the bi-directional form). To the best of our knowledge, this is one of the few papers in the field with a rich set of techniques and the rigorous corresponding theoretical analysis. We would like to kindly ask the Reviewer what methodological insights, orthogonal to any practical considerations the Reviewer would like to see in the paper. 2. Thank you very much for the comment ! (i) We would like to point out that one of the strategies in the L2L literature for improving performance is to scale-up meta-training to a large model close to the one seen in evaluation. This increases the tendency of LO to specialize on one task rather than be robustly applicable. On the contrary, we show that it is possible to restrict the meta-training budget and obtain LOs that can generalize to the training of models far larger than those seen in the meta-training. This highlights that the gradient update rules learned by our LO are generally useful rather than memorizing the loss landscape of the model to be evaluated on. (ii) Unlike standard optimizers, Mnemosyne LO does not require any hyper-parameters (such as learning rate etc.) by design. In the paper, we have shown ablation studies for the design choices in the Mnemosyne architecture such as: (1) the decay factor (Fig 13. left subfigure), (2) number of random features (Fig 13. right subfigure), (3) number of temporal encoders (Fig.14), (4) different linear attention mechanisms approximating softmax attention (Fig. 12 middle and right subfigures) and (5) different linear attention kernels beyond softmax (in the rebuttal pdf). These choices are made before meta-training and based on the comparison studies. After meta-training, the Mnemosyne LO is applied without any downstream hyper-parameter tuning on the evaluation tasks. 3. Thank you very much for your comment. We do acknowledge the limitations of the method in terms of scaling up to fine-tuning all the parameters of the 1B+ size models. However we would like to emphasize that this was never a target of this manuscript and we leave it to future work. The scalability to massive models is a challenge for all LOs. We do believe that this work is a significant step towards addressing it, as capable of fine-tuning much larger sets of parameters than previously proposed techniques. Given the above, parameter-efficient fine-tuning scenario is one of the most natural applications of the LOs and that is why we devote a considerable amount of time to discuss it in the paper. We want to emphasize that the current applicability of the method is not related to how “recent” the neural network is, but to: how many parameters should be directly fine-tuned. All the architectures we applied our method to are currently used by the Community.

Authorsrebuttal2023-08-21

4. Thank you for the comment. We will make sure to include all the relevant details in the final version of the paper. We have provided an explanation of the relevance of the presented theoretical analysis in the rebuttal and put it also here for Reviewer’s convenience. Our theoretical results show that the mechanism is able to memorize long histories of the optimization trajectories, since its corresponding associative memory model has on average an exponential-size capacity (Theorem 4.3). Furthermore, our detailed choices for the particular linear attention variant approximating softmax-kernel are driven by the variance analysis of the corresponding estimators. In particular, our default hyperbolic cosine variant is characterized by the provably reduced variance and the autoregressive FAVOR++ mechanism provides additional accuracy gains. This theoretical analysis is perfectly aligned with the practical observations regarding the downstream performance of these methods (middle and right sub-figure of Fig. 12). 5. Thank you for a good question. (i) In standard optimizer design, we observe the benefit from including time-series features that capture historical dependencies, such as first and second-order momentum in the case of ADAM. The LO field therefore focuses on memory based models such as RNNs for this task to enable the model to capture history features. It is well understood that attention mechanisms outperform RNNs in a wide range of sequence modeling tasks (the biggest example being language modeling). With this paper, we bring the same intuition to the field of LO. (ii) Fig. 12 (left) shows the performance benefit of including longer gradient histories in the regular attention based LO input. Note that regular attention mechanism is memory and compute intensive and can only process small history sizes. This motivates the use of linear auto-regressive attention proposed by us. The plot shows that Mnemosyne’s CAM mechanism (which can theoretically attend to unbounded history) outperforms all regular attention LOs with fixed size history input. This is evidence for the importance of long-horizon historical context in the optimization task.

Reviewer V1vR2023-08-22

Re: Official Comment by Authors

I would like to thank the authors for their detailed comments and clarifications. I’m happy to increase my score.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC