Files changed (1) hide show
  1. README.md +81 -0
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.