Summary
This paper studies few-shot molecular property prediction with an in-context learning formulation. A method called Context Aware Molecule Prediction (CAMP) is proposed, which employs a transformer encoder to learn representations for the set of (molecule, property) pairs from both support and query sets and an MLP head to predict the property of the query molecule based on its transformed representation. The proposed method is evaluated on the FS-Mol and BACE binary classification benchmarks. Ablation studies are conducted to explore the effects of positional embedding and (molecule, property) embedding concatenation.
Strengths
1. The formulation of in-context learning for molecular property prediction is interesting as the transformer encoder learns to extracts dynamic representations for the support (molecule, property) pairs alongside the query molecules, which is more powerful than the static representations in metric-based methods and more efficient than gradient-based methods.
2. The analysis and visualization of the embeddings and attention weights of CAMP are helpful for understanding the properties of the learned representations.
Weaknesses
1. My main concern is the experiments. In the abstract, the paper claims that
> On the FS-Mol and BACE molecular property prediction benchmarks, we find this method surpasses the performance of recent meta-learning algorithms at small support sizes and is competitive with the best methods at large support sizes.
However, the above claim is not an accurate description of the results in the paper, because the experimental setup in Section 6.1 is different from the setup of the FS-Mol benchmark. Specifically, the results in Section 6.1 are obtained with GNN features while the original FS-Mol benchmark and all subsequent works consider GNN + ECFP + PhysChem descriptors. I suspect this is why the performance of ADKF-IFT reported in Section 6.1 is significantly worse than that reported in the original paper [1]. The authors justify the removal of the ECFP and PhysChem descriptors by the observation that computing these features significantly slows down the inference speed, which does not seem to make sense because these features should be cheap to compute (actually much faster than computing the GNN features). Also, as far as I know, ECFP in particular is still the No.1 choice in industry. Therefore, the comparison would be more meaningful and useful if all methods are run with the original setup from the FS-Mol paper [2].
2. While the in-context learning formulation for molecular property prediction is interesting, the technical contribution of this work is quite limited as it is just an application of the transformer encoder to the few-shot molecular property prediction problem. The idea of concatenating (molecule, property) embedding (i.e., concatenating (x, y) embedding) and the permutation invariance property of transformers have already been thoroughly investigated in the set transformer paper [3].
3. There are a few questions that need to be clarified. See the Questions section below.
[1] Wenlin Chen, et al. "Meta-learning adaptive deep kernel gaussian processes for molecular property prediction." ICLR 2023.
[2] Megan Stanley, et al. "Fs-mol: A few-shot learning dataset of molecules." NeurIPS 2021 (Datasets and Benchmarks Track).
[3] Juho Lee, et al. "Set transformer: A framework for attention-based permutation-invariant neural networks." ICML 2019.
Questions
1. For the visualizations of the self-attention weights in Figure 2, the vertical striation pattern is interesting. Can the authors explain why there is such a vertical striation pattern in the self-attention weights? Intuitively, each molecule should attend to the ones that have the same property label as its own label, but it does not seem to be the case in the visualization.
2. It is sensible that ADKF-IFT is slower than CAMP due to the inner loop optimization of the GP parameters. However, can the authors explain why CNP and ProtoNet are much slower than CAMP as reported in Table 3? This does not seem to make sense because none of these methods require parameter updates at inference time. In theory, the inference time cost of CAMP is $O(N^2)$ due to the self-attention computation, which is the same as computing the Mahalanobis distance for each pair of molecules (also $O(N^2)$) in ProtoNet. CNP should be even faster at inference time as it only costs $O(N)$.
Rating
3: reject, not good enough
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.