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

Abstract

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 $J_1$-$J_2$ 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.

arXiv comment10 pages, 5 figures, 1 table

References in corpus (34)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Solving the Quantum Many-Body Problem with Artificial Neural Networks
• Quantum Spin Liquids
• CoAtNet: Marrying Convolution and Attention for All Data Sizes
• Computational Studies of Quantum Spin Systems
• Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks
• Finite-Size Scaling of the Ground State Parameters of the Two-Dimensional Heisenberg Model
• Plaquette Ordered Phase and Quantum Phase Diagram in the Spin-1/2 J1-J2 Square Heisenberg Model
• Green Function Monte Carlo with Stochastic Reconfiguration
• Wave function optimization in the variational Monte Carlo method
• Neural-Network Quantum States, String-Bond States, and Chiral Topological States
• Magnetic order and disorder in the frustrated quantum Heisenberg antiferromagnet in two dimensions
• Direct evidence for a gapless $Z_2$ spin liquid by frustrating Néel antiferromagnetism
• Observation of a 4-spin Plaquette Singlet State in the Shastry-Sutherland compound SrCu2(BO3)2
• NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems
• Dirac-type nodal spin liquid revealed by refined quantum many-body solver using neural-network wave function, correlation ratio, and level spectroscopy
• On the origin of long-range correlations in texts
• Fermionic Wave Functions from Neural-Network Constrained Hidden States
• Transformer variational wave functions for frustrated quantum spin systems
• Empowering deep neural quantum states through efficient optimization
• Rethinking Positional Encoding in Language Pre-training
• Helping restricted Boltzmann machines with quantum-state representation by restoring symmetry
• High-accuracy variational Monte Carlo for frustrated magnets with deep neural networks
• A simple linear algebra identity to optimize Large-Scale Neural Network Quantum States
• Gauge Invariant and Anyonic Symmetric Transformer and RNN Quantum States for Quantum Lattice Models
• Optimizing Design Choices for Neural Quantum States
• Variational Monte Carlo with Large Patched Transformers
• Mapping of attention mechanisms to a generalized Potts model
• A Self-Attention Ansatz for Ab-initio Quantum Chemistry
• Ab-initio variational wave functions for the time-dependent many-electron Schrödinger equation
• Fisher Information and Natural Gradient Learning of Random Deep Networks
• The NLP Task Effectiveness of Long-Range Transformers
• Fine-tuning Neural Network Quantum States
• Vision Transformers provably learn spatial structure

Cited by in corpus (3)

• Transformer Wave Function for two dimensional frustrated magnets: emergence of a Spin-Liquid Phase in the Shastry-Sutherland Model
• Design principles of deep translationally-symmetric neural quantum states for frustrated magnets
• Comparing Symmetrized Determinant Neural Quantum States for the Hubbard Model