# Are queries and keys always relevant? A case study on Transformer wave functions

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 66 references that are in the map.
- First posted on arXiv 29 May 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [A Primer in BERTology: What we know about how BERT works](https://mlsanity.com/paper/2002.12327/) (2020)
- [CoAtNet: Marrying Convolution and Attention for All Data Sizes](https://mlsanity.com/paper/2106.04803/) (2021)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [Rethinking Positional Encoding in Language Pre-training](https://mlsanity.com/paper/2006.15595/) (2020)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [A Self-Attention Ansatz for Ab-initio Quantum Chemistry](https://mlsanity.com/paper/2211.13672/) (2022)
- [Neural-Network Quantum States, String-Bond States, and Chiral Topological States](https://mlsanity.com/paper/1710.04045/) (2017)
- [A phase transition between positional and semantic learning in a solvable model of dot-product attention](https://mlsanity.com/paper/2402.03902/) (2024)
- [Fisher Information and Natural Gradient Learning of Random Deep Networks](https://mlsanity.com/paper/1808.07172/) (2018)
- [NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems](https://mlsanity.com/paper/2112.10526/) (2021)
- [Do Transformers Need Deep Long-Range Memory](https://mlsanity.com/paper/2007.03356/) (2020)
- [The Case for Translation-Invariant Self-Attention in Transformer-Based Language Models](https://mlsanity.com/paper/2106.01950/) (2021)
- [The NLP Task Effectiveness of Long-Range Transformers](https://mlsanity.com/paper/2202.07856/) (2022)
- [Mapping of attention mechanisms to a generalized Potts model](https://mlsanity.com/paper/2304.07235/) (2023)
- [Gauge Invariant and Anyonic Symmetric Transformer and RNN Quantum States for Quantum Lattice Models](https://mlsanity.com/paper/2101.07243/) (2021)
- [Geometry of learning neural quantum states](https://mlsanity.com/paper/1910.11163/) (2019)
- [Ab-initio variational wave functions for the time-dependent many-electron Schr"odinger equation](https://mlsanity.com/paper/2403.07447/) (2024)

## Built on by (most cited first, 6 of 6)

- [Hyperbolic recurrent neural network as the first type of non-Euclidean neural quantum state ansatz](https://mlsanity.com/paper/2505.22083/) (2025)
- [Scalable Physics-Inspired Transformers for Spin Glasses](https://mlsanity.com/paper/2606.22984/) (2026)
- [Explorative Curriculum Learning for Strongly Correlated Electron Systems](https://mlsanity.com/paper/2505.00233/) (2025)
- [New non-Euclidean neural quantum states from hyperbolic Lorentz recurrent architectures](https://mlsanity.com/paper/2604.24337/) (2026)
- [Fermionic neural Gibbs states](https://mlsanity.com/paper/2512.04663/) (2025)
- [Generalization of Transformer-Based Neural Quantum States via In-Context Learning](https://mlsanity.com/paper/2610.03463/) (2026)

## Abstract (arXiv)

The dot product attention mechanism, originally designed for natural language processing tasks, is a cornerstone of modern Transformers. It adeptly captures semantic relationships between word pairs in sentences by computing a similarity overlap between queries and keys. In this work, we explore the suitability of Transformers, focusing on their attention mechanisms, in the specific domain of the parametrization of variational wave functions to approximate ground states of quantum many-body spin Hamiltonians. Specifically, we perform numerical simulations on the two-dimensional J1-J2 Heisenberg model, a common benchmark in the field of quantum many-body systems on lattice. By comparing the performance of standard attention mechanisms with a simplified version that excludes queries and keys, relying solely on positions, we achieve competitive results while reducing computational cost and parameter usage. Furthermore, through the analysis of the attention maps generated by standard attention mechanisms, we show that the attention weights become effectively input-independent at the end of the optimization. We support the numerical results with analytical calculations, providing physical insights of why queries and keys should be, in principle, omitted from the attention mechanism when studying large systems.