Few-Shot Graph Learning for Molecular Property Prediction

In plain words

Many molecular properties, such as whether a compound is toxic to a given cell target, have only a handful of labelled examples. Meta-MGNN is built to learn a new property from just a few labelled molecules. It treats each molecule as a graph of atoms and bonds, starts from a network pre-trained on unlabelled molecules, and is trained across many different properties so that it learns to adapt quickly to a new one.

Two extra training signals make use of the structure of each molecule without needing labels: one asks the model to tell real bonds from non-bonds, the other to recover an atom's type from its neighbours. A learned weight lets some training properties count for more than others. On two public multi-property datasets, Tox21 and Sider, the authors report that the model beats the compared methods on average.

TL;DR

Meta-MGNN is a few-shot molecular property predictor: a 5-layer GIN encoder, initialised from a PreGNN checkpoint and trained with MAML-style meta-learning over property tasks. It adds bond-reconstruction and atom-type self-supervision and a softmax task weight. The author-reported average gains over the best listed baseline are 0.84 to 1.87 points of ROC-AUC. The official code differs from the paper in its loss scaling, its evaluation metric and its task embedding, so read the Pitfalls before reproducing or comparing these numbers.

Problem

• Labelled data per property is scarce. The authors report that more than half of the MoleculeNet properties they examined are shared by fewer than 100 molecules §1.
• Deep models usually need far more labels than these properties have §1. The task is to predict new properties, unseen during training, from k labelled examples per class §3.1 §1.
• The paper names three challenges: learning a representation from few labels, exploiting unlabelled molecular structure, and weighting training properties by how useful they are for transfer §1.
• Datasets: Tox21 (7,831 instances, 12 toxicity tasks) and Sider (1,427 instances, 27 system-organ classes of adverse drug reactions) §5.1. Three Tox21 tasks and six Sider tasks are held out for meta-testing §5.1.

Method

Prose overview. The model is a graph neural network that maps a molecule to a vector, followed by a prediction head, trained with meta-learning. Training uses a support set (the few labelled examples used for adaptation) and a query set (examples used to score the adapted model) for each task §4.1.

Molecular graph encoder. Each atom is a node and each bond an edge. Node inputs are the atom number and chirality tag; edge inputs are bond type and bond direction §3.2 tab. 1. The encoder is a graph isomorphism network (GIN) with 5 layers and width 300 §5.4. Each layer aggregates neighbour states and edge features, then applies an update with a non-linearity eq. 1 eq. 2. The paper states that node embeddings are L2-normalised §3.2. The graph embedding is the mean of the final node embeddings eq. 3. The encoder is initialised from the supervised context-prediction GIN released with PreGNN §5.4.

\mathbf{h}^{(l)}_{\mathcal{N}(v)} = \mathrm{Agg}_l\big(\{\mathbf{h}^{(l-1)}_u : u \in \mathcal{N}(v)\},\ \{\mathbf{h}^{(l-1)}_e : e = (v,u)\}\big) eq. 1

\mathbf{h}^{(l)}_v = \sigma\big(\mathbf{W}^{(l)} \cdot \mathrm{Concat}(\mathbf{h}^{(l-1)}_v, \mathbf{h}^{(l)}_{\mathcal{N}(v)})\big) eq. 2

\mathbf{h}_G = \mathrm{Mean}\big(\{\mathbf{h}^{(l)}_v : v \in \mathcal{V}\}\big) eq. 3

Prediction head. A multi-layer perceptron on the graph embedding gives the property logit §4.2.1.

Meta-learning setup. The model is trained with MAML §4.1. For each task, the support set is used to take one gradient step from the shared weights θ to task weights θ′ eq. 4. The paper writes one step but notes that several can be used §4.2. The query set is then scored with θ′, and the shared weights are updated from those query losses eq. 10. Test tasks are never used in training. At test time the support set adapts the model and the query set is scored §4.2.3.

\theta'_\tau = \theta - \alpha \nabla_{\theta} \mathcal{L}_{\mathcal{T}_\tau}(\theta) eq. 4

