# Generalization of Transformer-Based Neural Quantum States via In-Context Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 40 of its 78 references that are in the map.
- First posted on arXiv 2 Oct 2026.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Multimodal Transformer for Unaligned Multimodal Language Sequences](https://mlsanity.com/paper/1906.00295/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Deep Interest Network for Click-Through Rate Prediction](https://mlsanity.com/paper/1706.06978/) (2017)
- [MetaICL: Learning to Learn In Context](https://mlsanity.com/paper/2110.15943/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Predicting Many Properties of a Quantum System from Very Few Measurements](https://mlsanity.com/paper/2002.08953/) (2020)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Behavior Sequence Transformer for E-commerce Recommendation in Alibaba](https://mlsanity.com/paper/1905.06874/) (2019)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [A Self-Attention Ansatz for Ab-initio Quantum Chemistry](https://mlsanity.com/paper/2211.13672/) (2022)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [A survey on the complexity of learning quantum states](https://mlsanity.com/paper/2305.20069/) (2023)
- [Ab-initio quantum chemistry with neural-network wavefunctions](https://mlsanity.com/paper/2208.12590/) (2022)
- [Attention-based Quantum Tomography](https://mlsanity.com/paper/2006.12469/) (2020)
- [Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond](https://mlsanity.com/paper/2407.10005/) (2024)
- [Learning hard quantum distributions with variational autoencoders](https://mlsanity.com/paper/1710.00725/) (2017)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)

## Abstract (arXiv)

Neural quantum states based on modern deep learning architectures have emerged as powerful representations for quantum many-body systems. In particular, Transformer-based neural quantum states provide expressive models capable of capturing long-range correlations, and their empirical generalization performance has recently been demonstrated. However, a theoretical understanding of their generalization behavior remains largely unexplored. In this paper, we develop a theoretical framework to analyze the generalization properties of Transformer-based neural quantum states under in-context learning. We establish a rigorous inference-time generalization error bound in terms of mean squared error (MSE), showing that the pointwise prediction error decreases inversely with both the number of in-context examples and the depth of the Transformer. We further show that the Transformer depth required to achieve this guarantee scales only linearly with the system size--namely, the number of particles in continuous systems or the number of qudits in discrete systems. Building on this result, we extend our analysis to full quantum states formulated as rank-one density operators, and derive MSE-based generalization bounds over both continuous and discrete domains under physical constraints. Finally, numerical simulations corroborate our theoretical analysis.