Atom Type Embedding


Here is an overview of the deepmd-kit algorithm. Given a specific centric atom, we can obtain the matrix describing its local environment, named as R. It is consist of the distance between centric atom and its neighbors, as well as a direction vector. We can embed each distance into a vector of M1 dimension by a embedding net, so the environment matrix R can be embed into matrix G. We can thus extract a descriptor vector (of M1*M2 dim) of the centric atom from the G by some matrix multiplication, and put the descriptor into fitting net to get predicted energy E. The vanilla version of deepmd-kit build 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 decline training complexity largely.


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




Type of centric atom


Type of neighbor atom


Distance between centric atom and neighbor atom


Origin embedding net, take s_ij as input and output embedding vector of M1 dim


Shared embedding net


Matrix multiplication and flattening, output the descriptor vector of M1*M2 dim


Origin fitting net, take the descriptor vector as input and output energy


Shared fitting net


Atom type embedding net, input is atom type, output is type embedding vector of dim nchanl

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

Energy = F_i( Multi( G_ij( s_ij ) ) )

Deepmd-kit applying atom type embedding:

Energy = F( [ Multi( G( [s_ij, A(i), A(j)] ) ), A(j)] )


Energy = F( [ Multi( G( [s_ij, A(j)] ) ), A(j)] )

The difference between 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 atom type embedding algorithm automatically. An example of type_embedding is like

       "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/

In, 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/

When building the operation graph of the model in 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, which output the type embedding vector of each atom type (of [ntypes * nchanl] dimension). 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 R as input and output matrix G. Functions called in this process by 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 (G^T·R ), do further multiplication involved in Multi(·), and finally output the result of descriptor vector of M1*M2 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 R (the column of s_ij) with the atom type embedding information. It will decide whether using the atom type embedding vector of 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 G for further matrix multiplication stage.

fitting net (fit/

In fitting net, it take the descriptor vector as input, whose dimension is [natoms, (M1*M2)]. Because we need to involve information of centric atom in this step, we need to generate a matrix named as 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 centric atom. In the build phrase of fitting net, it will check whether type embedding exist in input_dict and fetch them. After that calling embed_atom_type function to lookup embedding vector for type vector of centric atom to obtain atype_embed, and concat input with it ([input, atype_embed]). The modified input go through fitting net to get predicted energy.

P.S.: You can’t apply compression method while using atom type embedding