MBFormer: Transformer Model for Excited States

• MBFormer is a transformer-based approach that learns many-body interactions from mean-field inputs to predict excited-state observables within the GW-BSE framework.
• It employs a symmetry-aware, grid-free tokenization of Kohn–Sham states using E(2)-equivariant embeddings to handle variable system sizes in 2D materials.
• The model achieves high accuracy in predicting quasiparticle and exciton properties, with interpretability provided by attention mechanisms linking physical electron–hole interactions.

Searching arXiv for the specified paper and closely related context. arXiv search query: (Hou et al., 7 Jul 2025) MBFormer transformer many-body materials GW BSE MBFormer is a transformer-based learning paradigm for quantum many-body interactions in real materials, introduced as a symmetry-aware, grid-free model that learns the many-body hierarchy directly from mean-field inputs and predicts excited-state quantities within the GW plus Bethe–Salpeter equation (GW-BSE) formalism (Hou et al., 7 Jul 2025). It is designed for regimes in which the effective interactions are nonlocal in space, strongly energy-dependent, and dependent on the full manifold of occupied and unoccupied mean-field states across the Brillouin zone, conditions under which nearsightedness-based surrogates for ground-state Hamiltonians are not adequate. In the reported implementation, MBFormer is trained on 721 two-dimensional materials from the C2DB database and targets state-resolved quasiparticle energies, exciton energies, exciton oscillator strengths, and exciton wavefunction distributions.

1. Problem setting and conceptual scope

MBFormer is motivated by the difficulty of predicting excited-state properties in real materials when the relevant interactions are both nonlocal and energy-dependent. In many-body perturbation theory, the quasiparticle self-energy $\Sigma$ depends on the entire manifold of occupied and unoccupied mean-field states across the Brillouin zone and on frequency $\omega$, while the electron–hole Bethe–Salpeter kernel $K$ couples electron–hole basis states nonlocally in both $k$-space and band space (Hou et al., 7 Jul 2025). These dependencies imply variable, system-dependent basis sizes and invalidate assumptions that underlie tight-binding models and many ML surrogates for ground-state Hamiltonians.

Within this setting, MBFormer learns a post-DFT mapping from Kohn–Sham states and metadata to excited-state observables. The learned functional is expressed as

$$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$

The predicted quantities are explicitly state-resolved: quasiparticle energies $E^{\mathrm{QP}}_{n\mathbf{k}}$, exciton energies $\Omega_S$, exciton oscillator strengths $f_S$, and exciton wavefunction distributions represented either as amplitudes $A_{vc\mathbf{k}}^{(S)}$ or derived spatial electron–hole probability densities.

A common misconception would be to regard MBFormer as a direct replacement for ab initio GW-BSE solvers. The reported formulation is more specific: it is trained to mimic post-DFT many-body perturbation theory functionals from mean-field inputs. This distinction matters because the model inherits its targets, and therefore its domain of validity, from the GW-BSE data on which it is supervised.

2. Relation to GW and the Bethe–Salpeter equation

The GW component is framed around the frequency-dependent self-energy operator

$\Sigma(\omega)\ =\ i\, G(\omega)\, W(\omega),$

or, more explicitly,

$\omega$0

The quasiparticle energy for a Kohn–Sham state $\omega$1 is written as

$\omega$2

with renormalization factor

$\omega$3

In the MBFormer experiments, the learned target is the scalar G$\omega$4W$\omega$5 correction per $\omega$6,

$\omega$7

The excitonic component is based on the Bethe–Salpeter eigenvalue problem in the electron–hole basis $\omega$8:

$\omega$9

where $K$0 contains screened direct and typically unscreened exchange electron–hole interactions. Oscillator strength is represented as

$K$1

and optical absorption can be expressed through

$K$2

The exciton wavefunction distribution is

$K$3

with marginal spatial densities obtained by integration over one coordinate.

The role of MBFormer in this formal structure is to approximate the action of the many-body self-energy and electron–hole interaction kernels on the relevant basis. This suggests that the model is not merely fitting scalar spectra, but is organized around the same hierarchy that distinguishes one-particle and two-particle sectors in MBPT.

3. Architecture: symmetry-aware, grid-free transformer

MBFormer is described as grid-free because it operates in the mean-field basis rather than on fixed real-space grids (Hou et al., 7 Jul 2025). Its inputs are per-state tokens carrying learned descriptors of Kohn–Sham wavefunctions together with metadata, specifically energies $K$4, band indices $K$5, and momenta $K$6. The architecture uses two disjoint token sets that mirror the structure of MBPT. The source set is a large collection of Kohn–Sham states, $K$7, representing the ground-state environment. The target set is task-dependent: single-particle states $K$8 for GW and electron–hole pairs $K$9 for BSE. Because the model is input-size agnostic, it can accommodate variable numbers of source and target tokens across materials.

