PACIA: Parameter-Efficient Adapter for Few-Shot Molecular Property Prediction

Few-shot learning (FSL) aims to generalize to a task with a few labeled samples Wang et al. (2020). In terms of adaptation mechanism, existing FSL methods can be classified into three main categories: (i) gradient-based approaches Finn et al. (2017); Grant et al. (2018) learn a model which can be generalized to new task by gradient descents, (ii) metric-based approaches Vinyals et al. (2016); Snell et al. (2017) learn to embed samples into a space where similar and dissimilar samples can be easily discriminated by a distance function, and (iii) amortization-based approaches Requeima et al. (2019); Lin et al. (2021); Przewiezlikowski et al. (2022) use hypernetworks to map the labeled samples in the task to a few parameters to adjust the main networks to be task-specific. Recent works Requeima et al. (2019) found that amortization-based approaches can reduce the risk of overfitting compared with gradient-based approaches. They also have faster inference speed as the adapted parameters are generated by a single forward pass without taking optimization steps. Besides, the main networks can approximate various functions in addition to distance-based ones.

Hypernetworks Ha et al. (2017) refer to neural networks which learn to generate parameters of the main network which handles the target tasks. Hypernetworks have been successfully used in various applications like cold-start recommendation Lin et al. (2021) and image classification Przewiezlikowski et al. (2022). Designing appropriate hypernetworks is challenging, requiring domain knowledge to decide what information to be fed into hypernetworks, how to adapt the main network, and what is the appropriate architecture of hypernetworks. For general GNNs, hypernetworks are developed to modulate weight matrix in aggregation function in message passing Brockschmidt (2020), or to facilitate node-specific message passing Nachmani and Wolf (2020). In contrast to them, we particularly consider designing parameter-efficient modulators for GNNs used in encoder-predictor framework for few-shot MPP.

In a few-shot MPP task $\mathcal{T}_{\tau}$ with respect to a specific property, each sample $\mathcal{X}_{{\tau},i}$ is a molecular graph and its label $y_{{\tau},i}\in\{0,1\}$ records whether the molecule is active or inactive on the target property. Only a few labeled samples are available in $\mathcal{T}_{\tau}$. Following earlier works Altae-Tran et al. (2017); Stanley et al. (2021); Chen et al. (2022); Schimunek et al. (2023), we model a $\mathcal{T}_{\tau}$ as a $2$-way classification task $\mathcal{T}_{\tau}$, associating with a support set $\mathcal{S}_{\tau}=\{(\mathcal{X}_{{\tau},s},y_{{\tau},s})\}_{s=1}^{N_{\tau}}$ containing labeled samples from active/inactive class, and a query set $\mathcal{Q}_{\tau}=\{({\mathcal{X}}_{{\tau},q},{y}_{{\tau},q})\}_{q=1}^{M_{\tau}}$ containing $M_{\tau}$ samples whose labels are only used for evaluation. We consider both (i) balanced support sets, i.e., $\mathcal{S}_{\tau}$ contains $\frac{N_{\tau}}{2}$ samples per class which is consistent with the standard $N$-way $K$-shot FSL setting Altae-Tran et al. (2017), and (ii) imbalanced support sets which exist in real-world applications Stanley et al. (2021). Our target is to learn a model from a set of tasks $\{\mathcal{T}_{\tau}\}_{j=1}^{N}$ that can generalize to new task given the few-shot support set. Specifically, the target properties are different across tasks.

In the past, molecules are encoded with certain properties (fingerprint vectors Rogers and Hahn (2010)) and fed to deep networks for prediction Unterthiner et al. (2014); Ma et al. (2015). While recently, GNNs are popularly taken as molecular encoders Li et al. (2018); Yang et al. (2019); Xiong et al. (2019); Hu et al. (2019) due to their superior performance on learning from of topological data. In either case, existing works can be summarized within an encoder-predictor framework.

Consider a molecular graph $\mathcal{X}=\{\mathcal{V},\mathcal{E}\}$ with node feature $\mathbf{h}_{v}$ for each atom $v\in\mathcal{V}$ and edge feature $\mathbf{b}_{vu}$ for each chemical bond $e_{vu}\in\mathcal{E}$ between atoms $v,u$. A GNN encoder maps $\mathcal{X}$ to molecular representation $\mathbf{r}$ which is a fixed-length vector. At the $l$th layer, GNN updates atom embedding $\mathbf{h}_{v}^{l}$ of $v$ as

