# Sampling at intermediate temperatures is optimal for training large language models in protein structure prediction

arXiv: https://arxiv.org/abs/2603.29529 · HTML: https://mlsanity.com/paper/2603.29529/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 35 references that are in the map.
- First posted on arXiv 31 Mar 2026.

## Builds on (most cited first, 12 of 12)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [BERTology Meets Biology: Interpreting Attention in Protein Language Models](https://mlsanity.com/paper/2006.15222/) (2020)
- [Unreasonable Effectiveness of Learning Neural Networks: From Accessible States and Robust Ensembles to Basic Algorithmic Schemes](https://mlsanity.com/paper/1605.06444/) (2016)
- [Subdominant Dense Clusters Allow for Simple Learning and High Computational Performance in Neural Networks with Discrete Synapses](https://mlsanity.com/paper/1509.05753/) (2015)
- [Memorizing without overfitting: Bias, variance, and interpolation in over-parameterized models](https://mlsanity.com/paper/2010.13933/) (2020)
- [Mapping of attention mechanisms to a generalized Potts model](https://mlsanity.com/paper/2304.07235/) (2023)
- [Geometry of energy landscapes and the optimizability of deep neural networks](https://mlsanity.com/paper/1808.00408/) (2018)
- [Controlled Langevin Dynamics for Sampling of Feedforward Neural Networks Trained with Minibatches](https://mlsanity.com/paper/2603.15367/) (2026)

## Abstract (arXiv)

Using a statistical mechanics framework, we investigate the parameter space of transformer models trained on protein sequence data. We sample the loss landscape at varying temperatures using Langevin dynamics to characterize the low-loss manifold, and to understand the mechanisms underlying transformers' superior performance in protein structure prediction. We find that, at variance with networks not based on the attention mechanism, the lack of a first--order--like transition in the loss of the transformer produces a range of intermediate temperatures with good learning properties; this is true both for synthetic and natural protein sequences. We also show that the parameters of most layers are highly conserved at these temperatures if the dimension of the embedding is optimal, and we provide an operative way to find this dimension. Additionally, we show that the attention matrix is more predictive of the contact maps of the protein at higher temperatures and for higher dimensions of the embedding than those optimal for learning. Finally, we showed that the models sampled at intermediate temperatures can predict the free-energy variation upon mutation, better than models obtained through standard optimization techniques.