SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shengchao Liu $^{1,2}$ , Weitao Du $^{3}$ , Yanjing Li $^{4}$ , Zhuoxinran Li $^{5}$ , Zhiling Zheng $^{6}$ , Chenru Duan $^{7}$
Zhiming Ma<sup>3</sup>, Omar Yaghi<sup>6</sup>, Anima Anandkumar<sup>8</sup>, Christian Borgs<sup>6</sup>,
Jennifer Chayes<sup>6</sup>, Hongyu Guo<sup>9</sup>, Jian Tang<sup>1,10,11</sup>
<sup>1</sup>Mila - Québec Artificial Intelligence Institute <sup>2</sup>Université de Montréal
<sup>3</sup>University of Chinese Academy of Sciences  <sup>4</sup>Carnegie Mellon University
$^{5}$ University of Toronto  $^{6}$ University of California, Berkeley  $^{7}$ Massachusetts Institute of Technology
$^{8}$ California Institute of Technology  $^{9}$ National Research Council Canada
10HEC Montréal 11CIFAR AI Chair
Correspondence: shengchao.liu@umontreal.ca, jian.tang@hec.ca
Artificial intelligence for scientific discovery has recently generated significant interest within the machine learning and scientific communities, particularly in the domains of chemistry, biology, and material discovery. For these scientific problems, molecules serve as the fundamental building blocks, and machine learning has emerged as a highly effective and powerful tool for modeling their geometric structures. Nevertheless, due to the rapidly evolving process of the field and the knowledge gap between science (e.g., physics, chemistry, & biology) and machine learning communities, a benchmarking study on geometrical representation for such data has not been conducted. To address such an issue, in this paper, we first provide a unified view of the current symmetry-informed geometric methods, classifying them into three main categories: invariance, equivariance with spherical frame basis, and equivariance with vector frame basis. Then we propose a platform, coined Geom3D, which enables benchmarking the effectiveness of geometric strategies. Geom3D contains 16 advanced symmetry-informed geometric representation models and 14 geometric pretraining methods over 52 diverse tasks, including small molecules, proteins, and crystalline materials. We hope that Geom3D can, on the one hand, eliminate barriers for machine learning researchers interested in exploring scientific problems; and, on the other hand, provide valuable guidance for researchers in computational chemistry, structural biology, and materials science, aiding in the informed selection of representation techniques for specific applications. The source code is available on the GitHub repository.
Artificial intelligence (AI) for molecule discovery has recently seen many developments, including small molecular property prediction [13, 17, 23, 38, 66, 78, 101, 103, 104, 119, 130, 132, 135], small molecule design and optimization [6, 54, 57, 85, 137], small molecule reaction and retrosynthesis [40, 111, 116], protein property prediction [27, 141], protein folding and inverse folding [48, 64, 92], protein design [15, 41, 46, 88, 91], and crystalline material design [33, 125, 128]. One of the most fundamental building blocks for these tasks is the geometric structure of molecules. Exploring effective methods for robust representation learning to leverage such geometric information fully remains an open challenge that interests both machine learning (ML) and science researchers.
To this end, symmetry-informed geometric representation [1] has emerged as a promising approach. By leveraging physical principles (i.e., group theory for depicting symmetric particles) into spatial
Figure 1: Pipeline for Geom3D, including dataset preprocessing, feature extraction, geometric pretraining and representation, and target tasks. We additionally demonstrate the SE(3)-equivariant force prediction task.
representation, they facilitate a more robust representation of small molecules, proteins, and crystalline materials. Nevertheless, pursuing geometric learning research is still challenging due to its evolving nature and the knowledge gap between science (e.g., physics) and machine learning communities. These factors contribute to a substantial barrier for machine learning researchers to investigate scientific problems and hinder efforts to reproduce results consistently. To overcome this, we introduce Geom3D, a benchmarking of the geometric representation with four advantages, as follows.
(1) A unified and novel aspect in understanding symmetry-informed geometric models. The molecule geometry needs to satisfy certain physical constraints regarding the 3D Euclidean space. For instance, the molecules' force needs to be equivariant to translation and rotation (see SE(3)-equivariance in Fig. 1). In this work, we classify the geometric methods into three categories: invariant model, SE(3)-equivariant model with spherical frame basis and vector frame basis. The invariant models only consider features that are constant w.r.t. the SE(3) group, while the two families of equivariant models can be further unified using the frame basis to capture equivariant symmetry. An illustration of three categories is in Fig. 2. Building equivariant models on the frame basis provides a novel and unified view of understanding geometric models and paves the way for intriguing more ML researchers to explore scientific problems.
Figure 2: Three categories of geometric modules. (a) Invariant models only consider type-0 features. Equivariant models use either (b) spherical harmonics frames or (c) vector frames by projecting the coordinate vectors.
(2) A unified platform for various scientific domains. There exist multiple platforms and tools for molecule discovery, but they are (1) mainly focusing on molecule's 2D graph representation [77, 102, 145]; (2) using 3D geometry with customized data structures or APIs [3, 105]; or (3) covering only a few geometric models [76]. Thus, it is necessary to have a platform benchmarking the geometric models, especially for researchers interested in solving scientific problems. In this work, we propose Geom3D, a geometric modeling framework based on PyTorch Geometric (PyG) [31], one of the most widely-used platforms for graph representation learning. Geom3D benchmarks 16 geometric models on solving 52 scientific tasks, and these tasks include the three most fundamental molecule types: small molecules, proteins, and crystalline materials. Each of them requires distinct domain-specific preprocessing steps, e.g., crystalline materials molecules possess periodic structures and thus need a particular periodic data augmentation. By leveraging such a unified framework, Geom3D serves as a comprehensive benchmarking tool, facilitating effective and consistent analysis components to interpret the existing geometric representation functions in a fair and convenient comparison setting.
(3) A framework for a wider range of ML tasks. The geometric models in Geom3D can serve as a building block for exploring extensive ML tasks, including but not limited to studying the molecule dynamic simulation and scrutinizing the transfer learning effect on molecule geometry. For example, pre
In what follows, we may use "molecule" to refer to "small molecule" for brevity.
(a) An example of small molecule structure.
(c) An example of protein sequence and structure.
(b) An illustration of the potential energy surface.
(d) An example of crystalline material.
Figure 3: Fig. 3(a) illustrates 2D topology and 3D conformation for molecule Glycine. Fig. 3(c) displays the 3D structure of protein. Fig. 3(d) shows a simple cubic crystal of the element Po. Fig. 3(b) is a demo of PES.
training is an important strategy to quickly transfer knowledge to target tasks, and recent works explore geometric pretraining on 3D conformations (including supervised and self-supervised) [59, 80, 136] and multi-modality pretraining on 2D topology and 3D geometry [30, 79, 86]. Other transfer learning venues include multi-task learning [82, 84] and out-of-distribution or domain adaptation [58, 133, 134], yet no geometry information has been utilized. All of these directions are promising for future exploration, and Geom3D serves as an auxiliary tool to accomplish them. For example, as will be shown in Sec. 4, we leverage Geom3D to effectively evaluate 14 pretraining methods with benchmarks.
(4) A framework for exploring data preprocessing and optimization tricks. When comparing different symmetry-informed geometric models, we find that in addition to the model architecture, there are two important factors affecting the performance: the data preprocessing (e.g., energy and force rescaling and shift) and optimization methods (e.g., learning rate, learning rate schedule, number of epochs, random seeds). In this work, we explore the effect of four preprocessing tricks and around 2-10 optimization hyperparameters for each model and task. In general, we observe that each model may benefit differently in different tasks regarding the preprocessing and optimization tricks. However, data normalization is found to help improve performance hugely in most cases. We believe that Geom3D is an effective tool for exploring and understanding various engineering tricks.
2 Data Structures for Geometric Data
Small molecule 3D conformation. Molecules are sets of points in the 3D Euclidean space, and they move in a dynamic motion, as known as the potential energy surface (PES). The region with the lowest energy corresponds to the most stable state for molecules, and molecules at these positions are called conformations, as illustrated in Fig. 3(b). For notation, we mark each 3D molecular graph as  $g = (X, R)$ , where  $X$  and  $R$  are for the atom types and positions, respectively.
Crystalline material with periodic structure. The crystalline materials or extended chemical structures possess a characteristic known as periodicity: their atomic or molecular arrangement repeats in a predictable and consistent pattern across all three spatial dimensions. This is the key aspect that differentiates them from small molecules. In Fig. 3(d), we show an original unit cell (marked in green) that can repeatedly compose the crystal structure along the lattice. To model such a periodic structure, we adopt the data augmentation from CGCNN [129]: for each original unit cell, we shift it along the lattice in three dimensions and connect edges within a cutoff value (hyperparameter). For more details on the two augmentation variants, please check Appendix A.
Protein with backbone structure. Protein structures can be classified into four primary levels, and the primary structure represents the linear arrangement of amino acids, and each amino acid is a molecule consisting of atoms. Geometric methods mainly focus on the tertiary structure, i.e., the 3D geometry of each atom, encompassing the complete organization of a single protein. However,
Figure 4: Pipelines for seven single-modal geometric pretraining methods. (a-c) conduct self-prediction. (d) maximizes the MI between nodes and graphs. (e-g) are GeoSSL, maximizing the MI between views  $\pmb{g}_1$  and  $\pmb{g}_2$ .
atom-level modeling for proteins is consuming due to the large volume of atoms and the GPU memory limit. One solution is modeling each amino acid's backbone structure. The backbone structure of each amino acid is  $N - C_{\alpha} - C$ , and the  $C_{\alpha}$  is bonded to the side chain. 20 common types of side chains corresponding to 20 amino acids, as illustrated in Fig. 3(c). Thus, modeling the backbone structure can balance the computational efficiency and the key geometric information.
3 Symmetry-Informed Geometric Representation
3.1 Group Symmetry and Equivariance
Symmetry means the object remains invariant after certain transformations [127], and it is everywhere on Earth, such as in animals, plants, and molecules. Formally, the set of all symmetric transformations satisfies the axioms of a group. Therefore, the group theory and its representation theory are common tools to depict such physical symmetry. Group is a set  $G$  equipped with a group product  $\times$  satisfying:
$$
(1) \exists e \in G, \boldsymbol {a} \times \boldsymbol {e} = \boldsymbol {e} \times \boldsymbol {a}, \forall \boldsymbol {a} \in G; (2) \boldsymbol {a} \times \boldsymbol {a} ^ {- 1} = \boldsymbol {a} ^ {- 1} \times \boldsymbol {a} = \boldsymbol {e}; (3) \boldsymbol {a} \times (\boldsymbol {b} \times \boldsymbol {c}) = \boldsymbol {a} \times \boldsymbol {b} \times \boldsymbol {c}. \tag {1}
$$
Group representation is a mapping from the group  $G$  to the group of linear transformations of a vector space  $X$  with dimension  $d$  (see [138] for more rigorous definition):
$$
\rho_ {X} (\cdot): G \rightarrow \mathbb {R} ^ {d \times d} \quad \text {s . t .} \quad \rho (\boldsymbol {e}) = 1 \wedge \rho_ {X} (\boldsymbol {a}) \rho_ {X} (\boldsymbol {b}) = \rho_ {X} (\boldsymbol {a} \times \boldsymbol {b}), \forall \boldsymbol {a}, \boldsymbol {b} \in G. \tag {2}
$$
During modeling, the  $X$  space can be the input 3D Euclidean space, the equivariant vector space in the intermediate layers, or the output force space. This enables the definition of equivariance as below.
Equivalence is the property for the geometric modeling function  $f: X \to Y$  as:
$$
f \left(\rho_ {X} (\boldsymbol {a}) \boldsymbol {x}\right) = \rho_ {Y} (\boldsymbol {a}) f (\boldsymbol {x}), \forall \boldsymbol {a} \in G, \boldsymbol {x} \in X. \tag {3}
$$
As displayed in Fig. 1, for molecule geometric modeling, the property should be rotation-equivariant and translation-equivariant (i.e., SE(3)-equivariant). More concretely,  $\rho_{X}(\pmb{a})$  and  $\rho_{Y}(\pmb{a})$  are the SE(3) group representations on the input (e.g., atom coordinates) and output space (e.g., force space), respectively. SE(3)-equivariant modeling in Eq. (3) is essentially saying that the designed deep learning model  $f$  is modeling the whole transformation trajectory on the molecule conformations, and the output is the transformed  $\hat{y}$  accordingly. Further, we want to highlight that, in addition to the network architecture or representation function, the input features can also be represented as an equivariant feature mapping from the 3D mesh to  $\mathbb{R}^{\tilde{d}}$  [11], where  $\tilde{d}$  depends on input data, e.g.,  $\tilde{d} = 1$  (for atom type dimension) + 3 (for atom coordinate dimension) on small molecules. Such features are called steerable features in [5, 11] when only considering the subgroup SO(3)-equivariance.
Invariance is a special type of equivariance, defined as:
$$
f \left(\rho_ {X} (\boldsymbol {a}) \boldsymbol {x}\right) = f (\boldsymbol {x}), \forall \boldsymbol {a} \in G, \boldsymbol {x} \in X, \tag {4}
$$
with  $\rho_{Y}(\pmb{a})$  as the identity  $\forall \pmb{a} \in G$ . The group representation helps define the equivariance condition for  $f$  to follow. Then, the question boils down to how to design such an equivariant  $f$ . In the following, we will discuss geometric modelings from a novel and unified perspective using the frame. In the next sections, we will provide a novel and unified aspect of understanding the advanced geometric representation and pretraining methods using the frame basis (details in Appendix H).
3.2 Invariant Geometric Representation Learning
One simple way of achieving SE(3) group symmetry is invariant modeling. It means the geometric model only considers the type-0 features [112], i.e., features that are invariant with respect to rotation and translation. Existing works have been adopting the invariant features for modeling, including pairwise distance (SchNet [109]), bond angles (DimeNet [68]), and torsion angles (SphereNet [89] and GemNet [67]). Note that the torsion angles are angles between two planes defined by pairwise bonds.
3.3 Equivariant Geometric Representation Learning
Invariant modeling only captures the type-0 features. However, equivariant modeling of higher-order particles may bring in extra expressiveness. For example, the elementary particles in high energy physics [98] inherit higher order symmetries in the sense of SO(3) representation theory, which makes the equivariant modeling necessary. Such higher-order particles include type-1 features like coordinates and forces in molecular conformation. There are many approaches to design such SE(3)-equivariant model satisfying Eq. (3). There are two main venues, as will be discussed below.
Spherical Frame Basis. This research line utilizes the irreducible representations [37] for building SO(3)-equivariant representations, and the first work is TFN [112]. Its main idea is to project the 3D Euclidean coordinates into the spherical harmonics space, which transforms equivariantly according to the irreducible representations of SO(3), and the translation-equivariant can be trivially guaranteed using the relative coordinates. Following this, there have been variants combining it with the attention module (Equiformer [73]) or with more expressive network architectures (SEGNN [4], Allegro [95]).
Vector Frame Basis. An alternative philosophy of equivariant modeling utilizes the vector (in physics) frame basis. It constructs three vectors bases, serving as a reference frame to help locate the vectors in each corresponding local environment. Works along this line for molecule discovery include DeePMD [140] for dynamics simulation, 3D-EMGP [59] and MoleculeSDE [79] for geometric pretraining, and ClofNet [20] for conformation generation. For macromolecules like protein, the equivariant vector frame has been used for protein design (StructTrans [53]) and protein folding (AlphaFold2 [64]). We also want to highlight that, from a mathematical perspective, equivariance and invariance can be transformed to each other by the scalarization technique. Please check [49] for details.
The spherical frame basis can be easily extended to higher-order particles, yet it may suffer from the high computational cost. On the other hand, the vector frame basis is specifically designed for the 3D point clouds; thus, it is more efficient but cannot generalize to higher-order particles. Meanwhile, we would like to acknowledge other equivariant modeling paradigms, including using orbital features [99] and elevating 3D Euclidean space to SE(3) group [32, 52]. Please check Appendix F for details.
3.4 Geometric Pretraining
Recent studies have started to explore single-modal of geometric pretraining on molecules. The GeoSSL paper [80] covers a wide range of geometric pretraining algorithms. The type prediction, distance prediction, and angle prediction predict the masked atom type, pairwise distance, and bond angle, respectively. The 3D InfoGraph predicts whether the node- and graph-level 3D representation are for the same molecule. GeoSSL is a novel geometric pretraining paradigm that maximizes the mutual information (MI) between the original conformation  $g_{1}$  and augmented conformation  $g_{2}$ , where  $g_{2}$  is obtained by adding small perturbations to  $g_{1}$ . RR, InfoNCE, and EBM-NCE optimize the objective in the latent representation space, either generative or contrastive. GeoSSL-DDM [80, 136] optimizes the same objective function using denoising score matching. 3D-EMGP [60] has the same strategy and utilizes an equivariant module to denoise the 3D noise directly. Another research line is the multi-modal of topological and geometric pretraining. GraphMVP [86] first proposes one contrastive objective (EBM-NCE) and one generative objective (VRR) to optimize the MI between the 2D topologies and 3D geometries in the representation space. 3D InfoMax [114] is a special case of GraphMVP, with the contrastive part only. MoleculeSDE [79] extends GraphMVP by introducing two SDE models for solving the 2D and 3D reconstruction. We illustrate these algorithms in Figs. 4 and 8.
3.5 Discussion: Reflection-antisymmetric in Geometric Learning
Till now, we have discussed the SE(3)-equivariance, i.e., the translation and rotation equivariance. As highlighted in the recent work [61, 79], the molecules needlessly satisfy the reflection-equivariant,
Table 1: Results of 26 models on 12 quantum mechanics prediction tasks in QM9, with 110K for training, 10K for validation, and 11K for testing. The task unit is specified, and the evaluation is the mean absolute error (MAE).
Featurization	Model	α↓ α03	∇E ↓ meV	EHOMO ↓ meV	ELUMO ↓ meV	μ ↓ D	Cv ↓ cal mol·K	G ↓ meV	H ↓ meV	R2 ↓ α02	U ↓ meV	U0 ↓ meV	ZPVE ↓ meV
1D FPs	MLP	2.231	196.72	131.27	164.94	0.526	0.919	2158.64	2358.23	68.621	2340.61	2314.77	155.921
RF	3.801	207.02	165.72	183.04	0.534	1.485	3391.79	3729.94	94.512	3705.75	3678.25	253.132
XGB	2.748	199.71	139.88	165.43	0.516	1.062	2563.93	2804.27	82.959	2786.28	2769.29	180.989
1D SMILES	CNN	0.364	165.22	124.65	114.81	0.566	0.173	156.66	170.59	20.403	166.18	169.89	10.070
BERT	0.313	117.50	84.93	98.88	0.446	0.176	170.01	183.43	18.002	183.84	188.60	13.410
1D SELFIES	CNN	0.345	157.04	115.51	113.00	0.499	0.168	136.42	146.56	20.080	143.00	140.01	10.149
BERT	0.348	123.11	91.15	90.80	0.461	0.203	168.20	187.50	19.125	204.93	195.98	17.328
2D Graph	GCN	1.338	145.82	96.21	106.66	0.434	0.526	1198.12	1291.57	37.585	1281.03	1303.39	85.103
ENN-S2S	1.401	270.59	129.18	132.84	0.577	0.760	1487.21	955.24	34.609	1800.79	1521.32	51.226
GraphSAGE	1.601	131.45	88.78	93.21	0.402	0.544	1473.42	1617.73	38.112	1553.01	1565.65	95.344
GAT	1.132	135.90	94.70	98.52	0.406	0.291	911.82	991.31	26.583	1161.29	592.67	55.061
GIN	1.165	175.82	90.66	110.74	0.539	0.691	848.24	1090.36	35.110	1498.23	1364.18	108.331
D-MPNN	0.568	118.42	85.01	86.20	0.441	0.241	423.14	458.39	24.816	470.01	445.91	29.291
PNA	0.681	148.88	88.72	97.31	0.361	0.409	664.98	692.74	23.855	616.70	694.92	57.217
Graphormer	2.836	79.27	54.24	52.42	0.330	0.080	2066.28	2546.01	131.158	2229.88	2525.51	144.595
AWARE	0.297	144.91	133.89	98.86	0.602	0.129	86.62	94.47	22.180	93.59	95.73	5.275
GraphGPS	0.209	75.98	54.75	54.53	0.288	0.089	528.50	693.19	12.488	296.00	411.16	49.888
3D Graph	SchNet	0.060	44.13	27.64	22.55	0.028	0.031	14.19	14.05	0.133	13.93	13.27	1.749
DimeNet++	0.044	36.22	20.01	16.66	0.028	0.022	7.45	6.14	0.323	6.33	7.18	1.118
SE(3)-Trans	0.137	56.52	34.65	34.41	0.050	0.063	65.28	70.70	1.747	68.92	68.88	5.428
EGNN	0.062	49.56	30.08	24.98	0.029	0.030	10.01	9.14	0.089	9.28	9.08	1.519
PaiNN	0.049	42.73	24.46	20.16	0.016	0.025	8.43	7.88	0.169	8.18	7.63	1.419
GemNet-T	0.041	35.46	17.85	15.86	0.021	0.023	7.61	7.08	0.271	6.42	5.88	1.232
SphereNet	0.047	38.93	21.45	18.25	0.027	0.025	8.16	13.68	0.288	6.77	7.43	1.295
SEGNN	0.048	33.61	17.66	17.01	0.021	0.026	11.60	12.45	0.404	11.29	12.20	1.590
Allegro	0.097	102.44	61.86	63.17	0.176	0.032	42.08	44.96	1.977	44.64	44.43	2.949
NequIP	0.066	61.94	42.00	31.64	0.036	0.028	22.08	23.36	0.415	23.23	23.02	1.899
Equiformer	0.051	33.46	17.93	16.85	0.015	0.023	14.49	14.60	0.433	14.88	13.78	2.342
but instead, they should be reflection-antisymmetric [79]. One classic example is that the energy of small molecules is reflection-antisymmetric in a binding system. Each of the two equivariant categories discussed in Sec. 3.3 can solve this problem easily. The spherical frame basis can achieve this by adding the reflection into the Wigner-D matrix [4], and the vector frame basis can accomplish this using the cross-product during frame construction [79].
4 Geometric Datasets and Benchmarks
In Sec. 3, we introduce a novel aspect for understanding symmetry-informed geometric models. In this section, we discuss utilizing Geom3D framework for benchmarking 16 geometric models over 52 tasks. For the detailed dataset acquisitions and task specifications (e.g., dataset size, splitting, and task unit), please check Appendix B. Geom3D also covers 7 1D models and 10 2D graph neural networks (GNNs) and benchmarks the 14 pretraining algorithms to learn a robust geometric representation. Additionally, we want to highlight Geom3D enables exploration of important data preprocessing and optimization tricks for performance improvement, as will be introduced next.
4.1 Small Molecules: QM9