where $\mathcal{H}(v)$ contains neighbors of $v$. AGG($\cdot$) and UPD($\cdot$) are aggregation and updating functions respectively. After $L$ layers, the query-level representation $\mathbf{r}$ for $\mathcal{X}$ is obtained as

where $\texttt{READOUT}(\cdot)$ function aggregates atom embeddings.

Then, a predictor $f(\cdot)$ assigns label for a query molecule $\mathcal{X}_{{\tau},q}$ given support molecules in $\mathcal{S}_{\tau}$:

The specific choice of $f(\cdot)$ is diverse, e.g., pair-wise similarity Altae-Tran et al. (2017), multi-layer perceptron (MLP) Guo et al. (2021); Wang et al. (2021) and Mahalanobis distance Stanley et al. (2021). Recently, GNN-based predictor which operates on relation graphs of molecules is found to effectively compensate for the lack of supervised information Wang et al. (2021). In particular, molecular representations are refined on relation graphs such that the similar molecules cluster closer. Initialize molecular representations as the output of the encoder, i.e., $\mathbf{h}_{\tau,i}^{0}=\mathbf{r}_{\tau,i}$. Denote the set of $N_{\tau}+1$ molecules as $\mathcal{R}_{{\tau},q}=({\mathcal{X}}_{{\tau},q},{y}_{{\tau},q})\cup\mathcal{S}_{\tau}$, which contains all information to make prediction for the query molecule ${\mathcal{X}}_{{\tau},q}$. The relation graph works by recurrently estimating the adjacency matrix and updating the molecular representations. At the $l$th layer, each element $a_{ij}^{l}$ in the adjacent matrix $\mathbf{A}_{{\tau},q}^{l}$ of the relation graph is learned to represent pair-wise similarities between any two molecules in $\mathcal{R}_{{\tau},q}$:

Then, each molecular representation is refined as

After $L$ layers of refinement, $\mathbf{h}_{{\tau},q}^{L}$ and $\mathbf{h}_{{\tau},s}^{L}$ (in place of $\mathbf{r}_{\tau,q}$ and $\mathbf{r}_{\tau,s}$) are fed to (3) to obtain final prediction $\hat{\bm{y}}_{{\tau},q}$ for $\mathcal{X}_{\tau}$.

To adapt GNN’s parameter-efficiently, we design a GNN adapter to modulate the node embedding and propagation depth, which are essential in message passing process.

Denote the collection of all model parameters in main network((1)-(5)) and hypernetwork ((9)-(11)) as $\bm{\Theta}$, excluding adaptive parameters. Our objective takes the form:

$\mathcal{L}_{\tau}$ is the loss in task $\mathcal{T}_{\tau}$, ${\bm{y}}_{{\tau},q}$ is one-hot ground-truth label vector and $\hat{\bm{y}}_{{\tau},q}$ is prediction obtained by (3).

Note that $\bm{\Theta}$ is shared across all tasks. While the adaptive parameter $\{\bm{\gamma}^{l}\}_{l=1}^{L}$ in (6) and $\{p^{l}\}_{l=1}^{L}$ in (7) are generated by hypernetworks. The size of adaptive parameter is far smaller than the main network. This realizes parameter-efficient adaptation and mitigates the risk of overfitting.

