SlowGuess/NAACL_2021_OCR huggingface dataset:
KNOWLEDGE ROUTER: Learning Disentangled Representations for Knowledge Graphs
Shuai Zhang $^{1}$ , Xi Rao $^{1}$ , Yi Tay $^{2}$  and Ce Zhang $^{1}$
$^{1}$ ETH Zurich, Switzerland
2Google Research, USA
The design of expressive representations of entities and relations in a knowledge graph is an important endeavor. While many of the existing approaches have primarily focused on learning from relational patterns and structural information, the intrinsic complexity of KG entities has been more or less overlooked. More concretely, we hypothesize KG entities may be more complex than we think, i.e., an entity may wear many hats and relational triplets may form due to more than a single reason. To this end, this paper proposes to learn disentangled representations of KG entities - a new method that disentangles the inner latent properties of KG entities. Our disentangled process operates at the graph level and a neighborhood mechanism is leveraged to disentangle the hidden properties of each entity. This disentangled representation learning approach is model agnostic and compatible with canonical KG embedding approaches. We conduct extensive experiments on several benchmark datasets, equipping a variety of models (DistMult, SimplE, and QuatE) with our proposed disentangling mechanism. Experimental results demonstrate that our proposed approach substantially improves performance on key metrics.
Knowledge graphs (KG) have emerged as a compelling abstraction for organizing structured knowledge. They have been playing crucial roles in many machine learning tasks. A knowledge graph represents a collection of linked data, describing entities of interest and relationships between them. To incorporate KGs into other machine learning systems, a prevalent way is mapping entities and relations of knowledge graphs into expressive representations in a low-dimensional space that preserves the relationships among objects, also known as knowledge graph embeddings. Representative work such as (Bordes et al., 2013; Wang et al., 2014; Yang et al.,
2014; Sun et al., 2019; Zhang et al., 2019; Chami et al., 2020) has gained intensive attention across the recent years.
The substantial effectiveness of recent work can be attributed to relational pattern modeling in which a suitable relational inductive bias is used to fit the structural information in data. Nevertheless, these methods ignore the fact that the origination and formation of KGs can be rather complex (Ehrlinger and Wöß, 2016). They may be collected, mined, handcrafted or merged in a complicated or convoluted process (Ji et al., 2017; Bosse-lut et al., 2019; Qin et al., 2018). To this end, entities in a knowledge graph may be highly entangled and relational triplets may form and be constructed for various reasons under a plethora of different circumstances or contexts. Contextual reasons and/or domains may be taken into account at the same time. As such, it is only natural that KG embedding methods trained in this fashion would result in highly entangled latent factors. Moreover, the existing holistic approaches fail to disentangle such factors and may result in sub-optimal solutions.
Recently, disentangled representation learning has achieved state-of-the-art performance and attracts much attention in the field of visual representation learning. A disentangled representation should separate the distinct, informative factors of variations in the data (Bengio et al., 2013). Disentangling the latent factors hidden in the observed data can not only increase the robustness, making the model less sensitive to misleading correlations but also enhance the model explainability. Disentanglement can be achieved using either supervised signals or unsupervised approaches. Zhu et al. (Zhu et al., 2014) propose to untangle the identity and view features in a supervised face recognition task. A bilinear model is adopted in (Tenenbaum and Freeman, 2000) to separate content from styles. There is also a large body of work on unsupervised disentangled representation learning (Chen et al.,
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1-10 June 6-11, 2021. ©2021 Association for Computational Linguistics
2016; Denton et al., 2017; Higgins et al., 2016). Generally, the disentanglement mechanism is integrated into unsupervised learning frameworks such as variational autoencoders (Kingma and Welling, 2013) and generative adversarial networks (Goodfellow et al., 2014). The quality of unsupervised disentangled representation can even match that learned from supervised label signals.
Inspired by the success of disentangled representation learning, we seek to enhance the disentanglement capability of entities representation in knowledge graphs. Our hope is that this idea can address the aforementioned challenge in learning entity embeddings, that is, enabling the entities embeddings to better reflect the their inner properties. Unlike learning disentangled representations in visual data, it is more challenging to disentangle the discrete relational data. Most KGs embedding approaches operate at the triplet level, which is uninformative for disentanglement. Intuitively, information about the entities resides largely within the graph encoded through neighborhood structures. Our assumption is that an entity connects with a certain group of entities for a certain reason. For example, Tim Robbins, as an actor, starred in films such as The Shawshank Redemption; as a musician, is a member of the folk music group The Highwaymen. We believe that relational triplets form because of different factors and this can be disentangled when looking it at the graph level.
To summarize, our key contributions are: (1) We propose Knowledge Router (KR), an approach that learns disentangled representations for entities in knowledge graphs. Specifically, a neighbourhood routing mechanism disentangles the hidden factors of entities from interactions with their neighbors. (2) KnowledgeRouter is model agnostic, which means that it can play with different canonical knowledge graph embedding approaches. It enables those models to have the capability in learning disentangled entity representations without incurring additional free parameters. (3) We conduct extensive experiments on four publicly available datasets to demonstrate the effectiveness of KnowledgeRouter. We apply KnowledgeRouter to models such as DistMult, SimplE, and QuatE and observe a notable performance enhancement. We also conduct model analysis to inspect the inner workings of KnowledgeRouter.
2.1 Learning Disentangled Representations
Learning representations from data is the key challenge in many machine learning tasks. The primary posit of disentangled representation learning is that disentangling the underlying structure of data into disjoint parts could bring advantages.
Recently, there is a growing interest in learning disentangled representations across various applications. A trending line of work is integrating disentanglement into generative models. (Tran et al., 2017) propose a disentangled generative adversarial network for face recognition and synthesis. The learned representation is explicitly disentangled from a pose variation to make it pose-invariant, which is critical for face recognition/synthesis task. (Denton et al., 2017) present a disentangled representation learning approach for videos. The proposed approach separates each frame into a time-independent component and a temporal dynamics aware component. As such, it can reflect both the time-invariant and temporal features of a video. (Ma et al., 2018) propose a disentangled generative model for personal image generation. It separates out the foreground, background, and pose information, and offers a mechanism to manipulate these three components as well as control the generated images. Some works (Higgins et al., 2016; Burgess et al., 2018) (e.g.,  $\beta$ -VAE) integrate disentanglement mechanism with variational autoencoder, a probabilistic generative model.  $\beta$ -VAE uses a regularization coefficient  $\beta$  to constrain the capacity of the latent information channel. This simple modification enables latent representations to be more factorised.
Drawing inspiration from the vision community, learning disentangled representations has also been investigated in areas such as natural language processing and graph analysis. (Jain et al., 2018) propose an autoencoders architecture to disentangle the populations, interventions, and outcomes in biomedical texts. (Liu et al., 2019) propose a prism module for semantic disentanglement in named entity recognition. The prism module can be easily trained with downstream tasks to enhance performance. For graph analysis, (Ma et al., 2019a) propose to untangle the node representation of graph-structured data in graph neural networks. (Ma et al., 2019b) present a disentangled variational autoencoder to disentangle the user's diverse interests for recommender systems.
2.2 Knowledge Graph Embeddings
Learning effective representations for knowledge graphs is extensively studied because of its importance in downstream tasks such as knowledge graph completion, natural language understanding, web search, and recommender systems. Among the large body of related literature, two popular lines are translational approaches and semantic matching approaches. The groundbreaking TransE (Bordes et al., 2013) sets the fundamental paradigm for translational models. Typically, the aim is to reduce the distance between translated (by relation) head entity and tail entity. Successors such as TransH (Wang et al., 2014), TransR (Lin et al., 2015) all follow this translational pattern. Semantic matching methods calculate the semantic similarities between entities. A representative semantic model is DistMult (Yang et al., 2014) which measures the plausibility of triplets with vector multiplications. To model more complex relation patterns, (Trouillon et al., 2016; Zhang et al., 2019; Sun et al., 2019; Zhang et al., 2021) extend the embedding spaces to complex number space or hyperbolic space. A fully expressive model named SimpleE (Kazemi and Poole, 2018) could achieve the same level of capability of ComplEx (Trouillon et al., 2016) with lower calculation cost.
Inspired by the success of disentangled representations, we explore methods to factorize different components/aspects of entangled entities in a knowledge graph. To the best of our knowledge, our work is one of the first efforts to induce disentangled representations in knowledge graphs. Our disentangled embedding algorithm can be easily integrated into existing knowledge graph embedding models (model agnostic).
3 The Proposed Knowledge Router
3.1 Notation and Problem Formulation
Suppose we have an entity set  $\mathcal{E}$  and a relation set  $\mathcal{R}$ , where  $|\mathcal{E}| = N$  and  $|\mathcal{R}| = M$ . A knowledge graph  $\mathcal{G} = (\mathcal{E}, \mathcal{R})$  is made up of a collection of facts  $\mathcal{F}$  in triplet form  $(h, r, t)$ , where  $h, t \in \mathcal{E}$  and  $r \in \mathcal{R}$ . The triplet  $(h, r, t) \in \mathcal{F}$  means that entities  $h$  and  $r$  are connected via a relation  $r$ . The facts are usually directional, which means exchanging the head entity and tail entity does not necessarily result in a legitimate fact.
We are concerned with the link prediction task. The goal is to embed the entities and relations of a knowledge graph into low-dimensional rep
Notation	Description
E	Entity set.
R	Relation set.
E	The entity embedding matrix.
W	The relation embedding matrix.
Ee	The ethrow of the entity embedding matrix.
Wr	The rthrow of the relation embedding matrix.
d	The length of the embedding vector.
N(e)	Neighbourhood entities set of entity e.
K	The number of independent components.
T	The number of routing iterations.
xe,k	The kthinitial vector for entity e.
pe,k	The kthvector of entity e after disentanglement.
Se,i,k	The similarity score between entity eand entity i w.r.t the kthcomponent.
wi,k	The extent to which the model attends tothe kthcomponent of entity i.
Table 1: The notations and denotations.
resentations that can preserve the facts in the graph. A classical setting is using an embedding matrix  $\mathbf{E} \in \mathbb{R}^{N \times d}$  to represent all the entities and an embedding matrix  $\mathbf{W} \in \mathbb{R}^{M \times d}$  to represent all the relations.
3.2 Disentangled Knowledge Graph Embeddings
Instead of directly modeling triplet facts, we propose to disentangle the entities with their neighbors in a message passing setting. The neighborhood entities could form several clusters for different reasons and the entity is updated by the information accepted from its neighborhood clusters.
Figure 1 illustrates the overall process of Knowledge Router. It consists of two stages: (1) disentangling the entities from a graph perspective using neighbourhood routing; (2) scoring the facts using relations and the disentangled entities representations.
Let us build an undirected graph from the training data. The relations are anonymized, which means we do not need to know under which conditions two entities are linked. We denote the neighbourhood of entity  $e$  as  $\mathcal{N}(e)$ , regardless of the relations. Our neighborhood routing approach operates on this graph.
Given an entity  $e$ , we aim to learn a disentangled embedding that encodes various attributes of the entity. In this regard, we suppose that each entity is composed of  $K$  independent components, with each component denoted by  $\mathbf{p}_{e,k} \in \mathbb{R}^{\frac{d}{K}}$ , where  $\forall k = 1,2,\dots,K$ . Each component stands for one aspect of the entity, e.g., a role of a person. A
Figure 1: The overall procedure of the proposed Knowledge Router algorithm for learning disentangled entity representations. In this example, we disentangle the entity embedding into four components  $(K = 4)$  via neighborhood routing (iterate  $T$  times). These components are then concatenated to represent the corresponding entity.
major challenge here is to make the learned  $K$  components to be independent of one another so that different facets can be separately encoded. To this end, we adopt routing mechanisms that are inspired by capsule networks (Hinton et al., 2011). Specifically, we aim to learn the  $K$  components from both the entity  $e$  and its neighbourhoods  $\mathcal{N}(e)$ . Next, we describe this procedure in detail.
For each entity  $e$ , we first initialize the  $\mathbf{E}_e$  randomly and evenly split it into  $K$  parts. The  $k^{\mathrm{th}}$  part is denoted by  $\mathbf{x}_{e,k} \in \mathbb{R}^{\frac{d}{K}}$ . By doing so, the embedding is projected into different subspaces. To ensure computation stability, each part is also normalized as follows:
$$
\mathbf {x} _ {e, k} = \frac {\mathbf {x} _ {e , k}}{\| \mathbf {x} _ {e , k} \| _ {2}} \tag {1}
$$
This is used for the initialization of  $\mathbf{p}_{e,k}$ . Obviously, the information contained is limited and it cannot reach the goal of disentanglement. To enrich the information, we use a graph message passing mechanism and define the update rule for the  $k^{\mathrm{th}}$  component of  $\mathbf{p}_e$  as follows:
$$
\mathbf {p} _ {e, k} = \mathbf {x} _ {e, k} + \text {A G G R E G A T E} \left(\left\{\mathbf {x} _ {i, k}, \forall i \in \mathcal {N} (e) \right\}\right), \tag {2}
$$
where AGGREGATE represents the neighborhood aggregation function (defined in equation 5). The same  $\ell_2$  normalization as (1) is applied to  $\mathbf{p}_{e,k}$  afterwards.
In this way,  $\mathbf{p}_{e,k}$  contains information from the  $k^{\mathrm{th}}$  aspect of both entity  $e$  and all of its neighbors. Common aggregating functions such as mean pooling and sum pooling are viable, but treating
each neighbor equally when determining one component of the representation is undoubtedly not sensible. As such, an attention mechanism is used to obtain weights for each neighbor. In particular, a scaled dot-product attention method is applied. We first get the dot product between  $\mathbf{p}_{e,k}$  and  $\mathbf{x}_{i,k},\forall i\in \mathcal{N}(e)$ . For each  $k$ , we get the following similarity score:
$$
s _ {e, i, k} = \frac {\mathbf {p} _ {e , k} ^ {\top} \mathbf {x} _ {i , k}}{\sqrt {d / k}}, \tag {3}
$$
which provides information on how entity  $e$  interacts with its neighbour entity  $i$  pertaining to the aspect  $k$ . Then the softmax function is applied to get the weight distribution over different components for each neighbour.
$$
w _ {i, k} = \frac {\exp \left(s _ {e , i , k}\right)}{\sum_ {k = 1} ^ {K} \exp \left(s _ {e , i , k}\right)}, \tag {4}
$$
and  $w_{i,k}$  indicates the extent to which the model attends to the  $k^{\mathrm{th}}$  component of entity  $i$ .
Now, we formulate the definition of the AGGREGATE function as follows:
$$
\text {A G G R E G A T E} (\{\mathbf {x} _ {i, k}, \forall i \in \mathcal {N} (e) \}) := \sum_ {i \in \mathcal {N} (e)} w _ {i, k} \mathbf {x} _ {i, k} \tag {5}
$$
The above process, including equations (2), (3), (4), (5) for learning  $\mathbf{p}_{e,k},\forall k = 1,2,\dots,K$  is repeated for  $T$  iterations, which is the same as that of a routing mechanism. Like capsule networks (Sabour et al., 2017), we also assume that entity (object) is composed of entity (object) parts. This routing method enables it to model part-whole
relationships and enlarge the differences between parts after several routing iterations.
Afterwards, the concatenation of all  $K$  components of an entity is used to represent that entity. That is, the disentangled representation  $\mathbf{p}_e$  of the entity  $e$  is defined as:
$$
\mathbf {p} _ {e} = \left[ \mathbf {p} _ {e, 1}, \mathbf {p} _ {e, 2}, \dots , \mathbf {p} _ {e, K} \right] \tag {6}
$$
This neighborhood routing algorithm is model agnostic as our aim is to learn an entity embedding matrix which is necessary for most knowledge graph embedding methods. It is worth noting that this model will not introduce additional free parameters to the model.
The intuition behind the "routing mechanism" is that each facet in an entity has a separate route to contribute to the meaning of this entity. The routing algorithm will coordinately infer  $\mathbf{p}_{e,k}$  (we can view it as the center of each cluster) and  $w_{i,k}$  (the probability that factor  $k$  is the reason why entity  $e$  is connected with entity  $i$ ). They are coordinately learned and under the constraint that each neighbor should belong to one cluster. It is reminiscent of the iterative method used in the EM algorithm (Bishop, 2006) and is expected to lead to convergence and meaningful disentangled representations (Ma et al., 2019a).
Until now, the relation embeddings are not utilized as all relations are anonymous during graph construction. This algorithm will be jointly trained with the following facts scoring algorithms.
3.3 Facts Scoring using Disentangled Entities
Using disentangled entity embeddings alone cannot recover the facts in a knowledge graph. It shall be further updated simultaneously with the relation embeddings for the fact scoring process. To predict whether a triplet  $\langle h, r, t \rangle$  holds or not, we first fetch the learned disentangled representation of the head and tail entities,  $\mathbf{p}_h$  and  $\mathbf{p}_t$ . Then we adopt three methods for triplet scoring including DistMult (Yang et al., 2014), SimpleE (Kazemi and Poole, 2018), and QuatE (Zhang et al., 2019). We denote the model after disentanglement as: KR-DistMult, KR-SimpleE, and KR-QuatE.
The scoring function of KR-DistMult is defined as follows:
$$
\phi (h, r, t) = \left\langle \mathbf {W} _ {r}, \mathbf {p} _ {h}, \mathbf {p} _ {t} \right\rangle \tag {7}
$$
where  $\langle *,*,* \rangle$  denotes the standard component-wise multi-linear dot product.
SimpleE needs an additional entity embedding matrix  $\mathbf{H} \in \mathbb{R}^{N \times d}$  and an additional relation embedding matrix  $\mathbf{V} \in \mathbb{R}^{M \times d}$ . We perform the same disentanglement process on  $\mathbf{H}$  and denote the disentangled representation of entity  $e$  as  $\mathbf{q}_e$ , the scoring function of KR-SimplE (SimpleE-avg is adopted since it outperforms SimplE-ignr) is:
$$
\phi (h, r, t) = \left(\left\langle \mathbf {W} _ {r}, \mathbf {p} _ {h}, \mathbf {q} _ {t} \right\rangle + \left\langle \mathbf {V} _ {r}, \mathbf {q} _ {h}, \mathbf {p} _ {t} \right\rangle\right) \times \frac {1}{2} \tag {8}
$$
For QuatE, entities and relations are represented with quaternions. Each quaternion is composed of a real component and three imaginary components. Let  $\mathbf{Q} \in \mathbb{H}^{N \times d}$  denote the quaternion entity embedding and  $\mathbf{W} \in \mathbb{H}^{M \times d}$  denote the quaternion relation embedding, where  $\mathbb{H}$  is the quaternion space. Each entity is represented by  $\mathbf{Q}_e$ . We apply the Knowledge Router algorithm on each component of  $\mathbf{Q}_e$ . The scoring function of KR-QuatE is:
$$
\phi (h, r, t) = \mathbf {Q} _ {h} ^ {\mathrm {K R}} \otimes \frac {\mathbf {W} _ {r}}{| \mathbf {W} _ {r} |} \cdot \mathbf {Q} _ {t} ^ {\mathrm {K R}} \tag {9}
$$
where “ $\otimes$ ” is Hamilton product; “.” represents the quaternion inner product;  $\mathbf{Q}^{\mathrm{KR}}$  denotes the entity representation after disentanglement.
As Knowledge Router is model agnostic, other scoring functions are also applicable.
3.4 Objective Functions
To learn a disentangled KG model, we adopt the following negative log-likelihood loss:
$$
\mathcal {L} = - \frac {1}{S} \sum_ {i = 1} ^ {S} \left(y ^ {(i)} \log \left(\phi^ {(i)}\right) + \left(1 - y ^ {(i)}\right) \log \left(1 - \phi^ {(i)}\right)\right) \tag {10}
$$
where  $S$  is the number of training samples (triplets);  $y^{(i)}$  is a binary label indicating whether the  $i^{\mathrm{th}}$  triplet holds or not;  $\phi^{(i)}$  is the prediction for the  $i^{\mathrm{th}}$  triplet. Our model can be trained with commonly used minibatch gradient descent optimizers.
3.5 Complexity Analysis
The disentanglement process of each node needs  $\mathcal{O}(|\mathcal{N}(e)|\frac{d}{K} K + T(|\mathcal{N}(e)|\frac{d}{K} K + \frac{d}{K} K))$  time complexity, where  $|\mathcal{N}(e)|$  is neighborhood size. After simplification, the time complexity is  $\mathcal{O}(T|\mathcal{N}(e)|d)$ . This will not incur a high computational cost since  $T$  is usually a small number (e.g., 3), and the neighborhood size is determined by the average degree and can usually be constrained by a constant value (e.g., 10). With regard to fact
Datasets	N	M	|train|	|validation|	|test|
FB15k-237	14,541	237	272,115	17,535	20,466
WIKIDATA	11,153	96	53,252	11,894	11,752
ICEWS14	7,128	230	42,690	7,331	7,419
ICEWS05-15	10,488	251	368,962	46,275	46,092
Table 2: Statistics of datasets used in our experiments.
scoring, it requires  $\mathcal{O}(d)$  time complexity for each triplet in general.
In this section, we conduct experiments on several benchmark datasets to verify the effectiveness of the proposed approach. We target at answering: RQ I: whether the disentanglement method can enhance the traditional knowledge graph embedding methods? RQ II: Model-agnosticism: can it effectively work with different baseline models? RQ III: How do certain important hyper-parameters impact the model performance and what has the disentanglement algorithm learned? Are they meaningful?
4.1 Datasets Description
We use four publicly available datasets including ICEWS14, ICEWS05-15, WikiData, and FB15k-237. The reason for using these is that their entities are complicated and highly entangled. The WordNet dataset is not appropriate to evaluate the proposed method as the entities in WordNet are already disentangled<sup>1</sup>.
FB15k-237 is a subset of the Freebase knowledge base which contains general information about the world. We adopt the widely used version generated by (Dettmers et al., 2018) where inverse relations are eliminated to avoid data leakage.
WikiData is sampled from Wikidata², a collaborative open knowledge base. The knowledge is relatively up-to-date compared with FB15k-237. We use the version provided by (García-Durán et al., 2018). Timestamp is discarded.