Symmetry-awareness is introduced through an E(2)-equivariant variational autoencoder, denoted E2-VAE, which integrates E2-CNN layers into a U-Net. This embedding enforces equivariance to planar rotations and translations and robustness to coordinate permutations and changes in lattice size, properties that are directly relevant for 2D materials and varying simulation cells and $k$0-grids. Sinusoidal positional embeddings are applied to momentum coordinates $k$1, band indices $k$2, and energies $k$3 so that reciprocal-space meta-information is included without imposing an arbitrary token ordering.

Tokenization combines latent wavefunction information and metadata. Real-space information from $k$4 is compressed into a latent vector $k$5, which is fused with encoded $k$6, $k$7, and $k$8. For the GW task, target tokens are single-particle tokens with initialization $k$9. For the BSE task, electron–hole pair tokens are assembled by concatenation or tensor product:

$$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$0

The transformer itself uses encoder-decoder stacks with attention, residual connections, and Add & Norm (LayerNorm) blocks. The encoder performs self-attention over source tokens to produce a latent representation of the ground-state environment. The decoder processes target tokens with self-attention and cross-attends to the encoder latent. Attention is the mechanism through which nonlocal, energy-dependent coupling is represented:

$$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$1

Because the learned projections depend on descriptors that include band energies, $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$2-vectors, and compressed wavefunction features, attention provides a route for coupling target excitations to the entire source manifold. In the BSE decoder, the normalized attention weights are directly used to predict wavefunction moduli $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$3.

4. Output heads, objectives, and training protocol

The architecture supports multiple output heads within a unified encoder-decoder framework (Hou et al., 7 Jul 2025). For GW, the output is a scalar $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$4 for each target single-particle token, with quasiparticle energies obtained as

$$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$5

For BSE, the model predicts exciton eigenvalues $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$6, oscillator strengths $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$7, and distributional outputs $$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$8 derived from attention rows.

The training objectives are correspondingly heterogeneous. Energies are optimized with mean squared error or mean absolute error,

$$f:\ \Big(\{\phi_{n\mathbf{k}}(\mathbf{r}),\,\epsilon_{n\mathbf{k}},\,n,\,\mathbf{k}\}\Big)\ \mapsto\ \Big(E_{n\mathbf{k}}^{\mathrm{QP}},\ \Omega_S,\ f_S,\ A_{vc\mathbf{k}}^{(S)}\Big).$$9

Oscillator strengths use MSE or MAE on max-normalized dipoles,

$E^{\mathrm{QP}}_{n\mathbf{k}}$0

Wavefunction amplitudes are trained with Jensen–Shannon divergence between normalized distributions:

$E^{\mathrm{QP}}_{n\mathbf{k}}$1

Multi-task training uses weighted sums such as

$E^{\mathrm{QP}}_{n\mathbf{k}}$2

The dataset comprises 721 two-dimensional materials from the Computational 2D Materials Database. Ground-state DFT supplies Kohn–Sham orbitals and energies, while GW-BSE calculations provide supervision for quasiparticle corrections, exciton energies, oscillator strengths, and exciton amplitudes. Preprocessing extracts $E^{\mathrm{QP}}_{n\mathbf{k}}$3 in plane-wave representation, compresses them via E2-VAE, and attaches positional encodings for $E^{\mathrm{QP}}_{n\mathbf{k}}$4, $E^{\mathrm{QP}}_{n\mathbf{k}}$5, and $E^{\mathrm{QP}}_{n\mathbf{k}}$6. Tokens are then partitioned into source and target sets according to the task. For GW, source tokens include many occupied and unoccupied states and targets are selected single-particle states. For BSE, the source is analogous and the targets are electron–hole pairs formed near the band edges.

The reported train-test protocols are task-specific. For GW, the split is 80% train and 20% test across the 721 materials. For BSE, it is 90% train and 10% test. An additional augmentation study is carried out for monolayer hBN, where 900 BSE Hamiltonians are generated over different $E^{\mathrm{QP}}_{n\mathbf{k}}$7-mesh sizes and small shifts in the Brillouin zone; training on 6 coarse meshes with 84 augmentations, approximately 10% of the dataset, is used for inference on fine unseen $E^{\mathrm{QP}}_{n\mathbf{k}}$8-grids. Optimization follows a standard transformer pipeline with AdamW-style decoupled weight decay, LayerNorm, residual connections, feed-forward sublayers, weight decay, and early stopping on validation metrics.

5. Quantitative performance, ablations, and interpretability

On the GW task, MBFormer achieves mean absolute error of approximately $E^{\mathrm{QP}}_{n\mathbf{k}}$9 and $\Omega_S$0 across test materials (Hou et al., 7 Jul 2025). The mean of the G$\Omega_S$1W$\Omega_S$2 corrections in the test set is approximately $\Omega_S$3 with standard deviation approximately $\Omega_S$4. The model reproduces smooth quasiparticle bandstructures and shows the expected dependence on the number of source states, described in the paper as “empty-band” convergence behavior. In a cross-attention ablation with variable source size, increasing source states from 1 to 60 conduction bands reduces validation MSE from approximately $\Omega_S$5 to approximately $\Omega_S$6, corresponding to approximately 43% improvement.