Losses. The property loss is the mean cross-entropy over the k support molecules eq. 5. The bond-reconstruction loss scores a pair of nodes by the inner product of their embeddings and applies binary cross-entropy to a set of sampled bonds and non-bonds §4.2.1 eq. 6. The atom-type loss takes the 15% of nodes sampled at random, predicts each node's type from the mean embedding of its one-hop neighbours (excluding the node itself) with an MLP, and applies cross-entropy §4.2.1 eq. 7 eq. 8. The joint loss adds these with weights λ1 = λ2 = 0.1 §4.2.1 eq. 9.

\mathcal{L}_{label}(\theta) = -\frac{1}{k}\sum_{i=1}^{k}\mathrm{CrossEntropy}(y_i, \hat{y}_i) eq. 5

\mathcal{L}_{edge}(\theta) = -\frac{1}{|\mathcal{E}_s|}\sum_{e_{uv} \in \mathcal{E}_s}\mathrm{BinaryCrossEntropy}(e_{uv}, \hat{e}_{uv}) eq. 6

\hat{v}_i = \mathrm{MLP}\big(\mathrm{Mean}(\{\mathbf{h}_u : u \in \mathcal{N}_l(v)\})\big) eq. 7

\mathcal{L}_{node}(\theta) = -\frac{1}{|\mathcal{V}_c|}\sum_{i=1}^{|\mathcal{V}_c|}\mathrm{CrossEntropy}(v_i, \hat{v}_i) eq. 8

\mathcal{L}_{\mathcal{T}_\tau}(\theta) = \mathcal{L}_{node}(\theta) + \lambda_1\mathcal{L}_{edge}(\theta) + \lambda_2\mathcal{L}_{label}(\theta) eq. 9

Task-aware attention. Each task gets a weight η(T): a softmax over all training tasks of an MLP applied to the task embedding H_T, which is the mean of the task's molecule embeddings §4.2.2 eq. 11. The outer update weights each task's query loss by η eq. 10. The paper's figure caption says the task embedding is taken from the query set fig. 1.

\theta \leftarrow \theta - \beta\nabla_{\theta}\sum_{\mathcal{T}_\tau \sim p(\mathcal{T})}\eta(\mathcal{T}_\tau)\cdot\mathcal{L}'_{\mathcal{T}_\tau}(\theta'_\tau) eq. 10

\eta(\mathcal{T}_\tau) = \frac{\exp(\mathrm{MLP}(\mathbf{H}_{\mathcal{T}_\tau}))}{\sum_{\mathcal{T}_{\tau'}\in\mathcal{T}}\exp(\mathrm{MLP}(\mathbf{H}_{\mathcal{T}_{\tau'}}))},\quad \mathbf{H}_{\mathcal{T}_\tau} = \mathrm{Mean}\big(\{\mathbf{h}_{\mathcal{T}_\tau,i}\}_{i=1}^{k}\big) eq. 11

Algorithm. The paper's Algorithm 1 starts from PreGNN weights, samples a batch of tasks, and for each task takes the inner step on k support molecules and scores n query molecules. It then computes the task weights eq. 11 and takes the outer step eq. 10 alg. 1. The pseudocode below follows the paper. The code's differences are listed under Paper vs code.

# Meta-MGNN meta-training step  [alg. 1] [§4.2]
# Input: theta (encoder, head, self-supervision heads, task-weight MLP)
# Input: tasks, each with a support set and a query set
# Output: updated theta

def meta_train_step(theta, tasks):
    # Keep one query loss per task  [eq. 4] [eq. 9]
    query_losses = []
    task_embeddings = []
    for task in tasks:
        # Task embedding: mean of molecule embeddings  [§4.2.2] [eq. 11]
        task_embeddings.append(mean_molecule_embedding(theta, task.support))
        # Joint loss on the support set  [eq. 9]
        support_loss = joint_loss(theta, task.support)
        # One inner gradient step with step size alpha  [eq. 4]
        adapted = theta - alpha * grad(support_loss, theta)
        # Joint loss on the query set at the adapted weights  [eq. 9]
        query_losses.append(joint_loss(adapted, task.query))
    # Self-attentive task weights from the task embeddings  [eq. 11]
    weights = softmax(task_weight_mlp(task_embeddings))
    # Outer update of theta with the weighted query losses  [eq. 10]
    outer_loss = sum(weights * stack(query_losses))
    theta = theta - beta * grad(outer_loss, theta)
    return theta

