xnn.gnn.models.dimenet#
DimeNet and DimeNet++ (Gasteiger et al.) – directional message passing.
Faithful implementations of the architectures described in
J. Gasteiger, J. Gross, S. Guennemann, “Directional Message Passing for Molecular Graphs”, ICLR 2020 (arXiv:2003.03123) – DimeNet; and
J. Gasteiger, S. Giri, J. T. Margraf, S. Guennemann, “Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules”, NeurIPS-W 2020 (arXiv:2011.14115) – DimeNet++,
built from the papers’ equations on the xnn abstractions
(BesselRBF,
PolynomialCutoff,
SphericalBesselBasis,
scatter_sum(), …) and consistent with the
conventions of the authors’ reference code (parameter layout, initialization,
basis ordering), without copying it.
Directional message passing (ICLR paper section 4). Instead of atoms, the
network embeds the directed edges j -> i of the cutoff graph: the message
m_ji carries the direction from j to i. A message is updated from
the messages m_kj arriving at j from its other neighbors
k != i, using the distance d_ji and, jointly, the distance d_kj
and the angle alpha_(kj,ji) = angle(x_k, x_j, x_i) at j (eq 4),
m_ji^(l+1) = f_update(m_ji^(l), sum_k f_int(m_kj^(l), e_RBF(d_ji), a_SBF(d_kj, alpha))).
Representations (section 5): the distances enter through the Bessel radial
basis e_RBF,n(d) = sqrt(2/c) sin(n pi d / c) / d (eq 7, frequencies
fine-tuned by backpropagation) and distance/angle pairs through the 2D
spherical Fourier-Bessel basis a_SBF,ln (eq 6), both multiplied by the
polynomial envelope u(d) of eq 8 (p = 6) so that the model is twice
continuously differentiable and the autograd forces of
ForceStressOutput are continuous.
Architecture (ICLR Fig. 4, hyperparameters of its Appendix B):
embedding block (eq 9):
m_ji = sigma([h_j || h_i || sigma(W_e e_RBF + b_e)] W + b)from learned atom-type embeddingshof widthF;Tinteraction blocks:m_jiandm_kjeach pass a dense layer; thekjbranch is gated element-wise by a linear map ofe_RBF(d_ji)and, per triplet, contracted with a linear map ofa_SBFthrough a bilinear weight tensor of sizeN_bilinear x F x F; the triplet terms are summed overk, added to thejibranch, refined by residual blocks (two dense layers and a skip), added to the block input through a skip connection and refined by further residual blocks;output blocks after the embedding and after every interaction block:
m_jigated by a linear map ofe_RBF(d_ji), summed over the incoming messages of every atomi(h_i = sum_j m_ji), passed through dense layers and a final linear layer to the atom-wise outputt_i^(l); the per-atom energy issum_l t_i^(l)and the total energy the sum over atoms.
DimeNet++ (NeurIPS-W paper Fig. 1 and section 2) keeps the embedding, the
bases and the skip structure and changes the expensive parts: the bilinear
layer becomes a Hadamard product, the basis representations pass through
two-layer linear maps (N_rbf -> N_basis -> F and
N_sbf -> N_basis -> N_triplet), the triplet branch is down-projected to
N_triplet features before the aggregation and up-projected after it, the
output block up-projects the atom features to N_out before its dense
layers, and four blocks replace the six of DimeNet. The activation is the
self-gated swish sigma(x) = x sigmoid(x) throughout (torch.nn.SiLU).
Conventions shared with the reference implementation, for transplantable
weights: the atom embedding is indexed by the atomic number; the columns of
a_SBF run degree-major (l * N_SRBF + n); dense layers are initialized
with Glorot-scaled orthogonal matrices and zero biases, the bilinear tensor
from N(0, (2/F)^2), the atom embeddings uniformly on
[-sqrt(3), sqrt(3)], and the final linear layer of each output block
either to zero or Glorot-orthogonally (output_init; the reference uses
the latter for energies). Where the papers and the reference code differ, the model follows the paper
unless an option says otherwise (see DimeNet): the angle is the
paper’s alpha_(kj,ji) (the code’s angle is its supplement pi - alpha,
a sign flip of the odd-l basis functions that a weight transplant
absorbs); the envelope of a_SBF is the paper’s u(d) unless
reference_basis is set (the code’s carries an extra c / d_kj); and
the radial gate of the incoming messages is the code’s e_RBF(d_kj) by
default (radial_gate="kj", what the published models were trained with),
while eq 4 as printed, e_RBF(d_ji), is radial_gate="ji".
The only addition is atom_ref, the learnable per-element reference energy
every xnn potential carries (zero by default, so inert unless set/trained);
the paper trains on atomization energies, i.e. with these references
subtracted. Periodic systems are handled by the edge vectors of the graph
(cell shifts included); a triplet excludes only the exact reverse image of
the message’s own edge, so an atom and its periodic images are distinct
neighbors.
Functions
|
Enumerate the message pairs |
Classes