On the BSE task, the reported exciton-energy performance is MAE approximately $\Omega_S$7 over all excitons, improving to approximately $\Omega_S$8 for the first 100 lowest-energy excitons, with $\Omega_S$9. The hBN $f_S$0-grid inference study yields exciton-energy performance of $f_S$1 and MAE approximately $f_S$2 on fine unseen $f_S$3-grids. For oscillator strengths, the reported dipole MAE is approximately $f_S$4 after max-normalization, and reconstructed absorption spectra $f_S$5 match BSE. For exciton amplitudes, attention-derived distributions $f_S$6 reproduce BSE amplitudes with low JSD below 0.05 for band-edge excitons, while higher-energy excitons show larger JSD due to mixing and degeneracy.

The ablation results identify attention as the critical mechanism. For BSE, including self-attention reduces validation MSE from approximately $f_S$7 to approximately $f_S$8, about 75% improvement, and reduces MAE from approximately $f_S$9 to approximately $A_{vc\mathbf{k}}^{(S)}$0, about 62.5% improvement. Increasing depth without attention, including 20-layer models, does not recover this gain. For GW, reducing the source-state count worsens validation loss, consistent with missing scattering channels and poorer approximation of the polarizability.

Interpretability is tied directly to the attention mechanism. Attention maps identify which mean-field states most strongly couple to a given target excitation. In the BSE setting, the attention row for exciton $A_{vc\mathbf{k}}^{(S)}$1 approximates the distribution $A_{vc\mathbf{k}}^{(S)}$2, linking learned weights to physical electron–hole amplitudes across $A_{vc\mathbf{k}}^{(S)}$3 and band indices. In the GW setting, cross-attention analysis shows improved prediction when more high-energy unoccupied states contributing to screening are included. A plausible implication is that MBFormer’s interpretability is not merely post hoc, but partly embedded in how the decoder represents many-body couplings.

6. Positioning, computational characteristics, and limitations

The paper positions MBFormer against earlier ML approaches for excited states in two ways (Hou et al., 7 Jul 2025). Relative to prior ML for GW bandstructures that handcraft state representations, exemplified by Knøsgaard–Thygesen, MBFormer is reported to attain state-of-the-art MAE of approximately $A_{vc\mathbf{k}}^{(S)}$4 while producing smooth bandstructures through learned latent representations of Kohn–Sham states. Relative to direct spectrum predictors such as GNNopt for absorption, it yields state-resolved excitonic quantities $A_{vc\mathbf{k}}^{(S)}$5, $A_{vc\mathbf{k}}^{(S)}$6, and $A_{vc\mathbf{k}}^{(S)}$7, which enables interpretability and downstream physics analyses. The important distinction is that MBFormer is structured around state-level many-body observables rather than solely around aggregate spectra.

The computational profile follows from token-based attention. Per layer, the attention mechanism scales as $A_{vc\mathbf{k}}^{(S)}$8, with $A_{vc\mathbf{k}}^{(S)}$9 given by the number of source plus target states selected for the task. In the reported formulation, this token count is typically orders of magnitude smaller than the full plane-wave grid. Compared to ab initio GW with naive $\Sigma(\omega)\ =\ i\, G(\omega)\, W(\omega),$0 scaling and BSE with $\Sigma(\omega)\ =\ i\, G(\omega)\, W(\omega),$1 scaling in system size, MBFormer inference is described as lightweight once trained and capable of delivering orders-of-magnitude speedups for excited-state prediction across materials. Variable numbers of states are handled through masking, padding, and batching, and the combination of grid-free tokenization, positional encodings, and $\Sigma(\omega)\ =\ i\, G(\omega)\, W(\omega),$2-grid augmentation supports interpolation across different samplings.

The stated limitations are domain-specific rather than incidental. The model is trained on 2D materials, so transfer to 3D or strongly correlated systems may require retraining and revised symmetry encoders, including E(3)-equivariant extensions. Predictions are sensitive to the quality of the mean-field inputs, so inaccurate Kohn–Sham states can degrade performance, especially for subtle many-body features. Exciton degeneracy and continuum states produce larger JSD for amplitudes and greater variability in oscillator strengths at higher energies. Frequency dependence is learned only implicitly through metadata, and the authors note that explicit $\Sigma(\omega)\ =\ i\, G(\omega)\, W(\omega),$3-dependent features may improve self-energy modelling.

The extension directions identified in the paper include 3D materials with an E(3)-equivariant embedder; spin–orbit coupling and magnetic order with appropriate symmetry encodings; phonons and temperature effects through electron–phonon couplings and finite-temperature screening; prediction of mobilities, lifetimes, and linewidths; defects, heterostructures, and interfaces; charged excitations such as trions; biexcitons; polaritons under light–matter coupling; and dynamical many-body phenomena through temporal or frequency channels. These directions indicate that MBFormer is intended not only as a model for GW-BSE observables in 2D materials, but as a general transformer-based template for many-body learning in materials systems, provided that the symmetry embedding and supervision targets are adapted appropriately.