SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiang Zhuang $^{1,2,3*}$ , Qiang Zhang $^{1,2,3*†}$ , Keyan Ding $^{2}$ , Yatao Bian $^{4}$ , Xiao Wang $^{5}$ , Jingsong Lv $^{6}$ , Hongyang Chen $^{6}$ , Huajun Chen $^{1,2,3†}$
<sup>1</sup>College of Computer Science and Technology, Zhejiang University
<sup>3</sup>Zhejiang University - Ant Group Joint Laboratory of Knowledge Graph
$^{4}$ Tencent AI Lab,  $^{5}$ School of Software, Beihang University,  $^{6}$ Zhejiang Lab {zhuangxiang, qiang.zhang.cs, dingkeyan, huajunsir} @zju.edu.cn yatao.bian@gmail.com, xiao_wang@buaa.edu.cn
{jingsonglv,hongyang}@zhejianglab.com
Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called "first-encoding-then-separation" to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at https://github.com/HICAI-ZJU/iMoLD.
Computer-aided drug discovery has played an important role in facilitating molecular design, aiming to reduce costs and alleviate the high risk of experimental failure [1, 2]. In recent years, the emergence of deep learning has led to a growing interest in molecular representation learning, which aims to encode molecules as low-dimensional and dense vectors [3, 4, 5, 6, 7, 8]. These learned representations have demonstrated their availability in various tasks, including target structure prediction [9], binding affinity analysis [10], drug re-purposing [11] and retrosynthesis [12].
Despite the significant progress in molecular representation methods, a prevailing assumption in the traditional approaches is that data sources are independent and sampled from the same distribution. However, in practical drug development, molecules exhibit diverse characteristics and may originate from different distributions [13, 14]. For example, in virtual screening scenarios [15], the distribution shift occurs not only in the molecule itself, e.g., size [13] or scaffold [16] changes, but also in the
target [14], e.g., the emergent COVID-19 leads to a new target from unknown distributions. This out-of-distribution (OOD) problem poses a challenge to the generalization capability of molecular representation methods and results in the degradation of performance in downstream tasks [17, 18].
Current studies mainly focus on regular Euclidean data for OOD generalization. Most studies [18, 19, 20, 21, 22] adopt the invariance principle [23, 24], which highlights the importance of focusing on the critical causal factors that remain invariant to distribution shifts while overlooking the spurious parts [23, 19, 20]. Although the invariance principle has shown effectiveness on Euclidean data, its application to non-Euclidean data necessitates further investigation and exploration. Molecules are often represented as graphs, a typical non-Euclidean data, where atoms as nodes and bonds as edges, thereby preserving rich structural information [3]. The complicated molecular graph structure makes it challenging to accurately distinguish the invariant causal parts from diverse spurious correlations [25].
Preliminary studies have made some attempts on molecular graphs [26, 25, 27, 28, 29]. They explicitly divide graphs to extract invariant substructures, at the granularity of nodes [28], edges [26, 27, 28, 29], or motifs [25]. These attempts can be summarized as the "first-separation-then-encoding" paradigm (Figure 1 (a)), which first divides the graph into invariant and spurious parts and then encodes each part separately. We argue that this practice is suboptimal for extremely complex and entangled graphs, such as real-world molecules [30, 31], since some intricate properties cannot be readily determined by analyzing a subset of the molecular structure [30, 31]. Besides, some methods [25, 27] require assumptions and inferences about the environmental distribution, which are often untenable for molecules due to the intricate environment. Additionally, the downstream tasks related to molecules are diverse, including regression and classification. However, some methods such as CIGA [26] and DisC [32] can only be applied to single-label classification
Figure 1: (a) First-Separation-Then-Encoding: the input is separated using a subgraph generator, then each subgraph is encoded respectively. (b): First-Encoding-Then-Separation: the input is encoded, then the representation is separated by a scorer.
tasks, due to the constraints imposed by the invariant learning objective function.
To fill these gaps, we present a novel molecule invariant learning framework to effectively capture the invariance of molecular graphs and achieve generalized representation against distribution shifts. In contrast to the conventional approaches, we propose a "first-encoding-then-separation strategy" (Figure 1 (b)). Specifically, we first employ a Graph Neural Network (GNN) [33, 34, 35] to encode the molecule (i.e., encoding GNN), followed by a residual vector quantization module to alleviate the over-fitting to training data distributions while preserving the expressivity of the encoder. We then utilize another GNN to score the molecule representation (i.e., scoring GNN), which measures the contributions of each dimension to the target in latent space, resulting in a clear separation between invariant and spurious representations. Finally, we design a self-supervised learning objective [36, 37, 38] that aims to encourage the identified invariant features to effectively preserve label-related information while discarding environment-related information. It is deserving of note that the objective is task-agnostic, which means that our method can be applied to various tasks, including regression and single- or multi-label classification.
- We propose a paradigm of "first-encoding-then-separation" using an encoding GNN and a scoring GNN, which enables us to effectively identify invariant features from highly complex graphs.
- We introduce a residual vector quantization module that strikes a balance between model expressivity and generalization. The quantization acts as the bottleneck to enhance generalization, while the residual connection complements the model's expressivity.
- We design a self-supervised invariant learning objective that facilitates the precise capture of invariant features. This objective is versatile, task-agnostic, and applicable to a variety of tasks.
Figure 2: An overview of distribution shifts in molecules. Distribution shifts occur when molecules originate from different scaffold, size or assay environments.
- We conduct comprehensive experiments on a diverse set of real-world datasets with various distribution shifts. The experimental results demonstrate the superiority of our method compared to state-of-the-art approaches.
OOD Generalization and Invariance Principle. The susceptibility of deep neural networks to substantial performance degradation under distribution shifts has led to a proliferation of research focused on out-of-distribution (OOD) generalization [39]. Three lines of methods have been studied for OOD generalization on Euclidean data, including group distributionally robust optimization [40, 41, 42], domain adaptation [43, 44, 45] and invariant learning [19, 20, 21]. Group distributionally robust optimization considers groups of distributions and optimize across all groups simultaneously. Domain adaptation aims to align the data distributions but may fail to find an optimal predictor without additional assumptions [19, 20, 46]. Invariant learning aims to learn invariant representation satisfying the invariant principle [20, 24], which includes two assumptions: (1) sufficiency, meaning the representation has sufficient predictive abilities, (2) invariance, meaning representation is invariant to environmental changes. However, most methods require environmental labels, which are expensive to obtain for molecules [26], and direct application of these methods to complicated non-Euclidean molecular structure does not yield promising results [13, 22, 26].
OOD Generalization on Graphs. Recently, there has been growing attention on the graph-level representations under distribution shifts from the perspective of invariant learning. Some methods [26, 32, 27, 28, 25] follow the "first-separation-then-encoding" paradigm, and they make attempt to capture the invariant substructures by dividing nodes and edges in the explicit structural space. However, these methods suffer from the difficulty in dividing molecule graphs in raw space due to the complex and entangled molecular structure [47]. Moreover, MoleOD [25] and GIL [27] require inference of unavailable environmental labels, which entails prior assumptions on the environmental distribution. And the objective of invariant learning may hinder the application, e.g., CIGA [26] and DisC [32] can only apply to single-label classification rather than to regression and multi-label tasks. Additionally, OOD-GNN [48] does not use the invariance principle. It proposes to learn disentangled graph representations, but requires computing global weights for all data, leading to a high computational cost. OOD generalization on graphs can also be improved by another line of relevant works [29, 49, 47] on GNN explainability [50, 51], which aims to provide a rationale for prediction. However, they may fail in some distribution shift cases [26]. And DIR [29] and GSAT [49] also divide graphs in the raw structural space. Although GREA [47] learns in the latent feature space, it only conducts separation on each node while neglecting the different significance of each representation dimensionality. In this work, we focus on the OOD generalization of molecular graphs, against multi-type of distribution shifts, e.g., scaffold, size, and assay, as shown in Figure 2.
Vector Quantization. Vector Quantization (VQ) [52, 53] acts as a bottleneck of representation learning. It discretizes continuous input data in the hidden space by assigning them to the nearest vectors in a predefined codebook. Some studies [53, 54, 55, 56] have demonstrated its effectiveness to enhance model robustness against data corruptions. Other studies [57, 58, 59] find that taking VQ as an inter-component communication within neural networks can contribute to the model generalization ability. However, we posit that while VQ can improve generalization against distribution shifts, it may also limit the model's expressivity and potentially lead to under-fitting. To address this concern, we propose to equip the conventional VQ with a residual connection to strike a balance between model generalization and expressivity.
We focus on the OOD generalization of molecules. Let  $\mathcal{G}$  be the molecule graph space and  $\mathcal{V}$  be the label space, the goal is to find a predictor  $f:\mathcal{G}\to \mathcal{V}$  to map the input  $G\in \mathcal{G}$  into the label  $Y\in \mathcal{V}$ . Generally, we are given a set of datasets collected from multiple environments  $E_{all}:D = \{D^{e}\}_{e\in E_{all}}$ . Each  $D^{e}$  contains pairs of an input molecule and its label:  $D^{e} = \{(G_{i},Y_{i})\}_{i = 1}^{n_{e}}$  that are drawn from the joint distribution  $P(G,Y|E = e)$  of environment  $e$ . However, the information of environment  $e$  is always not available for molecules, thus we redefine the training joint distribution as  $P_{train}(G,Y) = P(G,Y|E = e),\forall e\in E_{train}$ , and the testing joint distribution as  $P_{test}(G,Y) = P(G,Y|E = e),\forall e\in E_{all}\setminus E_{train}$ , where  $P_{train}(G,Y)\neq P_{test}(G,Y)$ . We denote joint distribution across all environment as  $P_{all}(G,Y) = P(G,Y|E = e),\forall e\in E_{all}$ . Formally, the goal is to learn an optimal predictor  $f^{*}$  based on training data and can generalize well across all distributions:
$$
f ^ {*} = \arg \min  _ {f} \mathbb {E} _ {(G, Y) \sim P _ {\text {a l l}}} [ \ell (f (G), Y) ], \tag {1}
$$
where  $\ell (\cdot ,\cdot)$  is the empirical risk function. Moreover, since the joint distribution  $P(G,Y)$  can be written as  $P(Y|G)P(G)$ , the OOD problem can be refined into two cases, namely covariate and concept shift [60, 61, 62, 63]. In covariate shift, the distribution of input differs. Formally,  $P_{train}(G)\neq P_{test}(G)$  and  $P_{train}(Y|G) = P_{test}(Y|G)$ . While concept shift occurs when the conditional distribution changes as  $P_{train}(Y|G)\neq P_{test}(Y|G)$  and  $P_{train}(G) = P_{test}(G)$ . We will consider and distinguish between both cases in our experiments.
3.2 Molecular Representation Learning
We denote a molecule graph by  $G = \{\mathcal{V},\mathcal{E}\}$ , where  $\mathcal{V}$  is the set of nodes (e.g., atoms) and  $\mathcal{E} \in \mathcal{V} \times \mathcal{V}$  is the set of edges (e.g., chemical bonds). Generally, the predictor  $f$  can be denoted as  $\rho \circ g$ , containing an encoder  $g: \mathcal{G} \to \mathbb{R}^d$  that extracts representation for each molecule and a downstream classifier  $\rho: \mathbb{R}^d \to \mathcal{V}$  that predicts the label with the molecular representation. In particular, the encoder  $g$  operates in two stages: firstly, by employing a graph neural network [33, 34, 35] to generate node representations  $\mathbf{H}$  according to the following equation:
$$
\mathbf {H} = \left[ \mathbf {h} _ {1}, \mathbf {h} _ {2}, \dots , \mathbf {h} _ {| \mathcal {V} |} \right] ^ {\top} = \operatorname {G N N} (G) \in \mathbb {R} ^ {| \mathcal {V} | \times d}, \tag {2}
$$
where  $\mathbf{h}_v\in \mathbb{R}^d$  is the representation of node  $v$ . Secondly, the encoder utilizes a readout operator to obtain the overall graph representation  $\mathbf{z}$ :
$$
\mathbf {z} = \operatorname {R E A D O U T} (\mathbf {H}) \in \mathbb {R} ^ {d}. \tag {3}
$$
The readout operator can be implemented using a simple, permutation invariant function such as average pooling.
This section presents the details of our proposed method that learns invariant Molecular representation in Latent Discrete space, called iMoLD. Figure 3 shows the overview of iMoLD, which mainly consists of three steps: 1) Using a GNN encoder and a residual vector quantization module to obtain the molecule representation (Section 4.1); 2) Separating the representation into invariant and spurious parts through a GNN scorer (Section 4.2); 3) Optimizing the above process with a task-agnostic self-supervised learning objective (Section 4.3).
4.1 Encoding with Residual Vector Quantization
The current mainstream methods [26, 25, 28, 27, 32] adopt the "first-separation-then-encoding" paradigm, which explicitly divides graphs into invariant and spurious substructures on the granularity of edge, node, or motif, and then encodes each substructure individually. In contrast, we use the opposite paradigm that first encodes the whole molecule followed by separation.
Figure 3: An overview of iMoLD. Firstly, given a batch of inputs, we learn invariant and spurious representations  $(\mathbf{z}^{\mathrm{Inv}}$  and  $\mathbf{z}^{\mathrm{Spu}})$  for each input in latent discrete space by a first-encoding-then-separation paradigm. An encoding GNN and an RVQ module are involved to obtain molecule representation, then the representation is separated through a scoring GNN. The invariant  $\mathbf{z}^{\mathrm{Inv}}$  is used to predict the label  $\widehat{y}$ . Then a task-agnostic self-supervised learning objective across the batch is designed to facilitate the acquisition of reliable invariant  $\mathbf{z}^{\mathrm{Inv}}$ .
Specifically, given an input molecule  $G$ , we first use a GNN to encode it, resulting in node representations  $\mathbf{H}$ :
$$
\mathbf {H} = \operatorname {G N N} _ {E} (G) \in \mathbb {R} ^ {| \mathcal {V} | \times d}, \tag {4}
$$
where  $\mathrm{GNN}_E$  represents the encoding GNN,  $|\mathcal{V}|$  is the number of nodes in  $G$ , and  $d$  is the dimensionality of features. Inspired by the studies [57, 59] that Vector Quantization (VQ) [52, 53] is helpful to improve the model generalization on computer vision tasks, we propose a Residual Vector Quantization (RVQ) module to refine the obtained representations.
In the RVQ module, VQ is used to discretize continuous representations into discrete ones. Formally, it introduces a shared learnable codebook as a discrete latent space:  $\mathcal{C} = \{\mathbf{e}_1,\mathbf{e}_2,\dots \mathbf{e}_{|\mathcal{C}|}\}$ , where each  $\mathbf{e}_k\in \mathbb{R}^d$ . For each node representation  $\mathbf{h}_v$  in  $\mathbf{H}$ , VQ looks up and fetches the nearest neighbor in the codebook and outputs it as the result. Mathematically,
$$
\mathrm {Q} \left(\mathbf {h} _ {v}\right) = \mathbf {e} _ {k}, \quad \text {w h e r e} \quad k = \underset {k \in \{1, \dots , | \mathcal {C} | \}} {\operatorname {a r g m i n}} \| \mathbf {h} _ {v} - \mathbf {e} _ {k} \| _ {2}, \tag {5}
$$
and  $\mathrm{Q}(\cdot)$  denotes the discretization operation which quantizes  $\mathbf{h}_v$  to  $\mathbf{e}_k$  in the codebook.
The VQ operation acts as a bottleneck to enhance generalization and alleviate the easy-over-fitting issue caused by distribution shifts. However, it also impairs the expressivity of the model by using a limited discrete codebook to replace the original continuous input, suffering from a potential under-fitting issue. Accordingly, we propose to equip the conventional VQ with a residual connection to strike a balance between model generalization and expressivity. In specific, we incorporate both the continuous and discrete representations to update node representations  $\mathbf{H}$  to  $\mathbf{H}'$ :
$$
\mathbf {H} ^ {\prime} = \left[ Q \left(\mathbf {h} _ {1}\right) + \mathbf {h} _ {1}, Q \left(\mathbf {h} _ {2}\right) + \mathbf {h} _ {2}, \dots , Q \left(\mathbf {h} _ {| \mathcal {V} |}\right) + \mathbf {h} _ {| \mathcal {V} |} \right] ^ {\top}. \tag {6}
$$
Similar to VQ-VAE [52, 53], we employ the exponential moving average updates for the codebook:
$$
N _ {k} ^ {(t)} = N _ {k} ^ {(t - 1)} * \eta + n _ {k} ^ {(t)} (1 - \eta), \quad \mathbf {m} _ {k} ^ {(t)} = \mathbf {m} _ {k} ^ {(t - 1)} * \eta + \sum_ {v} ^ {n _ {k} ^ {(t)}} \mathbf {h} _ {v} ^ {(t)} (1 - \eta), \quad \mathbf {e} _ {k} ^ {(t)} = \frac {\mathbf {m} _ {k} ^ {(t)}}{N _ {k} ^ {(t)}}, \tag {7}
$$
where  $n_k^{(t)}$  is the number of node representations in the  $t$ -th mini-batch that are quantized to  $\mathbf{e}_k$ , and  $\eta$  is a decay parameter between 0 and 1.
4.2 Separation at Nodes and Features
After encoding, we separate the representation into invariant parts and spurious parts. It is worth noting that our separation is not only performed at the node dimension but also takes into account the feature dimension in the latent space. The reasons are two-fold: 1) Distribution shifts on molecules
can occur at both the structure level and the attribute level [26], corresponding to the node dimension  $|\mathcal{V}|$  and the feature dimension  $d$  in  $\mathbf{H}'$ , respectively. 2) The resulting representation may be highly entangled, thus it is advisable to perform a separation on each dimension in the latent space.
Specifically, we use another GNN as a scorer to obtain the separating score S:
$$
\mathbf {S} = \sigma \left(\mathrm {G N N} _ {S} (G)\right) \in \mathbb {R} ^ {| \mathcal {V} | \times d}, \tag {8}
$$
where  $\mathrm{GNN}_S$  represents the scoring GNN,  $|\mathcal{V}|$  is the number of nodes in  $G$ , and  $d$  is the dimensionality of features.  $\sigma(\cdot)$  denotes the Sigmoid function to constrain each entry in  $\mathbf{S}$  falls into the range of  $(0,1)$ . Then we can capture the invariant and complementary spurious features at both structure and attribute granularity in the latent representation space by applying the separating scores to node representations:
$$
\mathbf {H} ^ {\text {I n v}} = \mathbf {H} ^ {\prime} \odot \mathbf {S}, \quad \mathbf {H} ^ {\text {S p u}} = \mathbf {H} ^ {\prime} \odot (1 - \mathbf {S}), \tag {9}
$$
where  $\mathbf{H}^{\mathrm{Inv}}$  and  $\mathbf{H}^{\mathrm{Spu}}$  denote the invariant and spurious node representations respectively, and  $\odot$  is the element-wise product. Finally, the invariant and spurious representation (denoted as  $\mathbf{z}^{\mathrm{Inv}}$  and  $\mathbf{z}^{\mathrm{Spu}}$  respectively) of  $G$  can be generated by a readout operator:
$$
\mathbf {z} ^ {\text {I n v}} = \operatorname {R E A D O U T} \left(\mathbf {H} ^ {\text {I n v}}\right) \in \mathbb {R} ^ {d}, \quad \mathbf {z} ^ {\text {S p u}} = \operatorname {R E A D O U T} \left(\mathbf {H} ^ {\text {S p u}}\right) \in \mathbb {R} ^ {d}. \tag {10}
$$
4.3 Learning Objective
Our OOD optimization objectives are composed of an invariant learning loss, a task prediction loss, and two additional regularization losses.
Task-agnostic Self-supervised Invariant Learning. Invariant learning aims to optimize the encoding  $\mathrm{GNN}_E$  and the scoring  $\mathrm{GNN}_S$  to produce precise invariant and spurious representations. In particular, we need  $\mathbf{z}^{\mathrm{Inv}}$  to be invariant under environmental changes. Additionally, we expect the objective to be independent of the downstream task, which allows the method to be not restricted to a specific type of task. To achieve these, we design a task-agnostic and self-supervised invariant learning objective. Specifically, we disturb  $\mathbf{z}_i^{\mathrm{Inv}}$  via concatenating a corresponding  $\mathbf{z}_j^{\mathrm{Spu}}$  in a shuffled batch, resulting in an augmented representation  $\widetilde{\mathbf{z}}_i^{\mathrm{Inv}}$ :
$$
\hat {\mathbf {z}} _ {i} ^ {\text {I n v}} = \mathbf {z} _ {i} ^ {\text {I n v}} \oplus \mathbf {z} _ {j \in [ 1, B ]} ^ {\text {S p u}}, \tag {11}
$$
where  $\oplus$  denotes concatenation operator and  $B$  is batch size.
Inspired by a simple self-supervised learning framework that takes different augmentation views as similar positive pairs and does not require negative samples [37], we treat  $\mathbf{z}_i^{\mathrm{Inv}}$  and  $\widetilde{\mathbf{z}}_i^{\mathrm{Inv}}$  as positive pairs and push them to be similar, using an MLP-based predictor (denoted as  $\omega$ ) that transforms the output of one view and aligns it to the other view. We minimize their negative cosine similarity:
$$
\mathcal {L} _ {\mathrm {i n v}} = - \sum_ {i = 1} ^ {B} \operatorname {s i m} \left(\operatorname {s g} \left[ \mathbf {z} _ {i} ^ {\text {I n v}} \right], \omega \left(\widetilde {\mathbf {z}} _ {i} ^ {\text {I n v}}\right)\right), \tag {12}
$$
where  $\mathrm{sim}(\cdot ,\cdot)$  represents the formula of cosine similarity and  $\mathrm{sg}[\cdot ]$  denotes stop-gradient operation to prevent collapsing [37]. We employ  $\mathcal{L}_{\mathrm{inv}}$  as our objective for invariant learning to ensure the invariance of  $\mathbf{z}^{\mathrm{Inv}}$  against distribution shifts.
Task Prediction. The objective of task prediction is to provide an invariant presentation  $\mathbf{z}^{\mathrm{Inv}}$  with sufficient predictive abilities. During training, the choice of prediction loss function depends on the type of task. For classification tasks, we employ the cross-entropy loss, while for regression tasks, we use the mean squared error loss. Take the binary classification task as an example, the cross-entropy loss is computed between the predicted  $\widehat{y_i} = \rho (\mathbf{z}_i^{\mathrm{Inv}})$  and the ground-truth label  $y_{i}$ :