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

TL;DR

This paper investigates the relevance of queries and keys in Transformer attention mechanisms within quantum many-body physics, demonstrating that simplified attention can perform competitively and offers physical insights into their necessity.

Contribution

The study shows that omitting queries and keys from Transformer attention mechanisms can still yield effective results in quantum physics applications, reducing computational costs.

Findings

Simplified attention performs comparably to standard attention in quantum simulations.

Attention weights become input-independent at optimization end.

Analytical insights explain why queries and keys may be omitted in large systems.

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.