SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuankai Luo
luoyk@buaa.edu.cn
Lei Shi
leishi@buaa.edu.cn
Veronika Thost
MIT-IBM Watson AI Lab,
veronika.thost@ibm.com
Self-supervised learning (SSL) has great potential for molecular representation learning given the complexity of molecular graphs, the large amounts of unlabelled data available, the considerable cost of obtaining labels experimentally, and the hence often only small training datasets. The importance of the topic is reflected in the variety of paradigms and architectures that have been investigated recently, most focus on designing views for contrastive learning. In this paper, we study SSL based on persistent homology (PH), a mathematical tool for modeling topological features of data that persist across multiple scales. It has several unique features which particularly suit SSL, naturally offering: different views of the data, stability in terms of distance preservation, and the opportunity to flexibly incorporate domain knowledge. We (1) investigate an autoencoder, which shows the general representational power of PH, and (2) propose a contrastive loss that complements existing approaches. We rigorously evaluate our approach for molecular property prediction and demonstrate its particular features in improving the embedding space: after SSL, the representations are better and offer considerably more predictive power than the baselines over different probing tasks; our loss increases baseline performance, sometimes largely; and we often obtain substantial improvements over very small datasets, a common scenario in practice.
Self-supervised learning (SSL) has great potential for molecular representation learning given the complexity of molecules, the large amounts of unlabelled data available, the considerable cost of obtaining labels experimentally, and the resulting often small datasets. The importance of the topic is reflected in the variety of paradigms and architectures that are investigated [Xia et al., 2023a].
Most existing approaches use contrastive learning (CL) as proposed in [You et al., 2020]. CL aims at learning an embedding space by comparing training samples and encouraging representations from positive pairs of samples to be close in the embedding space while representations from negative pairs are pushed away from each other. However, usually, each given molecule is considered as its own class, that is, a positive sample pair consists of two different views of the same molecule; and all other samples in a batch are used as the negative pairs during training. We observe that this represents a very coarse-grained comparison basically ignoring all commonalities and differences between the given molecules. Therefore also current efforts in molecular CL are put in constructing views that capture the possible relations between molecules: GraphCL [You et al., 2020] proposes four
Figure 1: Topological fingerprints are constructed using a filter (e.g., atom number) and recording resulting topological structures in a PD which is then vectorized.
simple augmentations (e.g., drop nodes), but fix two per dataset; JOAO [You et al., 2021] extends the former by automating the augmentation selection; GraphLoG [Xu et al., 2021] optimizes in a local neighborhood of the embeddings, based on similarity, and globally, using prototypes based on semantics; and SimGRACE [Xia et al., 2022] creates views using a second encoder, a perturbed version of the molecule's ones; also most recent works follow this paradigm [Wu et al., 2023].
Quality of Embedding Spaces is Underinvestigated. Although differences in performance are often observed, they seem marginal and are not yet fully comprehended. With a few exceptions, the models are usually evaluated over the MoleculeNet benchmark [Wu et al., 2018a] only, a set of admittedly rather diverse downstream tasks. While those certainly provide insights into model performance, recent evaluation papers have pointed out that the picture may be very different when the models are evaluated in more detail [Sun et al., 2022, Wang et al., 2022a, Akhondzadeh et al., 2023, Deng et al., 2022]. In particular, [Akhondzadeh et al., 2023] propose to use linear probing to analyse the actual power of the representations instead of a few, specific downstream tasks. This type of evaluation is also common in other areas of DL [Chen et al., 2020]. [Deng et al., 2022] consider other, smaller datasets, and compare to established machine learning (ML) approaches.
We study molecular SSL based on persistent homology (PH). PH is a mathematical tool for modeling topological features of data that persist across multiple scales. In a nutshell, the molecule graph can be constructed sequentially using a custom filter, such that atom nodes and the corresponding bond edges only appear once they meet a given criterion (e.g., based on atomic mass); see Figure 1. PH then captures this sequence in a persistence diagram (PD), and the area has developed various methods to translate these diagrams into topological fingerprints, which can be used for ML [Ali et al., 2022].
Topological fingerprints have been used for supervised ML over molecules by chemists [Krishnapriyan et al., 2021, Demir et al., 2022] and recent work in AI has shown promising results compared to the commonly used ECFP embeddings [Rogers and Hahn, 2010]. Moreover, many of these fingerprints have been shown to be stable, in the sense that distances between the PDs are reflected in the distances between the corresponding fingerprints. We point out that the unique features of topological fingerprints particularly suit molecular SSL, naturally offering: different views of the data (i.e., by switching the filter function), stability in terms of distances, and the opportunity to flexibly incorporate domain knowledge.
Figure 2: TDL may yield great improvements in the low-data scenario; ClinTox.
Our Idea. We consider the topological fingerprints of the molecules in the pre-training data as views and exploit their stability to model the distances between the given molecules (i.e., instead of between a molecule and/or its views). In particular, we use them for fine-grained supervision in SSL.
- We propose a topological distance contrastive loss (TDL) which, as outlined above, shapes the embedding space by providing supervision for the relationships between molecules, especially the ones in the data. TDL is complementary and can be flexibly and efficiently applied to complement any CL approach. We extensively evaluate it combining TDL with four established CL models.
- We also consider alternative paradigms for comparison: we study a straightforward topological fingerprints autoencoder (TAE).
- Our evaluation particularly focuses on the potential in improving the embedding space and we demonstrate the gain in representation power in detail empirically (see Section 4.1): TDL en
ables the models to learn, in a sense, calibrated distances; considerably improves linear probing performance; and the fine-grained supervision may mitigate deficiencies of individual models.
- Over downstream data (see Section 4.2) the performance increases depend more on the model and data. Notably, TDL is able to considerably improve GraphCL and make it competitive with SOTA. Moreover, TDL strongly improves various models in the low-data scenario (see Figure 2).
Our implementation is available at https://github.com/LUOyk1999/Molecular-homology.
Graph Homology. Molecules are graphs  $G = (V, E)$  with nodes  $V$ , the atoms, and bond edges  $E$ . In algebraic topology, graph homology considers such a graph  $G$  as a topological space. We focus on simplices: every node is a 0-simplex, and every edge is a 1-simplex. In the context of graphs, we obtain a 1-dimensional simplicial complex  $X = V \cup E$  in an easy way, by considering the simplices induced by  $G$ ; the dimension is determined by the maximal dimension of the contained simplices.
Persistent Homology (PH). We introduce the most important concepts and, for more details, refer to [Dey and Wang, 2022, Edelsbrunner and Harer, 2022]. Persistent homology is a mathematical tool for modeling topological features of data that persist across multiple scales, comparable to different resolutions. These features are captured in persistent diagrams which, in turn, can be vectorized in fingerprints. We outline the process simplified below and in Figure 1; see Appendix A for details.
First, the goal is to construct a nested sequence of subgraphs  $G_{1} \subseteq \ldots \subseteq G_{N} = G$  ( $1 \leq i \leq N$ ). As described above, these graphs can be considered as simplicial complexes, hence we have simplices which we can record in a persistence diagram. To this end, we consider one of the most common types of filtration methods: sublevel/superlevel filtrations. A sublevel filtration is a function  $f: X \to \mathbb{R}$  over all simplices. A simple such function  $f$  can be a node-valued function (e.g., map atom nodes to their atomic number) that is expanded to the edges as  $f(u, v) = \max(f(u), f(v))$ . Denote by  $X_{a}$  the sublevel set of  $X$ , consisting of simplices whose filtration function values  $\leq f(a)$ ,  $X_{a} = \{x \in X | f(x) \leq f(a)\}$ . As the threshold value  $a$  increases from  $\min_{v \in V} f(v)$  to  $\max_{v \in V} f(v)$ , let  $G_{a}$  be the subgraph of  $G$  induced by  $X_{a}$ ; i.e.,  $G_{a} = (V_{a}, E_{a})$  where  $V_{a} = \{v \in X_{a}\}$  and  $E_{a} = \{e_{rs} \in X_{a}\}$ . This process yields a nested sequence of subgraphs  $G_{1} \subseteq G_{2} \subseteq \ldots \subseteq G_{N} = G$ . As  $X_{a}$  grows to  $X$ , new topological structures gradually appear (born) and disappear (die).
Second, a persistence diagram (PD) is obtained as follows. For each topological structure  $\sigma$ , PH records its first appearance and its first disappearance in the filtration sequence. And this is represented by a unique pair  $(b_{\sigma}, d_{\sigma})$ , where  $1 \leq b_{\sigma} \leq d_{\sigma} \leq N$ . We call  $b_{\sigma}$  the birth time of  $\sigma$ ,  $d_{\sigma}$  the death time of  $\sigma$  and  $d_{\sigma} - b_{\sigma}$  the persistence of  $\sigma$ . PH records all these birth and death times of the topological structures in persistence diagram  $PD(G) = \{(b_{\sigma}, d_{\sigma}) | \sigma \in H_k(G_i)\}$ , where  $H_k(G_i)$  denotes the  $k$ -th homology group of  $G_i$ , and in practice,  $k$  typically takes values 0 or 1. This step is rather standard and corresponding software is available [Otter et al., 2017].
PH in ML. The vectorization of PDs for making them applicable in ML has been studied extensively [Ali et al., 2022], and proposals range from simple statistical descriptions to more complex persistence images (PIs) [Adams et al., 2017]. In a nutshell, for PIs, the PD is considered as a 2D surface, transformed using a Gaussian basis function, and finally discretized into a vector. Furthermore, the Euclidean distance between PIs is stable with respect to the 1-Wasserstein distance between PDs [Adams et al., 2017], which essentially means that the former is bounded by a constant multiple of the latter. We focus on PIs based on the promising results reported in supervised settings [Demir et al., 2022, Krishnapriyan et al., 2021]. Our study was inspired by the ToDD framework, applying pretrained vision transformers to custom 2D topological fingerprints [Demir et al., 2022]. Interestingly, they use the distances between the transformer's molecule representations to construct a suitable dataset for subsequent supervised training using triplet loss. Our focus is on SSL and we apply the ToDD fingerprints as more complex, expert fingerprints in comparison to the ones based on atomic mass only. Recently, various other approaches of integrating PH into ML are explored, but these are only coarsely related to our work (e.g., [Horn et al., 2022, Yan et al., 2022]).
Molecular SSL. Since the foundational work of [Hu* et al., 2020], who have proposed several effective methods, such as the node context prediction task ContextPred, the area is advancing at great pace [Xia et al., 2023a, Xie et al., 2022]. Our proposal falls into the category of contrastive learning as introduced in Section 1. Our study focuses on the potential of PH to complement existing
models such as [You et al., 2020] and follow-up approaches. There are other related CL approaches to which we do not aim to compare to directly; e.g., works including 3D geometry [Liu et al., 2022a, Stärk et al., 2022] or considerably scaling up the pre-training data [Ross et al., 2022, Zhou et al., 2023]. In contrast, our focus is on improvement through exploiting unused facets of the data.
SSL based on Distances. We found only few works that explicitly incorporate distances into SSL. With regard to graphs, [Kim et al., 2022] exploit the graph edit distance between a graph and views to similarly represent the actual distance in the embedding space inside the loss function (i.e., vs. invariance to the view transformation). This distance computation is feasible since it can be easily obtained based on the transformation. Our work focuses on exploring in how far distances in terms of PH can be applied towards the same goal; moreover, this makes it feasible to model the distances between given samples. [Wang et al., 2022b] consider a loss very similar to our proposal based on the more coarse-grained ECFPs, instead of PIs, but they focus on chemistry aspects instead of SSL more generally. Also [Zha et al., 2022] model distances between given samples in the context of CL in a similar way, but in the context of a supervised scenario, where the distances are the differences between given regression labels; note that they also list other coarser related works. [Taghanaki et al., 2021] apply custom distances inside a triplet loss but similarly exploit label information, in order to select the positive and negative samples. Beyond that, distances have been applied to obtain hard negatives [Zhang and Re, Demir et al., 2022], and are exploited in various other ways rather different from our method. For instance, several recent approaches aim to structure the embedding space by exploiting correlations between samples which are, in turn, obtained using nearest neighbors [Caron et al., 2020, Dwibedi et al., 2021, Ge et al., 2023]. Also methods considering equivariance between view transformations implicitly model distances [Chuang et al., 2022, Devillers and Lefort, 2023].
Others. While our focus on CL and distances hints at a close relationship to deep metric learning [KAYA and BILGE, 2019], these works usually do not have explicit distances for supervision but, for instance, exploit labels. Observe that this indicates the unique nature of the distances PH provides us with. Lastly, SSL more generally [Rethmeier and Augenstein, 2023] has naturally been inspiration for molecular SSL and is certainly one reason why the field was able to advance so fast. We use its insights by putting focus on linear probing and investigating dimensional collapse [Hua et al., 2021].
The main goal of SSL is to learn an embedding space that faithfully reflects the complexity of molecular graphs, captures the topological nature of the molecular representation space overall, and whose representations can be effectively adapted given labels during fine-tuning. We propose methods based on persistent homology and show that they naturally suit SSL. First, persistence images (or comparable topological fingerprints) offer great versatility in that they are able to flexibly represent knowledge about the graphs and allow for incorporating domain knowledge. Second, they capture this knowledge based on persistence diagrams, which is very different from - and hence likely complementary to - the common graph representation methods in deep learning. Third, and most importantly, their stability represents a unique feature, which makes them ideal views for SSL.
We study two SSL approaches based on PH and evaluate them in detail in terms of both their impact on representation power (see Section 4.1) and on downstream performance (see Section 4.2):
- In order to study the impact of PH on SSL in general, we consider a simple autoencoder architecture.
- Since we consider topological fingerprints to represent information that is complementary to that used in existing approaches (we also show that they are not ideal alone) and because their topological nature can be used to improve the latter in unique ways, we developed a loss function based on contrastive learning that complements existing approaches.
In this paper, our focus is on obtaining initial insights about the potential PH offers for molecular SSL, hence we chose one basic solution and one providing unique impact. There are certainly other promising ways to be investigated in the future.
3.1 Topological Fingerprints AutoEncoder (TAE)
Autoencoders are designed to reconstruct certain inputs given context information for the input graph  $G$ . We consider topological fingerprints  $I_{G}$  as the reconstruction targets, specifically, PIs.
Figure 4: Overview of Topological Distance Contrastive Learning.
The model itself employs a typical graph encoder  $\varepsilon (\cdot)$  for computing embeddings  $h_v$  for the individual nodes  $v\in V$  ; for simplicity, we write  $h_V = \varepsilon (G)$  , where  $h_V = \{h_v|v\in V\}$  represents all node representations. Next, we pass it through a projection head  $g(\cdot)$  and readout function  $R(\cdot)$  e.g., mean pooling) to obtain the graphlevel representation  $h_G = R(g(h_V))$  Hence the context is the particular knowledge the graph encoder was de
Figure 3: Comparison of molecule similarity based on PIs to similarity between corresp. ECFPs (blue:  $20\%$  most similar).
signed to exploit; in our evaluation, we consider the commonly used GIN [Xu et al.] in order to ease comparison with related works. As our loss formula for this topological fingerprints autoencoder (TAE) we use the regular mean squared error:  $\mathcal{L}_{\mathrm{TAE}} = \sum_{G} \mathrm{MSE}(h_G, I_G)$ .
Observations. Given stable fingerprints such as PIs, one main usual criticism for autoencoders, the fact that they fail to capture inter-molecule relationships, does not apply if the model is able to reliably learn the PIs, which we show TAE does; in particular, we observe a strong correlation between PIs and their reconstructions (see Table 7). The successful application of topological fingerprints in supervised learning and the simplicity of the model justify the study of TAE for analysis, however, we note that the PIs employed may not necessarily capture all information which is critical for a particular downstream task. For example, Figure 3 shows that PIs may offer a generally fitting embedding space, in that randomly chosen molecule pairs with similar standard fingerprints based on substructures (blue, Tanimoto similarity of ECFPs) have rather similar PIs (x-axis, cosine similarity). However, there is a clear difference for the two depicted datasets, sometimes PIs as shown here, based on the ToDD filtration Demir et al. [2022], fail to fully capture structural similarity; and this is directly reflected in performance (see Table 4). While optimization on a case-by-case basis w.r.t. the choice of fingerprints is possible, this is not in the spirit of foundational SSL, requires expert knowledge or extensive tuning, and may still not be sufficient. For that reason, we suggest to apply them together with existing approaches and show that they offer unique benefits.
3.2 Topological Distance Contrastive Loss (TDL)
Contrastive learning aims at learning an embedding space by comparing training samples and encouraging representations from positive pairs of examples to be close in the embedding space while representations from negative pairs are pushed away from each other. Current approaches usually consider each sample as its own class, that is, a positive pair consists of two different views of it; and all other samples in a batch are used as the negative pairs during training. We observe that this represents a very coarse-grained comparison basically ignoring all commonalities and differences between the given molecules; this is also why current efforts in graph CL focus on constructing views that capture the possible relations between molecules.
Our Idea, Figure 4. We exploit the stability of topological fingerprints such as PIs to model the distances between the given molecules (i.e., instead of just views of the same molecule) and to use them for fine-grained supervision in SSL (recall that stability means the distances between PIs reflect those between the topological persistence diagrams). This is very different from and complementary to related works in that it structures the embedding space in a different way. The fingerprints are usually constructed in a way so that they capture information about the molecular graph structure; even if they do not capture the entire complexity of the molecules, they represent some, probably important aspects. Moreover, based on the topological nature, they may capture aspects not represented by the commonly used graph embeddings. In particular, they offer a way to flexibly integrate expert knowledge into molecular SSL; observe that the usually employed GNNs have dimensions based on the available molecule information, hence including additional information, even if available only for some of the data, requires architecture changes.
We consider a batch of  $N$  graphs,  $\{G_i\}_{i \in [1,N]}$ . Similar as above, we first extract graph-level representation vectors  $h_{G_i}$  using a graph encoder  $\varepsilon(\cdot)$ , followed by a readout function  $R(\cdot)$ . Here, we apply the projection head  $g(\cdot)$  later, to map the graph representations to another latent space and obtain the final graph embedding  $z_i$ . Specifically, we apply a two-layer MLP and hence a non-linear transformation, known to enhance performance [Chen et al., 2020]:  $z_i = g(R(\varepsilon(G_i)))$
Let  $G_{n}$  be the sample under consideration. Instead of constructing an artificial view, we consider all possible positive pairs of samples  $(G_{n}, G_{m})$  together with a set of stable, topological fingerprints  $\{I_{i}\}_{i \in [1,N]}$  for all graphs. Note that these can be computed rather efficiently, and they have to be computed only once for the given pre-training data. In a nutshell, our loss adapts the regular NT-Xent [Sohn, 2016, Oord et al., 2018, Wu et al., 2018b] by considering only those negative pairs  $(G_{n}, G_{k})$  of samples where the Euclidean distance  $\text{dis}(I_{n}, I_{k})$  between the corresponding fingerprints is greater than the one between  $I_{n}$  and  $I_{m}$ . We compute the similarity score as  $\text{sim}(z_{n}, z_{m}) = z_{n}^{\top} z_{m} / \|z_{n}\| \|z_{m}\|$ , and consider a temperature parameter  $\tau$  and indicator function  $\mathbb{I}_{[\cdot]} \in \{0,1\}$  as usual. Our topological distance contrastive loss (TDL) is defined as follows, for the n-th sample:
$$
\mathcal{L}_{\mathrm{TDL}_{n}} = \frac{1}{N - 1}\sum_{\substack{m\in [1,N],\\ m\neq n}} - \log \frac{e^{sim(z_{n},z_{m}) / \tau}}{\sum_{\substack{k\in[1,N],\\ k\neq n}}\mathbb{I}_{[dis(I_{n},I_{k})\geq dis(I_{n},I_{m})]}\cdot e^{sim(z_{n},z_{k}) / \tau}}
$$
TDL can be flexibly and efficiently applied to complement any graph CL framework, e.g., in the form  $\mathcal{L}_n = \mathcal{L}_{\text{others}} + \lambda \mathcal{L}_{\text{TDL}_n}$ , where  $\lambda$  determines its impact. In our evaluation, we used  $\lambda = 1$ .