Implement the Psiformer, from A Self-Attention Ansatz for Ab Initio Quantum Chemistry, ICLR 2023.

PiperOrigin-RevId: 513235453
Change-Id: Ic38fe88da627debc51e3bd063d873a21f45e46b0

Author: Ingrid von Glehn

ferminet/base_config.py

@@ -176,14 +176,26 @@ def default() -> ml_collections.ConfigDict:
           'blocks': 1,  # Number of blocks to split the MCMC sampling into
       },
       'network': {
-          'detnet': {
+          'network_type': 'ferminet',  # One of 'ferminet' or 'psiformer'.
+          # Config specific to original FermiNet architecture.
+          # Only used if network_type is 'ferminet'.
+          'ferminet': {
               'hidden_dims': ((256, 32), (256, 32), (256, 32), (256, 32)),
-              'determinants': 16,
+              # Whether to use the last layer of the two-electron stream of the
+              # FermiNet.
+              'use_last_layer': False,
           },
+          # Only used if network_type is 'psiformer'.
+          'psiformer': {
+              'num_layers': 2,
+              'num_heads': 4,
+              'heads_dim': 64,
+              'mlp_hidden_dims': (256,),
+              'use_layer_norm': False,
+          },
+          # Config common to all architectures.
+          'determinants': 16,  # Number of determinants.
           'bias_orbitals': False,  # include bias in last layer to orbitals
-          # Whether to use the last layer of the two-electron stream of the
-          # DetNet
-          'use_last_layer': False,
           # If true, determinants are dense rather than block-sparse
           'full_det': True,
           # If specified, include a pre-determinant Jastrow factor.

ferminet/networks.py

@@ -169,6 +169,7 @@ def __call__(
       r_ae: jnp.ndarray,
       ee: jnp.ndarray,
       r_ee: jnp.ndarray,
+      spins: jnp.ndarray,
       charges: jnp.ndarray,
   ) -> jnp.ndarray:
     """Forward evaluation of the equivariant interaction layers.
@@ -179,6 +180,7 @@ def __call__(
       r_ae: electron-nuclear distances.
       ee: electron-electron vectors.
       r_ee: electron-electron distances.
+      spins: spin of each electron.
       charges: nuclear charges.
 
     Returns:
@@ -533,6 +535,7 @@ def apply(
       r_ae: jnp.ndarray,
       ee: jnp.ndarray,
       r_ee: jnp.ndarray,
+      spins: jnp.ndarray,
       charges: jnp.ndarray,
   ) -> jnp.ndarray:
     """Applies the FermiNet interaction layers to a walker configuration.
@@ -543,14 +546,15 @@ def apply(
       r_ae: electron-nuclear distances.
       ee: electron-electron vectors.
       r_ee: electron-electron distances.
+      spins: spin of each electron.
       charges: nuclear charges.
 
     Returns:
       Array of shape (nelectron, output_dim), where the output dimension,
       output_dim, is given by init, and is suitable for projection into orbital
       space.
     """
-    del charges  # Unused.
+    del spins, charges  # Unused.
 
     ae_features, ee_features = options.feature_layer.apply(
         ae=ae, r_ae=r_ae, ee=ee, r_ee=r_ee, **params['input']
@@ -578,7 +582,7 @@ def apply(
 def make_orbitals(
     nspins: Tuple[int, ...],
     charges: jnp.ndarray,
-    options: FermiNetOptions,
+    options: BaseNetworkOptions,
     equivariant_layers: Tuple[InitLayersFn, ApplyLayersFn],
 ) -> ...:
   """Returns init, apply pair for orbitals.
@@ -678,11 +682,15 @@ def apply(
       columns under the exchange of inputs of shape (ndet, nalpha+nbeta,
       nalpha+nbeta) (or (ndet, nalpha, nalpha) and (ndet, nbeta, nbeta)).
     """
-    del spins
-
     ae, ee, r_ae, r_ee = construct_input_features(pos, atoms, ndim=options.ndim)
     h_to_orbitals = equivariant_layers_apply(
-        params['layers'], ae=ae, r_ae=r_ae, ee=ee, r_ee=r_ee, charges=charges
+        params['layers'],
+        ae=ae,
+        r_ae=r_ae,
+        ee=ee,
+        r_ee=r_ee,
+        spins=spins,
+        charges=charges,
     )
 
     if options.envelope.apply_type == envelopes.EnvelopeType.PRE_ORBITAL:

ferminet/pbc/tests/hamiltonian_test.py

@@ -52,9 +52,8 @@ def test_periodicity(self):
         envelope=envelopes.make_multiwave_envelope(kpoints),
         feature_layer=feature_layer,
         bias_orbitals=cfg.network.bias_orbitals,
-        use_last_layer=cfg.network.use_last_layer,
         full_det=cfg.network.full_det,
-        **cfg.network.detnet
+        **cfg.network.ferminet,
     )
 
     key, subkey = jax.random.split(key)