# Joint loss on one set of molecules  [eq. 9]
# graphs: molecular graphs; labels: (B,) with values 0 or 1
# Returns a scalar

def joint_loss(theta, graphs, labels):
    # Node embeddings from the GIN encoder  [§3.2] [eq. 1] [eq. 2]
    nodes = gnn(theta, graphs)
    # Graph embedding: mean over each molecule's nodes  [eq. 3]
    graph_emb = mean_pool(nodes, graphs)
    # Property loss: cross-entropy over the molecules  [eq. 5]
    label_loss = bce_mean(head(graph_emb), labels)
    # Bond reconstruction: 5 bonds and 5 non-bonds  [§4.2.1] [eq. 6]
    pairs, is_bond = sample_bond_pairs(graphs, n_pos=5, n_neg=5)
    scores = sum(nodes[pairs[0]] * nodes[pairs[1]], axis=1)
    edge_loss = bce_mean(scores, is_bond)
    # Atom type: sample 15% of nodes  [§4.2.1]
    centres = sample_nodes(graphs, fraction=0.15)
    # Predict each sampled node's type from its 1-hop neighbours  [eq. 7]
    context = mean_neighbour_embedding(nodes, graphs, centres)
    # Atom-type cross-entropy  [eq. 8]
    node_loss = cross_entropy(mlp(context), atom_type[centres])
    # Weighted sum, lambda1 = lambda2 = 0.1  [eq. 9] [§4.2.1]
    return node_loss + 0.1 * edge_loss + 0.1 * label_loss

# Meta-test on one held-out property  [§4.2.3] [§5.4]
# Input: support (k molecules per class), query (remaining molecules)
# Output: ROC-AUC on the query set

def meta_test(theta, support, query):
    adapted = theta
    # Ten gradient steps on the support set  [§5.4] [eq. 4]
    for step in range(10):
        # Step size alpha: not stated in the paper [code: main.py#L60]
        loss = joint_loss(adapted, support)
        adapted = adapted - alpha * grad(loss, adapted)
    # Score the query set with the adapted weights  [§5.3]
    scores = sigmoid(head(graph_embeddings(adapted, query.graphs)))
    # ROC-AUC computed on the scores, not on thresholded labels  [§5.3]
    return roc_auc(query.labels, scores)

Hyperparameters

Values in the code are defaults or the values set in the __main__ call of main.py. "not stated" means the paper does not give the value.

Results

Author-reported. Setup: 2-way classification with 1 and 5 shots, ROC-AUC per test task, six baselines (GraphSAGE, GCN, MAML, Seq3seq, EGNN, PreGNN) §5.2 §5.3 tab. 2. In the table below, italics mark the best baseline in each row and the last column is Meta-MGNN. The paper's percentage gains are differences in percentage points of the averages.

1-shot

5-shot

Average gain over the best baseline, as stated in the paper: Tox21 +1.04 (1-shot) and +0.84 (5-shot); Sider +1.80 (1-shot) and +1.87 (5-shot) tab. 2 §5.5. Per-task gains range from −0.32 (SR-p53, 1-shot) to +2.64 (Si-T6, 5-shot) tab. 2.

Ablation. Eight variants (M1 to M8) switch on the pre-trained encoder, meta-learning, bond reconstruction, atom-type prediction and task-aware attention tab. 3. The paper reports the ablation only as plots, with no numbers in the text fig. 2. Its qualitative claims: M2 (no pre-training) is worst, M3 beats M1 and M2, each self-supervised term helps, M6 (both terms) is best among M4 to M6, M7 beats M3, and M8 is best in most cases §5.6. The paper's embedding plot for SR-MMP is qualitative only fig. 3 §5.7.

Paper vs code

Official code at commit e72ee6508f24b8eb2e528153c50b44d6e717a53c, linked from the paper's footnote README.md. Differences, each with both citations:

