SlowGuess/NAACL_2021_OCR huggingface dataset:
Xutan Peng Guanyi Chen Chenghua Lin \* Mark Stevenson
$\diamond$  Department of Computer Science, The University of Sheffield
Department of Information and Computing Sciences, Utrecht University {x.peng, c.lin, mark.stevenson}@shef.ac.uk g.chen@uu.nl
Knowledge Graph Embeddings (KGEs) have been intensively explored in recent years due to their promise for a wide range of applications. However, existing studies focus on improving the final model performance without acknowledging the computational cost of the proposed approaches, in terms of execution time and environmental impact. This paper proposes a simple yet effective KGE framework which can reduce the training time and carbon footprint by orders of magnitudes compared with state-of-the-art approaches, while producing competitive performance. We highlight three technical innovations: full batch learning via relational matrices, closed-form Orthogonal Procrustes Analysis for KGEs, and non-negative-sampling training. In addition, as the first KGE method whose entity embeddings also store full relation information, our trained models encode rich semantics and are highly interpretable. Comprehensive experiments and ablation studies involving 13 strong baselines and two standard datasets verify the effectiveness and efficiency of our algorithm.
The recent growth in energy requirements for Natural Language Processing (NLP) algorithms has led to the recognition of the importance of computationally cheap and eco-friendly approaches (Strubell et al., 2019). The increase in computational requirements can, to a large extent, be attributed to the popularity of massive pre-trained models, such as Language Models (e.g., BERT (Devlin et al., 2019) and GPT-3 (Brown et al., 2020)) and Knowledge Graph Embeddings (KGEs, e.g., SACN (Shang et al., 2019)), that require significant resources to train. A number of solutions have been proposed such as reducing the number of parameters the model contains. For instance, Sanh et al. (2019) introduced a distilled version of BERT and
Zhang et al. (2019) decreased the parameters used for training KGEs with the help of the quaternion. In contrast with previous work, this paper explores algorithmic approaches to the development of efficient KGE techniques.
Knowledge Graphs are core to many NLP tasks and downstream applications, such as question answering (Saxena et al., 2020), dialogue agents (He et al., 2017), search engines (Dong et al., 2014) and recommendation systems (Guo et al., 2020). Facts stored in a knowledge graph are always in the format of tuples consisting of one head entity, one tail entity (both are nodes in knowledge graphs) and a relation (an edge in knowledge graphs) between them. KGEs learn representations of relations and entities in a knowledge graph, which are then utilised in downstream tasks like predicting missing relations (Bordes et al., 2013; Sun et al., 2019; Tang et al., 2020). The application of deep learning has led to significant advances in KGE (Rossi et al., 2021). Nonetheless, such approaches are computationally expensive with associated environmental costs. For example, training the SACN model (Shang et al., 2019) can lead to emissions of more than  $5.3\mathrm{kgCO}_{2}$  (for more data of other algorithms, see Tab. 2).
To alleviate the computational cost we introduce ProCRUSTES, a lightweight, fast, and eco-friendly KGE training technique. ProCRUSTES is built upon three novel techniques. First, to reduce the batch-wise computational overhead, we propose to parallelise batches by grouping tuples according to their relations, which ultimately enables efficient full batch learning. Second, we turn to a closed-form solution for Orthogonal Procrustes Problem to boost the embedding training, which has never been explored in the context of KGEs. Third, to break though the bandwidth bottleneck, our algorithm is allowed to be trained without negative samples.
To verify the effectiveness and efficiency of our proposed method, we benchmark two popu
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2364-2375 June 6-11, 2021. ©2021 Association for Computational Linguistics
Figure 1: The by-relation partitioning architecture of PROCRUSTES for a toy graph (left). Matrices involved in the computation of Eq. (1) are divided into two relational matrices: the upper is for relation 1 (dashed) and the lower is for relation 2 (solid).
T2,1	T2,2	...	T2,d/ds
e2,1	e2,2	...	e2,d/ds
e4,1	e4,2	...	e4,d/ds
e5,1	e5,2	...	e5,d/ds
e4,1	e4,2	...	e4,d/ds
e6,1	e6,2	...	e6,d/ds
lar datasets (WN18RR and FB15k-237) against 13 strong baselines. Experimental results show that PROCRUSTEs yields performance competitive with the state-of-the-art while also reducing training time by up to  $98.4\%$  and the carbon footprint by up to  $99.3\%$ . In addition, we found that our algorithm can produce easily interpretable entity embeddings with richer semantics than previous approaches. Our code is available at https://github.com/Pzoom522/ProcrustEs-KGE.
Our contribution is three-fold: (1) We introduce three novel approaches to substantially reduce computational overhead of embedding large and complex knowledge graphs: full batch learning based on relational matrices, closed-form Orthogonal Procrustes Analysis for KGEs, and non-negative-sampling training. (2) We systemically benchmark the proposed algorithm against 13 strong baselines on two standard datasets, demonstrating that it retains highly competitive performance with just order-of-minute training time and emissions of less than making two cups of coffee. (3) We successfully encode both entity and relation information in a single vector space for the first time, thereby enriching the expressiveness of entity embeddings and producing new insights into interpretability.
We propose a highly efficient and lightweight method for training KGEs called PROCUSTES, which is more efficient in terms of time consumption and  $\mathrm{CO}_{2}$  emissions than previous counterparts
by orders of magnitude while retaining strong performance. This is achieved by introducing three novel optimisation strategies, namely, relational mini-batch, closed-form Orthogonal Procrustes Analysis, and non-negative sampling training.
2.1 Preliminaries: Segmented Embeddings
Our proposed PROCRUSTES model is built upon segmented embeddings, a technique which has been leveraged by a number of promising recent approaches to KGE learning (e.g., RotatE (Sun et al., 2019), SEEK (Xu et al., 2020), and OTE (Tang et al., 2020)). In contrast to conventional methods for KGEs where each entity only corresponds to one single vector, algorithms adopting segmented embeddings explicitly divide the entity representation space into multiple independent sub-spaces. During training each entity is encoded as a concatenation of decoupled sub-vectors (i.e., different segments, and hence the name). For example, as shown in Fig. 1, to encode a graph with 7 entities, the embedding of the  $t$ th entity is the row-wise concatenation of its  $d / d_{s}$  sub-vectors (i.e.,  $e_{t,1}\widehat{e}_{t,2}\widehat{\dots}\widehat{e}_{t,d / d_s}$ ), where  $d$  and  $d_{s}$  denote the dimensions of entity vectors and sub-vectors, respectively. Employing segmented embeddings permits parallel processing of the structurally separated sub-spaces, and hence significantly boosts the overall training speed. Furthermore, segmented embeddings can also enhance the overall expressiveness of our model, while substantially reducing the dimension of matrix calculations. We provide detailed discussion on the empirical influence of segmented embedding setups in § 3.4.
2.2 Efficient KGE Optimisation
Full batch learning via relational matrices. Segmented embeddings can speed up training process by parallelising tuple-wise computation. In this section, we propose a full batch learning technique via relational matrices, which can optimise batch-wise computation to further reduce training time. This idea is motivated by the observation that existing neural KGE frameworks all perform training based on random batches constructed from tuples consisting of different types of relations (Bordes et al., 2013; Trouillon et al., 2016; Schlichtkrull et al., 2018; Chami et al., 2020). Such a training paradigm is based on random batches which, although straightforward to implement, is difficult to parallelise. This is due to the nature of computer process scheduling: during the interval between a process reading and updating the relation embeddings, they are likely to be modified by other processes, leading to synchronisation errors and consequently result in unintended data corruption, degraded optimisation, or even convergence issues.
To tackle this challenge, we propose to construct batches by grouping tuples which contain the same relations. The advantage of this novel strategy is two-fold. For one thing, it naturally reduces the original tuple-level computation to simple matrix-level arithmetic. For another and more importantly, we can then easily ensure that the embedding of each relation is only accessible by one single process. Such a training strategy completely avoids the data corruption issue. In addition, it makes the employment of the full batch learning technique (via relational matrices) possible, which offers a robust solution for parallelising the KGEs training process and hence can greatly enhance the training speed. To the best of our knowledge, this approach has never been explored by the KGE community.
As illustrated in Fig. 1, we first separate the embedding space into segments (cf. § 2.1) and arrange batches based on relations. After that, for each training step, the workflow of PROCRUSTES is essentially decomposed into  $m \times d / d_s$  parallel optimisation processes, where  $m$  is the number of relation types. Let  $i$  and  $j$  denote the indices of relation types and sub-spaces, respectively, then the column-wise concatenations of the  $j$ th sub-vectors of all tuples of  $i$ th relations can be symbolised as  $H_{i,j}$  (for head entities) and  $T_{i,j}$  (for tail entities). Similarly,  $R_{i,j}$  denotes the corresponding relation embedding matrix in the  $j$ th sub-space. The final
objective function of PROCRUSTES becomes
$$
\mathcal {L} = \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {d / d _ {s}} | | H _ {i, j} R _ {i, j} - T _ {i, j} | | _ {2}. \tag {1}
$$
Orthogonal Procrustes Analysis. Our key optimisation objective, as formulated in Eq. (1), is to minimise the Euclidean distance between the head and tail matrices for each parallel process. In addition, following Sun et al. (2019) and Tang et al. (2020), we restrict the relation embedding matrix  $R_{i,j}$  to be orthogonal throughout model training, which has been shown effective in improving KGE quality. Previous KGE models use different approaches to impose orthogonality. For instance, RotatE (Sun et al., 2019) takes advantage of a corollary of Euler's identity and defines its relation embedding as
$$
R _ {i, j} = \left[ \begin{array}{l l} \cos \theta_ {i, j} & \sin \theta_ {i, j} \\ - \sin \theta_ {i, j} & \cos \theta_ {i, j} \end{array} \right], \qquad (2)
$$
which is controlled by a learnable parameter  $\theta_{i,j}$ . Although Eq. (2) holds orthogonality and retains simplicity, it is essentially a special case of segmented embedding where  $d_{s}$  equals 2. As a result,  $R_{i,j}$  is always two-dimensional, which greatly limits the modelling capacity (see § 3.4 for discussion on the impact of dimensionality). To overcome this limitation, OTE (Tang et al., 2020) explicitly orthogonalises  $R_{i,j}$  using the Gram-Schmidt algorithm per back-propagation step (see Appendix A for details). However, while this scheme works well for a wide range of  $d_{s}$  (i.e., the dimension for the sub-vector), similar to RotatE, OTE finds a good model solution based on gradient descent, which is computationally very expensive.
We address the computational issue by proposing a highly efficient method utilising the proposed parallelism of full batch learning. With full batch learning, comparing with existing methods which deal with heterogeneous relations, PROCRUSTES only needs to optimise one single  $R_{i,j}$  in each process, which becomes a simple constrained matrix regression task. More importantly, through Singular Value Decomposition (SVD), we can derive an closed-form solution (Schönemann, 1966) as
$$
R _ {i, j} ^ {\star} = U V ^ {\top}, \mathrm {w} / U \Sigma V ^ {\top} = \operatorname {S V D} \left(H _ {i, j} ^ {\top} T _ {i, j}\right), \tag {3}
$$
where  $R_{i,j}^{\star}$  denotes the optima. During each iteration, PROCRUSTES can directly find the globally
optimal embedding for each relation given the current entity embeddings by applying Eq. (3). Then, based on the calculated  $\mathcal{L}$ , PROCRUSTES updates entity embeddings through the back propagation mechanism (NB: the relation embeddings do not require gradients here). This process is repeated until convergence. As the optimisation of relation embeddings can be done almost instantly per iteration thanks to the closed-form Eq. (3), PROCRUSTES is significantly (orders of magnitude) faster than RotatE and OTE. In addition, compared with entity embeddings of all other KGE models which are updated separately with relation embedding, entity embeddings trained by PROCRUSTES can be used to restore relation embeddings directly (via Eq. (3)). In other words, PROCRUSTES can encode richer information in the entity space than its counterparts (see § 3.5).
Further optimisation schemes. As recently surveyed by Ruffinelli et al. (2020), existing KGE methods employ negative sampling as a standard technique for reducing training time, where update is performed only on a subset of parameters by calculating loss based on the generated negative samples. With our proposed closed-form solution (i.e., Eq. (3)), computing gradients to update embeddings is no longer an efficiency bottleneck for PROCRUSTES. Instead, the speed bottleneck turns out to be the extra bandwidth being occupied due to the added negative samples. Therefore, for PROCRUSTES, we do not employ negative sampling but rather update all embeddings during each round of back propagation with positive samples only, in order to further optimise the training speed (see Appendix B for bandwidth comparisons against baselines which adopts negative sampling).
We also discovered that if we do not apply any additional conditions during training, PROCUSTES tends to fall into a trivial optimum after several updates, i.e.,  $\mathcal{L} = 0$ , with all values in  $H_{i,j}$ ,  $T_{i,j}$  and  $R_{i,j}$  being zero. In other words, the model collapses with nothing encoded at all. This is somewhat unsurprising as such trivial optima often yields large gradient and leads to this behaviour (Zhou et al., 2019). To mitigate this degeneration issue, inspired by the geometric meaning of orthogonal  $R_{i,j}$  (i.e., to rotate  $H_{i,j}$  towards  $T_{i,j}$  around the coordinate origin, without changing vector length), we propose to constrain all entities to a high-dimensional hypersphere by performing two spherisation steps in every epoch. The first
FB15k-237	WN18RR
Entities	14,541	40,943
Relations	237	11
Train samples	272,115	86,835
Validate samples	17,535	3,034
Test samples	20,466	3,134
Table 1: Basic statistics of the two benchmark datasets.
technique, namely centring, respectively translates  $H_{i,j}$  and  $T_{i,j}$  so that the column-wise sum of each matrix becomes a zero vector (note that each row denotes a sub-vector of an entity). The second operation is length normalisation, which ensures the row-wise Euclidean norm of  $H_{i,j}$  and  $T_{i,j}$  to always be one. Employing these two simple constraints effectively alleviates the trivial optimum issue, as evidenced in our experiments (see § 3).
3.1 Setups
We assess the performance of PROCRUSTES on the task of multi-relational link prediction, which is the de facto standard of KGE evaluation.
Datasets. In this study, following previous works (e.g., baselines in Tab. 2), we employ two benchmark datasets for link prediction: (1) FB15K-237 (Toutanova and Chen, 2015), which consists of sub-graphs extracted from Freebase, and contains no inverse relations; and (2) WN18RR (Dettmers et al., 2018), which is extracted from WordNet. Tab. 1 shows descriptive statistics for these two datasets, indicating that FB15K-237 is larger in size and has more types relations while WN18RR has more entities. We use the same training, validating, and testing splits as past studies.
Evaluation metrics. Consistent with Sun et al. (2019) and Tang et al. (2020), we report Hit Ratio with cut-off values  $n = 1,3,10$  (i.e., H1, H3, and H10) and Mean Reciprocal Rank (MRR). Additionally, as to efficiency, we report the time cost and  $\mathrm{CO}_{2}$  emissions for each model, i.e., from the beginning of training until convergence.
Baselines. We compare PROCRUSTES to not only classical neural graph embedding methods, including TransE (Bordes et al., 2013), Dist-Multi (Yang et al., 2015), and ComplEx (Trouillon et al., 2016), but also embedding techniques recently reporting state-of-the-art performance on either WN18RR or FB15k-237, including R-GCN (Schlichtkrull et al., 2018), ConvE (Dettmers
WN18RR	FB15k-237
MRR	H1	H3	H10	MRR	H1	H3	H10
TransE (2013)	.226	-	-	.501	85	367	.294	-	-	.465	96	370
DistMult (2015)	.430	.390	.440	.490	79	309	.241	.155	.263	.419	91	350
ComplEx (2016)	.440	.410	.460	.510	130	493	.247	.158	.275	.428	121	534
R-GCN (2018)	.417	.387	.442	.476	138	572	.248	.151	.264	.417	152	598
ConvE (2018)	.430	.400	.440	.520	840	3702	.325	.237	.356	.501	1007	4053
A2N (2019)	.450	.420	.460	.510	203	758	.317	.232	.348	.486	229	751
SACN (2019)	.470	.430	.480	.540	1539	5342	.352	.261	.385	.536	1128	4589
TuckER (2019)	.470	.443	.482	.526	173	686	.358	.266	.392	.544	184	704
QuatE (2019)	.488	.438	.508	.582	176	880	.348	.248	.382	.550	180	945
InteractE (2020)	.463	.430	-	.528	254	1152	.354	.263	-	.535	267	1173
RotH (2020)	.496	.449	.514	.586	192	903	.344	.246	.380	.535	207	1120
RotatE (2019)	.439	.390	.456	.527	255	823	.297	.205	.328	.480	343	1006
OTE (2020)	.448	.402	.465	.531	304	1008	.309	.213	.337	.483	320	1144
PROCRUSTES (ours)	.453	.408	.491	.549	14	37	.295	.241	.310	.433	9	42
w/ NS (ours)	.457	.411	.494	.551	44	124	.302	.245	.333	.465	37	159
w/ TB (ours)	.468	.417	.498	.557	92	268	.326	.247	.354	.492	56	243
w/ NS+TB (ours)	.474	.421	.502	.569	131	346	.345	.249	.379	.541	85	285
Table 2: Model effectiveness and efficiency on link prediction benchmarks.  $\widehat{\mathbf{c}}$ : training time (minutes);  $\widehat{\mathbf{c}}_{\mathbf{0}_{2}}$ : carbon dioxide production (grams). NS: negative sampling; TB: traditional batch. The performance results of baselines are coloured heavily and lightly if they are below those of PROCRUSTES and "w/ NS+TB", respectively. State-of-the-art scores are in bold. Following Balazevic et al. (2019) and Zhang et al. (2019), for fair comparison, both RotatE and OTE results are reported with conventional negative sampling rather than the self-adversarial one.
et al., 2018), A2N (Bansal et al., 2019), RotatE (Sun et al., 2019), SACN (Shang et al., 2019), TuckER (Balazevic et al., 2019), QuatE (Zhang et al., 2019), InteractE (Vashisht et al., 2020), OTE (Tang et al., 2020), and RotH (Chami et al., 2020). For all these baselines, we use the official code and published hyper-parameters to facilitate reproducibility.
Implementation details. All experiments are conducted on a workstation with one NVIDIA GTX 1080 Ti GPU and one Intel Core i9-9900K CPU, which is widely applicable to moderate industrial/academic environments. We use the Experiment Impact Tracker (Henderson et al., 2020) to benchmark the time and carbon footprint of training. To reduce measurement error, in each setup we fix the random seeds, run PROCRUSTES and all baselines for three times and reported the average.
The key hyper-parameters of our model is  $d$  and  $d_{s}$ , which are respectively set at 2K and 20 for both datasets. The detailed selection process is described in § 3.4. We train each model for a maximum of 2K epochs and check if the validation MRR stops increasing every 100 epochs after 100 epochs. For WN18RR and FB15k-237 respectively, we report the best hyperparameters as fixed learning rates of 0.001 and 0.05 (Adam optimiser), and stopping epochs of 1K and 200.
Tab. 2 reports the results of both our PROCRUSTES and all other 13 baselines on both WN18RR and FB15k-237 datasets. We analyse these results from two dimensions: (1) Effectiveness: the model performance on link prediction task (MRR is our main indicator); (2) Efficiency: system training time and carbon footprint (i.e.,  $\mathrm{CO}_{2}$  emissions).
Regarding the performance on WN18RR, we found that PROCRUSTES performs as good as or even better than previous state-of-the-art approaches. To be concrete, out of all 13 baselines, it beats 11 in H10, (at least) 9 in H3 and 8 in MRR. The models outperformed by PROCRUSTES include not only all methods prior to 2019, but also several approaches published in 2019 or even 2020. Notably, when compared with the RotatE and OTE, two highly competitive methods which have similar architectures to PROCRUSTES (i.e., with segmented embeddings and orthogonal constraints), our PROCRUSTES can learn KGEs with higher quality (i.e., 0.014 and 0.005 higher in MRR, respectively). This evidences the effectiveness of the proposed approaches in § 2 in modelling knowledge tuples.
While PROCRUSTES achieves very competitive performance, it requires significantly less time for training: it converges in merely 14 minutes, more than 100 times faster than strong-performing counterparts such as SACN. Moreover, it is very envi
Avg. MRR gain per minute of training
Figure 2: Unified effectiveness-efficiency comparison between most competitive KGE models in Tab. 2. The left and right sub-figures are respectively for WN18RR and FB15k-237.
ronmentally friendly: from bootstrapping to convergence, PROCIRSTES only emits  $37\mathrm{g}$  of  $\mathrm{CO}_{2}$ , which is even less than making two cups of coffee<sup>1</sup>. On the contrary, the baselines emit on average  $1469\mathrm{g}$  and up to  $5342\mathrm{g}$ $\mathrm{CO}_{2}$ : the latter is even roughly equal to the carbon footprint of a coach ride from Los Angeles to San Diego<sup>2</sup>.
As for the testing results on FB15k-237, we found that although PROCRUSTES seems less outstanding (we investigate the reasons in § 3.3), it still outperforms at least 7 more complex baselines in H1 and almost all models prior to 2019 in MRR. Furthermore, similar to the observation on WN18RR, it demonstrates great advantage in terms of efficiency. While all baselines need 91 to 1128 minutes to coverage with  $350\mathrm{g}$  to  $4589\mathrm{gCO_2}$  produced, PROCRUSTES can learn embeddings of similar quality in just 9 minutes and with  $42\mathrm{g}$  emissions. By employing both traditional batch and negative sampling, we show that PROCRUSTES can achieve near-state-of-the-art performance on both datasets. We discuss this in detail in § 3.3.
To provide a unified comparisons between PROCRUSTES and the most strong-performing baselines on both effectiveness and efficiency, we further investigate the following question: How much performance gain can we obtain by spending unit time on training or making unit emissions? We did analysis by calculating MRR/(training time) and MRR/(carbon footprint) and the results are presented in Fig. 2. It is obvious that among all competitive KGE models, PROCRUSTES is the most economic algorithm in terms of performance-cost trade-off: it is more than 20 times more efficient than any past works, in terms of both performance per unit training time and per unit  $\mathrm{CO}_{2}$  emissions.
We also investigate baseline performance with a