The standalone attention module Attention(•, •, •) receives input query, key, and value vectors, of dimension d k , d k , and d v respectively and gathered in matrices Q, K, V , respectively. With ω(•) = softmax(•/ √ d k ), the attention operation computes a weighted sum of the values where a large querykey dot product assigns a larger weight to the corresponding value. In multi-head attention, Q, K, V are projected to new dimensions by learnt projections W Q i , W K i , W V i , respectively, for a total of m independent times. The results are processed by an attention module, with the concatenated output attention heads being projected with W O . Following the original Set Transformer implementation, we use the lower-level multihead and set attention blocks (MABs, respectively SABs) and pooling by multihead attention (PMA) to define an encoder-decoder architecture that embeds input atom vectors into a chosen dimension d, then learns to aggregate or compress the encoded representations into a single vector, the molecule representation. The encoder is defined as