# Atom Type Embedding

## Overview

Here is an overview of the DeePMD-kit algorithm. Given a specific centric atom, we can obtain the matrix describing its local environment, named \(\mathcal R\). It consists of the distance between the centric atom and its neighbors, as well as a direction vector. We can embed each distance into a vector of \(M_1\) dimension by an `embedding net`

, so the environment matrix \(\mathcal R\) can be embedded into matrix \(\mathcal G\). We can thus extract a descriptor vector (of \(M_1 \times M_2\) dim) of the centric atom from the \(\mathcal G\) by some matrix multiplication, and put the descriptor into `fitting net`

to get the predicted energy \(E\). The vanilla version of DeePMD-kit builds `embedding net`

and `fitting net`

relying on the atom type, resulting in \(O(N)\) memory usage. After applying atom type embedding, in DeePMD-kit v2.0, we can share one `embedding net`

and one `fitting net`

in total, which reduces training complexity largely.

## Preliminary

In the following chart, you can find the meaning of symbols used to clarify the atom-type embedding algorithm.

\(i\): Type of centric atom

\(j\): Type of neighbor atom

\(s_{ij}\): Distance between centric atom and neighbor atom

\(\mathcal G_{ij}(\cdot)\): Origin embedding net, take \(s_{ij}\) as input and output embedding vector of \(M_1\) dim

\(\mathcal G(\cdot)\): Shared embedding net

\(\text{Multi}(\cdot)\): Matrix multiplication and flattening, output the descriptor vector of \(M_1\times M_2\) dim

\(F_i(\cdot)\): Origin fitting net, take the descriptor vector as input and output energy

\(F(\cdot)\): Shared fitting net

\(A(\cdot)\): Atom type embedding net, input is atom type, the output is type embedding vector of dim `nchanl`

So, we can formulate the training process as follows. Vanilla DeePMD-kit algorithm:

DeePMD-kit applying atom type embedding:

or

The difference between the two variants above is whether using the information of centric atom when generating the descriptor. Users can choose by modifying the `type_one_side`

hyper-parameter in the input JSON file.

## How to use

A detailed introduction can be found at `se_e2_a_tebd`

. Looking for a fast start-up, you can simply add a `type_embedding`

section in the input JSON file as displayed in the following, and the algorithm will adopt the atom type embedding algorithm automatically. An example of `type_embedding`

is like

```
"type_embedding":{
"neuron": [2, 4, 8],
"resnet_dt": false,
"seed": 1
}
```

## Code Modification

Atom-type embedding can be applied to varied `embedding net`

and `fitting net`

, as a result, we build a class `TypeEmbedNet`

to support this free combination. In the following, we will go through the execution process of the code to explain our code modification.

### trainer (train/trainer.py)

In trainer.py, it will parse the parameter from the input JSON file. If a `type_embedding`

section is detected, it will build a `TypeEmbedNet`

, which will be later input in the `model`

. `model`

will be built in the function `_build_network`

.

### model (model/ener.py)

When building the operation graph of the `model`

in `model.build`

. If a `TypeEmbedNet`

is detected, it will build the operation graph of `type embed net`

, `embedding net`

and `fitting net`

by order. The building process of `type embed net`

can be found in `TypeEmbedNet.build`

, which output the type embedding vector of each atom type (of [\(\text{ntypes} \times \text{nchanl}\)] dimensions). We then save the type embedding vector into `input_dict`

, so that they can be fetched later in `embedding net`

and `fitting net`

.

### embedding net (descriptor/se*.py)

In `embedding net`

, we shall take local environment \(\mathcal R\) as input and output matrix \(\mathcal G\). Functions called in this process by the order is

```
build -> _pass_filter -> _filter -> _filter_lower
```

`_pass_filter`

: It will first detect whether an atom type embedding exists, if so, it will apply atom type embedding algorithm and doesn’t divide the input by type.

`_filter`

: It will call `_filter_lower`

function to obtain the result of matrix multiplication (\(\mathcal G^T\cdot \mathcal R\)), do further multiplication involved in \(\text{Multi}(\cdot)\), and finally output the result of descriptor vector of \(M_1 \times M_2\) dim.

`_filter_lower`

: The main function handling input modification. If type embedding exists, it will call `_concat_type_embedding`

function to concat the first column of input \(\mathcal R\) (the column of \(s_{ij}\)) with the atom type embedding information. It will decide whether to use the atom type embedding vector of the centric atom according to the value of `type_one_side`

(if set **True**, then we only use the vector of the neighbor atom). The modified input will be put into the `fitting net`

to get \(\mathcal G\) for further matrix multiplication stage.

### fitting net (fit/ener.py)

In `fitting net`

, it takes the descriptor vector as input, whose dimension is [natoms, \(M_1\times M_2\)]. Because we need to involve information on the centric atom in this step, we need to generate a matrix named `atype_embed`

(of dim [natoms, nchanl]), in which each row is the type embedding vector of the specific centric atom. The input is sorted by type of centric atom, we also know the number of a particular atom type (stored in `natoms[2+i]`

), thus we get the type vector of the centric atom. In the build phase of the fitting net, it will check whether type embedding exists in `input_dict`

and fetch them. After that, call `embed_atom_type`

function to look up the embedding vector for the type vector of the centric atom to obtain `atype_embed`

, and concat input with it ([input, atype_embed]). The modified input goes through `fitting`

net` to get predicted energy.

Note

You can’t apply the compression method while using atom-type embedding.