Algorithm  1¹¹1In Algorithm 1 and 4,“*” (resp. “+”) indicates the step is executed by the main network (resp. hypernetwork). summarizes the training procedure of PACIA. As mentioned above, our unified GNN adapter can simultaneously modulate the node embedding and propagation depth, and be cascaded to adapt both encoder and predictor. During training, molecular graph $\mathcal{X}_{{\tau},i}$ is first processed by encoder (line 4-9). At each layer, adaptive parameters $[\bm{\gamma}_{\tau}^{l},p^{l}_{\tau}]$ are obtained by (10) (line 5). Then, (6) modulates all atom embeddings $\mathbf{h}_{v}^{l}$ (line 6). After $L_{\text{enc}}$ layers of message passing (1), (8) is applied before (2), to get property-adaptive molecular representations and initialize node embeddings $\mathbf{h}_{{\tau},i}^{0}=\mathbf{r}_{\tau,i}$ in relation graph (line 9). Then in predictor, at each layer of GNN on relation graph, adaptive parameters $[\bm{\gamma}_{\tau,q}^{l},p^{l}_{\tau,q}]$ are obtained with (11) (line 13) and (6) modulates all node embeddings $\mathbf{h}_{\tau,i}^{l}$ (line 14). After $L_{\text{rel}}$ layers of message passing by (4)(5), (8) is applied (line 17). The final prediction $\hat{\mathbf{y}}_{\tau,q}$ is obtained by (3).

Testing procedure is provided in Appendix B. The process is similar. A noteworthy difference is the propagation depth is adapted by selecting the layer with maximal plausibility:

Only $\mathbf{h}^{l^{\prime}}$, rather than (8), is fed forward to the next module.

From the perspective of hypernetworks, the usage of hypernetworks for encoder is related to GNN-FiLM Brockschmidt (2020), which considers a GNN as main network. It builds hypernetworks with target node as input to generate parameters of FiLM layers, to equip different nodes with different aggregation functions in the GNN. What and how to adapt are similar to ours, but it is different that the input of our hypernetworks for encoder is $\mathcal{S}_{\tau}$ and how we encode a set of labeled graphs.

As for few-shot learning, some recent studies Requeima et al. (2019); Lin et al. (2021) use hypernetworks to transform the support set into parameters that modulate the main network. The functionality of their hypernetworks is akin to that of our hypernetwork for the encoder. However, there is a distinct difference in the architecture of main networks: while their approaches employ convolutional neural networks and MLPs, our method uniquely modulates the message passing process of GNN. Furthermore, the application of hypernetworks for the predictor, aimed at adjusting the model architecture based on a query sample (without label information) and a set of labeled samples, has not yet been explored in the literature. A detailed comparison of PACIA w.r.t. existing few-shot MPP works is in Appendix C.

We use Tox21 National Center for Advancing Translational Sciences (2017), SIDER Kuhn et al. (2016), MUV Rohrer and Baumann (2009) and ToxCast Richard et al. (2016) from MoleculeNet Wu et al. (2018), which are commonly used to evaluate the performance on few-shot MPP Altae-Tran et al. (2017); Wang et al. (2021). We adopt the public data split provided by Wang et al. (2021). The support sets are balanced, each of them contains $K$ labeled molecules per class, where $K=1$ and $K=10$ are considered. The performance is evaluated by ROC-AUC calculated on the query set of each meta-testing task and averaged across all meta-testing tasks.

We compare PACIA with the following baselines: 1) single-task method GNN-ST Gilmer et al. (2017); 2) multi-task pretraining method GNN-MT Corso et al. (2020); Gilmer et al. (2017); 3) self-supervised pretraining method MAT Maziarka et al. (2020); 4) meta-learning methods, including Siamese Koch et al. (2015), ProtoNet Snell et al. (2017), MAML Finn et al. (2017), EGNN Kim et al. (2019); and 5) methods proposed for few-shot MPP, including IterRefLSTM Altae-Tran et al. (2017), PAR Wang et al. (2021) and ADKF-IFT Chen et al. (2022). Note that MHNfs Schimunek et al. (2023) is not included as it uses additional reference molecules from external datasets, which leads to unfair comparison. GS-META Zhuang et al. (2023) has not been compared since that approach requires multiple properties of each molecule, which is not applicable when a molecule is only evaluated w.r.t. one property. Following earlier works Guo et al. (2021); Wang et al. (2021), we use GIN Xu et al. (2019) as encoder, which is trained from scratch.