1. Joint-loss weights. Paper weights the node (atom-type) loss by 1 and the label and edge losses by 0.1 eq. 9. Code weights the label loss by 1 and both self-supervised terms by 0.1 meta_model.py meta_model.py meta_model.py.
2. Bond samples. Paper uses 5 positive and 5 negative pairs §4.2.1. Code uses all positive bonds of the batch and the same number of negatives meta_model.py meta_model.py.
3. Bond-loss normalisation. Paper averages the bond BCE over |E_s| = 10 eq. 6. Code adds the mean BCE of the positive pairs and the mean BCE of the negative pairs, then divides by the number of negatives meta_model.py. See Pitfalls.
4. Atom-type head. Paper samples 15% of nodes and predicts each type from the mean embedding of 1-hop neighbours, excluding the node §4.2.1 eq. 7. Code samples batch_size (5) random nodes and applies a linear layer to each node's own embedding, with the atom number as the target and no input masking meta_model.py meta_model.py.
5. Prediction head. Paper uses an MLP §4.2.1. Code uses one linear layer from the graph embedding to a single logit model.py.
6. Input embeddings. Paper concatenates the atom-number and chirality embeddings, and the bond-type and direction embeddings §3.2. Code adds them model.py model.py. Paper's atom-type table lists 118 types tab. 1; code uses 120 embeddings, including mask tokens model.py, and a 119-way head meta_model.py.
7. Node normalisation. Paper L2-normalises node embeddings §3.2. The GIN path does not normalise; the only normalisation is in the GraphSAGE layer model.py model.py.
8. Task embedding. Paper averages all molecule embeddings of the task §4.2.2; the figure caption says the query set fig. 1. Code takes node-level means over the support batches, using a running average that starts at zero (e ← (e + mean)/2), at the pre-step weights meta_model.py meta_model.py. With 10 support molecules and batch size 5, the final vector is one quarter of the first batch's mean plus one half of the second's.
9. Outer loss scale. Paper's outer objective is the η-weighted sum over tasks eq. 10. Code divides it by the number of training tasks meta_model.py.
10. Meta-gradient order. Paper's Algorithm 1 and eq. 10 do not say whether the inner step is differentiated through. Code computes the inner gradient without create_graph and writes the adapted weights back with vector_to_parameters, which assigns .data. The meta-gradient is therefore first-order meta_model.py meta_model.py. This reading is from the code, not from the paper.
11. Optimiser. Algorithm 1 writes a plain gradient step with β alg. 1 eq. 10. Code uses Adam with learning rate meta_lr meta_model.py.
12. Tasks per step. Algorithm 1 samples a batch of tasks in each iteration alg. 1. Code uses all training tasks in every meta-step; the batch sampling is commented out meta_model.py meta_model.py.
13. Sampling frequency. Algorithm 1 samples support and query molecules for each task in each iteration alg. 1. Code builds the support and query loaders once per epoch and reuses them for all 5 meta-steps of that epoch meta_model.py meta_model.py.
14. Metric. Paper reports ROC-AUC §5.3. Code thresholds the predicted probabilities at 0.5 and passes the 0/1 values to roc_auc_score meta_model.py meta_model.py. See Pitfalls.
15. Reported epoch. The paper does not say which epoch is reported. Code records test scores after every epoch and keeps the per-task maximum main.py. The README gives the iteration of the best result for each run README.md.
16. Reruns versus the table. The README reports two runs per setting, with averages of 72.87/73.82 (Sider 1-shot), 76.59/77.14 (Tox21 1-shot), 74.74/74.70 (Sider 5-shot) and 77.99/78.05 (Tox21 5-shot) README.md README.md. The paper's averages are 73.34, 76.87, 74.72 and 78.02 tab. 2.
17. Test tasks. The paper does not name the held-out tasks. Code holds out folders 10 to 12 for Tox21 and 22 to 27 for Sider meta_model.py samples.py. The Tox21 folders are not mapped to column names in the repository.
18. Per-task dataset sizes. The paper gives 7,831 Tox21 instances §5.1. Code's per-task label counts range from 5,810 to 7,265 samples.py. Sider's per-task count is 1,427 in every task samples.py.
19. Ablation M1. The paper does not say how the pre-trained-only variant is adapted to new properties §5.6.
20. Baselines. The repository has no code for the six baselines, so their numbers are not reproducible from it.
21. Dataset split script. Original_datasets/splitdata.py sets name = 'tox21' but loops over 617 columns splitdata.py splitdata.py. It cannot run on the 12-task Tox21 file as shipped, so the committed folders were produced some other way.

