$U(1)$ Gauge-Equivariant Representation Learning of Bloch States

Year: 2026

Type: preprint

Abstract: Representation learning, or featurization, of mean-field Bloch states provides an important route for incorporating electronic information into machine-learning models of first-principles condensed-phase systems. A challenge in representing Bloch states is the gauge redundancy of each state. In this work, we propose to address this issue by explicitly incorporating gauge equivariance in the machine-learning framework. We develop a $U(1)$-equivariant autoencoder that compresses the plane-wave-basis Bloch states obtained from density functional theory into a low-dimensional latent representation. We achieve an average overlap of 0.970 between original wavefunctions and the reconstructions from a 64-dimensional latent space using a dataset of two-dimensional insulators. We further show that $U(1)$ equivariance endows the latent space with physically meaningful structures, including smoothness and topological information. Finally, we present a proof-of-principle application in which the latent representations are used to predict the off-diagonal elements of the $GW$ self energy matrix, and demonstrate the importance of enforcing gauge equivariance in the prediction.

Source: arXiv (Cornell University)

DOI: 10.48550/arxiv.2609.30845

Authors: Zhang, Chengyan, Xu, Xian, Hou, Bowen, Qiu, Diana Y.

Institution: Yale University

Language: English

Topic: Machine Learning in Materials Science

Subfield: Materials Chemistry

Field: Materials Science

Domain: Physical Sciences

Open Access status: green

Funders: U.S. Department of Energy, National Energy Research Scientific Computing Center, Office of Science, Basic Energy Sciences, Argonne National Laboratory, Oak Ridge National Laboratory, Lawrence Berkeley National Laboratory

OpenAlex ID: W7214682139