Update README.md
#2
by
heloise-chomet
- opened
README.md
CHANGED
|
@@ -1,3 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
## License summary
|
| 2 |
|
| 3 |
1. The Licensed Models are **only** available under this License for Non-Commercial Purposes.
|
|
|
|
| 1 |
+
|
| 2 |
+
# Nequip
|
| 3 |
+
|
| 4 |
+
## Reference
|
| 5 |
+
|
| 6 |
+
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa,
|
| 7 |
+
Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky.
|
| 8 |
+
E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.
|
| 9 |
+
Nature Communications, 13(1), May 2022. ISSN: 2041-1723. URL: https://dx.doi.org/10.1038/s41467-022-29939-5.
|
| 10 |
+
|
| 11 |
+
## Hyperparameters, model configurations and training strategies
|
| 12 |
+
|
| 13 |
+
### Model architecture
|
| 14 |
+
| Parameter | Value | Description |
|
| 15 |
+
|---------------------------|-----------------------------------------------|---------------------------------------------|
|
| 16 |
+
| `num_layers` | `5` | Number of NequIP layers. |
|
| 17 |
+
| `node_irreps` | `64x0e + 64x0o + 32x1e + 32x1o + 4x2e + 4x2o` | O3 representation space of node features. |
|
| 18 |
+
| `l_max` | `2` | Maximal degree of spherical harmonics. |
|
| 19 |
+
| `num_bessel` | `8` | Number of Bessel basis functions. |
|
| 20 |
+
| `radial_net_nonlinearity` | `swish` | Activation function for radial MLP. |
|
| 21 |
+
| `radial_net_n_hidden` | `64` | Number of hidden features in radial MLP. |
|
| 22 |
+
| `radial_net_n_layers` | `2` | Number of layers in radial MLP. |
|
| 23 |
+
| `radial_envelope` | `polynomial_envelope` | Radial envelope function. |
|
| 24 |
+
| `scalar_mlp_std` | `4` | Standard deviation of weight initialisation.|
|
| 25 |
+
| `atomic_energies` | `None` | Treatment of the atomic energies. |
|
| 26 |
+
| `avg_um_neighbors` | `None` | Mean number of neighbors. |
|
| 27 |
+
|
| 28 |
+
### Training
|
| 29 |
+
| Parameter | Value | Description |
|
| 30 |
+
|--------------------------|--------|--------------------------------------------------|
|
| 31 |
+
| `num_epochs` | `220` | Number of epochs to run. |
|
| 32 |
+
| `ema_decay` | `0.99` | The EMA decay rate. |
|
| 33 |
+
| `eval_num_graphs` | `None` | Number of validation set graphs to evaluate on. |
|
| 34 |
+
| `use_ema_params_for_eval`| `True` | Whether to use the EMA parameters for evaluation.|
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
### Optimizer
|
| 38 |
+
| Parameter | Value | Description |
|
| 39 |
+
|----------------------------------|----------------|-----------------------------------------------------------------|
|
| 40 |
+
| `init_learning_rate` | `0.002` | Initial learning rate. |
|
| 41 |
+
| `peak_learning_rate` | `0.002` | Peak learning rate. |
|
| 42 |
+
| `final_learning_rate` | `0.002` | Final learning rate. |
|
| 43 |
+
| `weight_decay` | `0` | Weight decay. |
|
| 44 |
+
| `warmup_steps` | `4000` | Number of optimizer warm-up steps. |
|
| 45 |
+
| `transition_steps` | `360000` | Number of optimizer transition steps. |
|
| 46 |
+
| `grad_norm` | `500` | Gradient norm used for gradient clipping. |
|
| 47 |
+
| `num_gradient_accumulation_steps`| `1` | Steps to accumulate before taking an optimizer step. |
|
| 48 |
+
| `algorithm` | `optax.amsgrad`| The AMSGrad optimizer. |
|
| 49 |
+
| `b1` | `0.9` | Exponential decay rate to track first moment of past gradients. |
|
| 50 |
+
| `b2` | `0.999` | Exponential decay rate to track second moment of past gradients.|
|
| 51 |
+
| `eps` | `1e-8` | Constant applied to denominator outside the square root. |
|
| 52 |
+
| `eps_root` | `0.0` | Constant applied to denominator inside the square root. |
|
| 53 |
+
|
| 54 |
+
### Huber Loss Energy weight schedule
|
| 55 |
+
| Parameter | Value | Description |
|
| 56 |
+
|------------------------|------------------------------------|-------------------------------------------------------------------------------------------------|
|
| 57 |
+
| `schedule` | `optax.piecewise_constant_schedule`| Piecewise constant schedule with scaled jumps at specific boundaries. |
|
| 58 |
+
| `init_value` | `40` | Initial value. |
|
| 59 |
+
| `boundaries_and_scale` | `{115: 25}` | Dictionary of {step: scale} where scale is multiplied into the schedule value at the given step.|
|
| 60 |
+
### Huber Loss Force weight schedule
|
| 61 |
+
| Parameter | Value | Description |
|
| 62 |
+
|------------------------|------------------------------------|-------------------------------------------------------------------------------------------------|
|
| 63 |
+
| `schedule` | `optax.piecewise_constant_schedule`| Piecewise constant schedule with scaled jumps at specific boundaries. |
|
| 64 |
+
| `init_value` | `1000` | Initial value. |
|
| 65 |
+
| `boundaries_and_scale` | `{115: 0.04}` | Dictionary of {step: scale} where scale is multiplied into the schedule value at the given step.|
|
| 66 |
+
|
| 67 |
+
### Dataset
|
| 68 |
+
| Parameter | Value | Description |
|
| 69 |
+
|-----------------------------|-------|--------------------------------------------|
|
| 70 |
+
| `graph_cutoff_angstrom` | `5` | Graph cutoff distance (in Å). |
|
| 71 |
+
| `max_n_node` | `32` | Maximum number of nodes allowed in a batch.|
|
| 72 |
+
| `max_n_edge` | `288` | Maximum number of edges allowed in a batch.|
|
| 73 |
+
| `batch_size` | `16` | Number of graphs in a batch. |
|
| 74 |
+
|
| 75 |
+
This model was trained on the [SPICE2_curated dataset](https://huggingface.co/datasets/InstaDeepAI/SPICE2-curated).
|
| 76 |
+
|
| 77 |
+
## How to Use
|
| 78 |
+
|
| 79 |
+
For complete usage instructions and more information, please refer to our [documentation](https://instadeep.github.io/mlip)
|
| 80 |
+
|
| 81 |
+
|
| 82 |
## License summary
|
| 83 |
|
| 84 |
1. The Licensed Models are **only** available under this License for Non-Commercial Purposes.
|