Pitfalls

• The reported "best" is a test-set maximum. The training script evaluates the test tasks after every epoch and keeps the best score per task main.py. The paper does not state this selection rule. Any comparison should use a validation split chosen before testing.
• The metric is not ROC-AUC on scores. Predictions are turned into hard 0/1 labels before roc_auc_score, so the value is the area under a two-point curve, which equals balanced accuracy meta_model.py. Implement the metric on the sigmoid scores, as §5.3 describes.
• The bond loss is scaled by the number of bonds. nn.BCEWithLogitsLoss() already averages, and the code sums the two averages and divides by the number of negatives, which equals the number of bonds in the batch meta_model.py meta_model.py meta_model.py. The same term appears in the query and test passes meta_model.py meta_model.py. As written, the bond term shrinks as the batch contains more bonds, so it does not match the paper's averaged term eq. 6. Ablation claims about bond reconstruction should not be read off the paper's formula.
• Label losses are divided by batch size twice. Each label loss is already a mean and is then divided by the number of graphs in the batch, and the per-batch values are summed over batches meta_model.py meta_model.py meta_model.py. The scale depends on how many batches there are. With the defaults, a 10-molecule support set gives 2 batches and a 128-molecule query set gives 26, so the support and query label terms are on different scales.
• Signs in the equations. Equations 5, 6 and 8 each start with a minus sign in front of a cross-entropy, which would make the loss negative. The code uses positive losses meta_model.py meta_model.py meta_model.py. Implement the positive form.
• Negative-bond sampling. Negatives are drawn only from pairs (num, num_back) with num_back >= num, and existing bonds are excluded only in the direction stored in all_edge meta_model.py. Bonds are stored as (begin, end) in RDKit order loader.py. A true bond whose begin atom has the larger index can therefore be sampled as a negative. The card does not measure how often this happens.
• Cost of negative sampling. The negatives are built with a Python double loop over all node pairs in each batch meta_model.py, and this runs for every support and query batch.
• Task weights are trained on the loss they weight. The task-weight MLP receives gradient from the weighted query loss meta_model.py. That gradient favours tasks with low query loss, not a chosen notion of importance. This is a reading of the code, not something the paper discusses.
• First-order meta-gradient. The inner gradient is not differentiated through (see Paper vs code, item 10). Results may differ from a second-order MAML implementation.
• Test-time adaptation is noisy. Each test task takes 10 full-batch steps on its support set, and the query set is all remaining molecules meta_model.py samples.py. For Sider, that is about 1,400 query molecules per test task, scored every epoch.
• Small margins and unstable runs. The README calls the meta-learning results unstable on some properties README.md. Its two Sider 1-shot averages differ by about 0.95 points (72.87 and 73.82), which is the same order as the paper's gains of 0.84 to 1.87 points README.md tab. 2.
• Table versus text. The paper says Meta-MGNN is best on every task §5.5. The table shows two exceptions: EGNN on 1-shot SR-p53 (78.03 against 77.71) and Seq3seq on 5-shot SR-MMP (80.40 against 80.26) tab. 2.
• Random seeds. Torch and numpy are seeded, but Python's random module is not, and it drives all support, query and negative sampling main.py samples.py meta_model.py. Runs are not reproducible as shipped.
• Pre-trained checkpoint on CPU. from_pretrained calls torch.load with map_location='cuda:0' model.py, so a CPU-only machine needs a change before it can load the checkpoint.
• Memory and time. Each epoch runs 5 meta-steps over all training tasks, each with a 128-molecule query pass per task, plus a test pass. The default is 2000 epochs main.py meta_model.py. The paper gives no run time or hardware.

Lineage

• Builds on MAML (Finn et al., 2017) for the meta-learning loop §4.1.
• Builds on PreGNN (Hu et al., 2020) for the pre-trained GIN initialisation §3.2 §5.4.
• Uses the GIN encoder of Xu et al. (2019) §3.2 model.py.
• Uses self-attentive weighting from Lin et al. (2017) for the task weights §4.2.2.
• Compared against GraphSAGE, GCN, MAML, Seq3seq, EGNN and PreGNN §5.2.
• Superseded: not stated.