Table 1 shows the results. Results of Siamese and IterRefLSTM are copied from Altae-Tran et al. (2017) as their codes are unavailable, and their results on ToxCast are unknown. GNN-FiLM is a general GNN whose target is not few-shot MPP, which explains its bad performance. PACIA obtains the highest ROC-AUC scores on all cases except the 10-shot case on MUV, where ADKF-IFT outperforms the others by a large margin. This can be a special case where ADKF-IFT works well but may not be generalizable. Moreover, depending on the number of local-update steps of ADKF-IFT, PACIA is about 5 times faster than ADKF-IFT (both meta-training and inference time is about 1/5). In terms of average performance, PACIA significantly outperforms the second-best method ADKF-IFT by 3.25%. We also provide results obtained with a pretrained encoder in Appendix D.3. Similar observations can be made: our PACIA with pretrained encoder (Pre-PACIA) performs the best, and its performance gain is more significant when fewer labeled samples are provided. Further, results in Appendix D.4 shows the performance comparison between PACIA and a fine-tuned GNN with varying support set size. This shows that PACIA nicely achieves its goal: handling few-shot MPP problem in a parameter-efficient way.

We also use FS-Mol Stanley et al. (2021), a new benchmark consisting of a large number of diverse tasks for model pretraining and a set of few-shot tasks with imbalanced classes. We adopt the public data split Stanley et al. (2021). Each support set contains 64 labeled molecules, and can be imbalanced where the number of labeled molecules from active and inactive is not equal. All remaining molecules in the task form the query set. Testing tasks are divided into categories with support size 16 Schimunek et al. (2023), which is close to real-world scenario. The performance is evaluated by $\Delta$AUPRC (change in area under the precision-recall curve) w.r.t. a random classifier Stanley et al. (2021), averaged across meta-testing tasks.

We use the same baselines that were applied to MoleculeNet. Table 2 shows the results. We find that PACIA performs the best. Besides, the time-efficiency of PACIA is much higher since its adaptation only needs a single forward pass. While IterRefLSTM and ADKF-IFT take multiple local-update steps, they are much slower to generalize.

We consider various variants of PACIA, including (i) fine-tuning: using the same model structure and fine-tuning all parameters to adapt to each property without hypernetworks; (ii) w/o T: removing task-level adaptation, thus the GNN encoder will not be adapted by hypernetworks w.r.t. each property; and (iii) w/o Q: removing query-level adaptation, such that all molecules are processed by the same predictor.

Figure 2(a) provides performance comparison on Tox21. As shown, the performance gain of PACIA over “w/o Q” shows the necessity of query-level adaptation. The gap between PACIA and “w/o T” indicates the effect of adapting the model to be task-specific. One can also notice that without query-level adaptation, “w/o Q” still obtains better performance than gradient-based baselines like PAR, which indicates the advantage of designing the amortization-based hypernetwork. The poor performance of “fine-tuning” is possibly because of the overfitting caused by updating all parameters with only a few samples. In sum, every component of PACIA is important for achieving good performance.

Now that the effectiveness of task-level and query-level adaptation are validated, we further investigate modulation functions, i.e, modulating node embedding (N), modulating propagation depth (D), modulating both (ND), for encoder and predictor. There are $3\times 3$ combinations, whose performance is reported in Figure 2(b). We find that only modulating node embedding in encoder while only modulating propagation depth in predictor obtains the best performance. As the GIN encoder has highly non-linearity across layers, truncation would lead to non-explainability and somehow perturb the black-box. While the operation of relation graph in predictor updates node embedding in a linear way (5), adapting the propagation depth is harmonious with message passing process.

Figure 3 shows the t-SNE visualization Van der Maaten and Hinton (2008) of molecular representations learned on a 10-shot support set and a query molecule with ground-truth label “active” in task SR-p53 from Tox21. As shown, molecular representations obtained without adaptation (Figure 3(a)) are mixed up, since the encoder has not been adapted to the target property of the task. Molecular representations being processed by our property-adaptive encoder (Figure 3(b)) becomes more distinguishable, indicating that adapting molecular representation in task-level takes effect. Molecular representations in Figure 3(c) and Figure 3(d) form clear clusters as we encourage similar molecules to be connected during relation graph refinement by (5). The difference is that molecular representations in Figure 3(c) are refined by the best propagation depth number for all tasks in 10-shot case, while molecular representations in Figure 3(d) are refined by 4 depths of propagation which are selected for the specific query molecule. As shown, we can conclude that our molecular-adaptive refinement steps help better separate molecules of different classes. Ablation study of configurations of hypernetworks is in Appendix D.5.