To design an expressive pooling function that considers the entire context of the molecule (i.e. all computed atom representations) and is self-contained (does not require any additional inputs), we leverage the existing Set Transformer framework introduced by Lee et al. for set modelling , and proposed by Buterez et al. for use in node aggregation . In other words, the pooling operation is reframed as a set summarisation task, with an output that corresponds to the desired molecular embedding. This is achieved by assembling building blocks defined using a standard multi-head attention mechanism: