Set2Set Readout for Molecular Representation

• Set2Set readout is a neural module that aggregates unordered atom features into a fixed-size, permutation-invariant molecular representation using an LSTM-based recurrent attention mechanism.
• It replaces traditional pooling by dynamically focusing on distinct molecular substructures such as functional groups, conjugated rings, or heteroatoms through learned attention weights.
• Empirical performance shows that Set2Set improves predictions for molecular properties by reducing out-of-distribution errors, especially for large molecules, compared to sum or average pooling.

Set2Set is a neural readout module introduced to aggregate sets of per-node (e.g., per-atom) feature vectors into a single, fixed-size global graph or molecule representation while preserving permutation invariance and enabling adaptive attention over elements of the input set. In the context of molecular machine learning, Set2Set functions as an expressive alternative to conventional aggregation methods such as sum- or average-pooling. As implemented in Gaul & Cuesta-Lopez (2022), Set2Set augments SchNet for the prediction of frontier molecular orbital energies (HOMO/LUMO) and achieves superior out-of-distribution generalization on large molecules relative to standard pooling approaches (Gaul et al., 2022).

1. Architectural Overview

Set2Set receives as input a set of per-atom feature vectors $\{x_1, \dots, x_{N_a}\}$, where each $x_i \in \mathbb{R}^n$ is derived from a sequence of graph convolutional or message-passing layers (e.g., SchNet). The output is a single graph/molecule feature vector $h_G$ that is both permutation-invariant and capable of selectively focusing on contextually relevant atoms via a recurrent attention mechanism.

The core mechanism is as follows:

• A small Long Short-Term Memory (LSTM) network is initialized with zero hidden/cell states $(q_0, c_0) = (0, 0)$.
• For $t=1,\ldots,T$ steps:
  • The LSTM’s current hidden state $q_{t-1}$ is used to compute attention weights $\alpha_{i,t}$ over the atoms.
  • These weights define a softmax-weighted sum of atom features to obtain a context vector $r_t$.
  • The LSTM updates its state using $r_t$ as input, yielding $(q_t, c_t)$.
• After $T$ iterations, the final molecular representation is $h_G = [q_T; r_T] \in \mathbb{R}^{2n}$.

This aggregated feature is subsequently passed to a multilayer perceptron (MLP) to predict scalar molecular properties.

2. Mathematical Formulation

The Set2Set readout is mathematically defined as:

• Initialization:

$q_0 = \mathbf{0}, \quad c_0 = \mathbf{0}$

• Attention weights (for atom $i$ at step $t$):

$$e_{i,t} = x_i^\top q_{t-1}, \qquad \alpha_{i,t} = \frac{\exp(e_{i,t})}{\sum_{j=1}^{N_a} \exp(e_{j,t})}$$

• Context vector:

$r_t = \sum_{i=1}^{N_a} \alpha_{i,t} x_i$

• LSTM update:

$$(q_t, c_t) = \mathrm{LSTM}(r_t, (q_{t-1}, c_{t-1}))$$

• Final readout:

$$h_G = \begin{bmatrix} q_T \ r_T \end{bmatrix} \in \mathbb{R}^{2n}$$

• Output prediction:

$\hat{y} = \mathrm{MLP}(h_G)$

In practice, the MLP consists of two dense layers with shifted-softplus nonlinearities to ensure numerical stability and smooth gradients.

3. Hyperparameters and Implementation Choices

The Set2Set readout implementation used the following configuration (Gaul et al., 2022):

Parameter	Value/Choice	Notes
Atomic embedding $n$	$128$	Dimension of per-atom feature vectors
Set2Set steps $T$	$3$	LSTM hidden/cell in $\mathbb{R}^{128}$
Final feature $h_G$	$\mathbb{R}^{256}$	Concatenation of $q_T$ and $r_T$
Output MLP	2 layers, size $32$	Shifted-softplus activations
SchNet backbone	6 interaction layers	25 Gaussian basis functions, $\gamma$ per SchNetPack
Optimizer	Adam	LR = $1\times10^{-4}$, decay on plateau, early stopping
Training time	$\approx$350 epochs	Quadro GP100, $\sim$5 days

Set2Set was positioned in the “aggregation” slot of SchNetPack, operating directly on atomwise embeddings output by SchNet’s interaction blocks. PyTorch Geometric’s torch_scatter and built-in LSTM were used, with masking for padded atoms in batched execution. For multitask learning, parallel MLP output heads receive the shared Set2Set graph embedding $h_G$.

4. Expressivity Beyond Sum and Average Pooling

The expressivity of Set2Set exceeds that of commutative aggregation functions such as sum and average, which are limited to producing a single fixed function of the multiset of atomwise features and cannot distinguish molecules with coincident set sums or averages. Set2Set’s recurrent LSTM-query mechanism enables dynamic, stepwise attention over subsets of atoms through learned, context-dependent attention weights $\alpha_{i,t}$, which can shift focus to distinct functional groups, conjugated subrings, or heteroatoms across steps.

For properties such as molecular orbital energies, which are sensitive to particular atomic environments, a naive sum- or mean-pooling operation cannot differentiate molecular structure-function relationships when feature sums overlap. Set2Set allows global features whose scale and composition adapt nontrivially with graph size. Empirically, Set2Set aggregation reduced out-of-distribution errors, notably for large molecules, by up to a factor of two compared to sum or average pooling (as detailed in Sec. 4.1 of the source) (Gaul et al., 2022).

5. Practical and Implementation Considerations

Implementation of Set2Set entails several computational and integration factors:

• Each iteration requires an $N_a \times n$ matrix-vector multiplication to form attentional logits, a softmax normalization over $N_a$ atoms, and a weighted sum over $N_a$ vectors, yielding an overall cost of $\mathcal{O}(T N_a n)$.
• For typical molecules (tens to hundreds of atoms) and the employed parameters ($n=128$, $T=3$), the computational overhead remains minor compared to the SchNet backbone.
• Zero-vector initialization is used for LSTM states; only LSTM weights are learned.
• The shifted-softplus activation function, $\mathrm{softplus}(x) = \ln(1 + e^x)$, ensures numerical stability across MLP layers.
• For multitask settings (different computational chemistry levels of theory), a shared Set2Set embedding enables parallel task-specific prediction heads, each with independent loss signals.
• Code implementations leverage maltose v1.0.1 and PyTorch Geometric toolkits, with masking for padded nodes in batch-processed graphs.

6. Empirical Performance and Significance

Set2Set readout, when integrated with SchNet, contributed to improvements in the predictive accuracy for molecular HOMO and LUMO energies, particularly on datasets containing larger and chemically diverse molecules (Gaul et al., 2022). The ability to adaptively focus on relevant substructures enables effective modeling of complex electronic properties. The framework's multitask extension further supported robust learning across data sources with differing theoretical fidelity, bringing model predictions near chemical accuracy limits.

The significance of Set2Set thus lies in its combination of permutation invariance, dynamic attention, and powerful aggregation—properties essential for machine learning tasks involving unordered sets such as atoms in a molecule, where both size and compositional complexity vary across samples.