IMPLICIT REASONET: MODELING LARGE-SCALE STRUCTURED RELATIONSHIPS WITH SHARED MEM- ORY
Recent studies on knowledge base completion, the task of recovering missing relationships based on recorded relations, demonstrate the importance of learning embeddings from multi-step relations. However, due to the size of knowledge bases, learning multi-step relations directly on top of observed instances could be costly. In this paper, we propose Implicit ReasoNets (IRNs), which is designed to perform large-scale inference implicitly through a search controller and shared memory. Unlike previous work, IRNs use training data to learn to perform multi-step inference through the shared memory, which is also jointly updated during training. While the inference procedure is not operating on top of observed instances for IRNs, our proposed model outperforms all previous approaches on the popular FB15k benchmark by more than 5.7%.
1 INTRODUCTION
Knowledge bases such as WordNet (Fellbaum, 1998), Freebase (Bollacker et al., 2008), or Yago (Suchanek et al., 2007) contain many real-world facts expressed as triples, e.g., (Bill Gates, FounderOf, Microsoft). These knowledge bases are useful for many downstream applications such as question answering (Berant et al., 2013; Yih et al., 2015) and information extraction (Mintz et al., 2009). However, despite the formidable size of knowledge bases, many important facts are still missing. For example, West et al. (2014) showed that 21% of the 100K most frequent PERSON entities have no recorded nationality in a recent version of Freebase. We seek to infer unknown relations based on the observed triples. Thus, the knowledge base completion (KBC) task has emerged an important open research problem (Nickel et al., 2011).
Neural-network based methods have been very popular for solving the KBC task. Following Bordes et al. (2013), one of the most popular approaches for KBC is to learn vector-space representations of entities and relations during training, and then apply linear or bi-linear operations to infer the missing relations at test time. However, several recent papers demonstrate limitations of prior approaches relying upon vector-space models alone. By themselves, there is no straightforward way to capture the structured relationships between multiple triples adequately (Guu et al., 2015; Toutanova et al., 2016; Lin et al., 2015a). For example, assume that we want to fill in the missing relation for the triple (Obama, NATIONALITY, ?), a multi-step search procedure might be needed to discover the evidence in the observed triples such as (Obama, BORNIN, Hawaii) and (Hawaii, PARTOF, U.S.A). To address this issue, Guu et al. (2015); Toutanova et al. (2016); Lin et al. (2015a) propose different approaches of injecting structured information by directly operating on the observed triplets. Unfortunately, due to the size of knowledge bases, these newly proposed approaches suffer from some limitations, as most paths are not informative for inferring missing relations, and it is prohibitive to consider all possible paths during the training time with expressive models.
In this paper, we take a different approach from prior work on KBC by addressing the challenges of performing large-scale inference through the design of search controller and shared memory. Our inference procedure centers around the search controller, which only operates on the shared memory instead of directly manipulating the observed triples in knowledge base. IRNs use training data to
∗Equal contribution.
learn to perform multi-step inference through the shared memory. First, input module generates a representation of the query. Then, the search controller repeatedly interacts with the shared memory and checks the termination gate. After each iteration, if the termination condition is met, the model stops the search process and calls the output module to generate a prediction. The shared memory is designed to store key information about the overall structures it learned during training, and hence the search controller only needs to access the shared memory instead of operating on the observed triples.
There are several advantages of using IRNs. First, the cost of inference can be controlled because the search controller only needs to access the shared memory. Second, all the modules, including the search controller and memory, are jointly trained, and hence alleviate the needs to inject structured relationships between instances manually. Finally, we can easily extend IRNs to other tasks that require modeling structured relationships between instances by switching the input and output modules.
The main contributions of our paper are as follows:
• We propose Implicit ReasoNets (IRNs), which use a shared memory guided by a search controller to model large-scale structured relationships implicitly.
• We evaluate IRNs and demonstrate that our proposed model achieves the state-of-the-art results on the popular FB15k benchmark, surpassing prior approaches by more than 5.7%.
• We analyze the behavior of IRNs for shortest path synthesis. We show that IRNs outperform a standard sequence-to-sequence model and execute meaningful multi-step inference.
2 REASONET FOR IMPLICIT INFERENCE
In this section, we describe the general architecture of IRNs in a way that is agnostic to KBC. IRNs are composed of four main components: an input component, an output component, a shared memory, and a search controller, as shown in Figure 1. In this section, we briefly describe each component.
Input/Output Modules: These two modules are task-dependent. The input module takes a query and converts the query into a vector representation q. The output module is a function fo, which converts the hidden state received from the search controller (s) into an output O. We optimize the whole
model using the output prediction O with respect to a ground-truth target using a task-specified loss function.
Shared Memory: The shared memory is denoted as M . It consists of a list of memory vectors, M = {mi}i=1...I , where mi is a fixed dimensional vector. The memory vectors are randomly initialized and automatically updated through back-propagation. The shared memory component is shared across all instances.
Search Controller: The search controller is a recurrent neural network and controls the search process by keeping internal state sequences to track the current search process and history. The search controller uses an attention mechanism to fetch information from relevant memory vectors in M , and decides if the model should output the prediction or continue to generate the next possible output.
• Internal State: The internal state of the search controller is denoted as S, which is a vector representation of the search process. The initial state s1 is usually the vector representation of the input vector q. The internal state at t-th time step is represented by st. The sequence of internal states is modeled by an RNN: st+1 = RNN(st, xt; θs).
• Attention to memory: The attention vector xt at t-th time step is generated based on the current internal state st and the shared memory M : xt = fatt(st,M ; θx). Specifically, the attention score at,i on a memory vector mi given a state st is computed as at,i = softmaxi=1,...,|M |λ cos(W1mi,W2st), where λ is set to 10 in our experiments and the weight matrices W1 and W2 are learned during training. The attention vector xt can be written as xt = fatt(st,M ; θx) = ∑|M | i at,imi.
• Termination Control: The terminate gate produces a stochastic random variable according to the current internal state, tt ∼ p(·|ftc(st; θtc))). tt is a binary random variable. If tt is true, the IRN will finish the search process, and the output module will execute at time step t; otherwise the IRN will generate the next attention vector xt+1 and feed into the state network to update the next internal state st+1. In our experiments, the termination variable is modeled by a logistical regression: ftc(st; θtc) = sigmoid(Wtcst + btc), where the weight matrix Wtc and bias vector btc are learned during training.
Compared IRNs to Memory Networks (MemNN) (Weston et al., 2014; Sukhbaatar et al., 2015; ?) and Neural Turing Machines (NTM) (Graves et al., 2014; 2016), the biggest difference between our model and the existing frameworks is the search controller and the use of the shared memory. We build upon our previous work (Shen et al., 2016) for using a search controller module to dynamically perform a multi-step inference depending on the complexity of the instance. MemNN and NTM explicitly store inputs (such as graph definition, supporting facts) in the memory. In contrast, in IRNs, we do not explicitly store all the observed inputs in the shared memory. Instead, we directly operate on the shared memory, which modeling the structured relationships implicitly. We randomly initialize the memory and update the memory with respect to task-specific objectives. The idea of exploiting shared memory is proposed by Munkhdalai & Yu (2016) independently. Despite of using the same term, the goal and the operations used by IRNs are different from the one used in Munkhdalai & Yu (2016), as IRNs allow the model to perform multi-step for each instance dynamically.
2.1 STOCHASTIC INFERENCE PROCESS
The inference process of an IRN is as follows. First, the model converts a task-dependent input to a vector representation through the input module. Then, the model uses the input representation to initialize the search controller. In every time step, the search controller determines whether the process is finished by sampling from the distribution according to the terminate gate. If the outcome is termination, the output module will generate a task-dependent prediction given the search controller states. If the outcome is continuation, the search controller will move on to the next time step, and create an attention vector based on the current search controller state and the shared memory. Intuitively, we design whole process by mimicking a search procedure that iteratively finds its target through a structure and output its prediction when a satisfying answer is found. The detailed inference process is described in Algorithm 1.
The inference process of an IRN is considered as a Partially Observable Markov Decision Process (POMDP) (Kaelbling et al., 1998) in the reinforcement learning (RL) literature. The IRN produces
Algorithm 1: Stochastic Inference Process in an IRN Input :Randomly initialized shared memory M ; Input vector q; Maximum step Tmax Output :Output vector o
1 Define s1 = q; t = 1; 2 Sample tt from the distribution p(·|ftc(st; θtc)); 3 if tt is false, go to Step 4; otherwise Step 7; 4 Generate an attention vector xt = fatt(st,M ; θx); 5 Update the internal state st+1 = RNN(st, xt; θs); 6 Set t = t+ 1; if t < Tmax go to Step 2; otherwise Step 7; 7 Generate output ot = fo(st; θo); 8 Return o = ot;
the output vector oT at the T -th step, which implies termination gate variables t1:T = (t1 = 0, t2 = 0, ..., tT−1 = 0, tT = 1), and then takes prediction action pT according to the probability distribution given oT . Therefore, the IRN learns a stochastic policy π((t1:T , pT )|q; θ) with parameters θ to get a distribution over termination actions, and over prediction actions. The termination step T varies from instance to instance. The parameters of the IRN θ are given by the parameters of the embedding matrices W for the input/output module, the shared memory M , the attention network θx, the search controller RNN network θs, the output generation network θo, and the termination gate network θtc. The parameters θ = {W,M, θx, θs, θo, θtc} are trained to maximize the total expected reward that the IRN when interacting with the environment. The expected reward for an instance is defined as:
J(θ) = Eπ(t1:T ,pT ;θ) [ T∑ t=1 rt ]
The reward can only be received at the final termination step when a prediction action pT is performed. The rewards on intermediate steps are zeros, {rt = 0}t=1...T−1. We employ the approach from our previous work (Shen et al., 2016), REINFORCE (Williams, 1992) based Contrastive Reward method, to maximize the expected reward. The gradient of J can be written as:
∇θJ(θ) = ∑
(t1:T ,pT )∈A† π(t1:T , pT ; θ)
[ ∇θlogπ(t1:T , pT ; θ)(
rT bi − 1)
]
where A† is all the possible episodes, the baseline bi = ∑
(t1:T ,pT )∈A† π(t1:T , pT ; θ)rT is the expected reward on the |A†| episodes for the i-th training instance.
3 APPLYING IRNS TO KNOWLEDGE BASE COMPLETION
The goal of KBC tasks (Bordes et al., 2013) is to predict a head or a tail entity given the relation type and the other entity, i.e. predicting h given (?, r, t) or predicting t given (h, r, ?), where ? denotes the missing entity. For a KBC task, the input to our model is a subject entity (a head or tail entity) and a relation. The task-dependent input module first extracts the embedding vectors for the entity and relation from an embedding matrix. We then represent the query vector q for an IRN as the concatenation of the two vectors. We randomly initialize the shared memory component. At each step, a training triplet is processed through the model by Algorithm 1, where no explicit path information is given. The IRN updates the shared memory implicitly with respect to the objective function. For the task dependent output module, we use a nonlinear projection to project the search controller state into an output vector o: fo(st; θo) = tanh(Wost + bo), where the Wo and bo are the weight matrix and bias vector, respectively. We define the ground truth target (object) entity embedding as y, and use the L1 distance measure between the output o and target entity y, namely d(o, y) = |o− y|1. We sample a set of incorrect entity embeddings N = {y−i } |N | i=1 as negative examples. The probability of
selecting a prediction ŷ ∈ D can be approximated as
p(ŷ|o) = exp(−γd(o, ŷ))∑ yk∈D exp(−γd(o, yk))
where D = N ∪ {y}. We set |N | and γ to 20 and 5, respectively, for the experiments on FB15k and WN18 datasets. The IRN performs a prediction action pT on selecting ŷ with probability p(ŷ|o) . We define the reward of the prediction action as one if the ground truth entity is selected, and zero otherwise.
4 EXPERIMENTAL RESULTS
In this section, we evaluate the performance of our model on the benchmark FB15k and WN18 datasets for KBC tasks (Bordes et al., 2013). These datasets contain multi-relations between head and tail entities. Given a head entity and a relation, the model produces a ranked list of the entities according to the score of the entity being the tail entity of this triple. To evaluate the ranking, we report mean rank (MR), the mean of rank of the correct entity across the test examples, and hits@10, the proportion of correct entities ranked in the top-10 predictions. Lower MR or higher hits@10 indicates a better prediction performance. We follow the evaluation protocol in Bordes et al. (2013) to report filtered results, where negative examples N are removed from the dataset. In this case, we can avoid some negative examples being valid and ranked above the target triplet.
We use the same hyper-parameters of our model for both FB15k and WN18 datasets. Entity embeddings (which are not shared between input and output modules) and relation embedding are both 100-dimensions. We use the input module and output module to encode subject and object entities, respectively. There are 64 memory vectors with 200 dimensions each, initialized by random vectors with unit L2-norm. We use single-layer GRU with 200 cells as the search controller. We set the maximum inference step of the IRN to 5. We randomly initialize all model parameters, and use SGD as the training algorithm with mini-batch size of 64. We set the learning rate to a constant number, 0.01. To prevent the model from learning a trivial solution by increasing entity embeddings norms, we follow Bordes et al. (2013) to enforce the L2-norm of the entity embeddings as 1. We use hits@10 as the validation metric for the IRN. Following the work (Lin et al., 2015a), we add reverse relations into the training triplet set to increase the training data.
Following Nguyen et al. (2016), we divide the results of previous work into two groups. The first group contains the models that directly optimize a scoring function for the triples in a knowledge base without using extra information. The second group of models make uses of additional information from multi-step relations. For example, RTransE (García-Durán et al., 2015) and PTransE (Lin et al., 2015a) models are extensions of the TransE (Bordes et al., 2013) model by explicitly exploring multi-step relations in the knowledge base to regularize the trained embeddings. The NLFeat model (Toutanova et al., 2015) is a log-linear model that makes use of simple node and link features.
Table 1 presents the experimental results. According to the table, our model significantly outperforms previous baselines, regardless of whether previous approaches use additional information or not. Specifically, on FB15k, the MR of our model surpasses all previous results by 12, and our hit@10 outperforms others by 5.7%. On WN18, the IRN obtains the highest hit@10 while maintaining similar MR results compared to previous work.1
To better understand the behavior of IRNs, we report the results of IRNs with different memory sizes and different Tmax on FB15K in Table 2. We find the performance of IRNs increases significantly if the number of inference step increases. Note that an IRN with Tmax = 1 is the case that an IRN without the shared memory. Interestingly, given Tmax = 5, IRNs are not sensitive to memory sizes. In particular, larger memory always improves the MR score, but the best hit@10 is obtained by |M | = 64 memory vectors. A possible reason is that the best memory size is determined by the complexity of the tasks.
We analyze hits@10 results on FB15k with respect to the relation categories. Following the evaluation in Bordes et al. (2013), we evaluate the performance in four types of relation: 1-1 if a head entity
1Nguyen et al. (2016) reported two results on WN18, where the first one is obtained by choosing to optimize hits@10 on the validation set, and second one is obtained by choosing to optimize MR on the validation set. We list both of them in Table 1.
can appear with at most one tail entity, 1-Many if a head entity can appear with many tail entities, Many-1 if multiple heads can appear with the same tail entity, and Many-Many if multiple head entities can appear with multiple tail entities. The detailed results are shown in Table 3. The IRN significantly improves the hits@10 results in the Many-1 category on predicting the head entity (18.8%), the 1-Many category on predicting the tail entity (16.5%), and the Many-Many category (over 8% in average).
To analyze the behavior of IRNs, we pick some examples for the tail entity prediction in Table 4. Interestingly, we observed that the model can gradually increase the ranking score of the correct tail entity during the inference process.
5 ANALYSIS: APPLYING IRNS TO A SHORTEST PATH SYNTHESIS TASK
We construct a synthetic task, shortest path synthesis, to evaluate the inference capability over a shared memory. The motivations of applying our model to this task are as follows. First, we want to evaluate IRNs on another task requiring multi-step inference. Second, we select the sequence generation task so that we are able to analyze the inference capability of IRNs in details.
In the shortest path synthesis task, as illustrated in Figure 2, a training instance consists of a start node and an end node (e.g., 215 493) of an underlying weighted directed graph that is unknown to models. The output of each instance is the shortest path between the given start and end nodes of the underlying graph (e.g., 215→ 101→ 493). Specifically, models can only observe the start-end node
pairs as input and their shortest path as output. The whole graph is unknown to the models and the edge weights are not revealed in the training data. At test time, a path sequence is considered correct if it connects the start node and the end node of the underlying graph, and the cost of the predicted path is the same as the optimal path.
Note that the task is very difficult and cannot be solved by dynamic programming algorithms since the weights on the edges are not revealed to the algorithms or the models. To recover some of the shortest paths at the test time, the model needs to infer the correct path from the observed instances. For example, assume that we observe two instances in the training data, “A D: A→ B → G→ D” and “B E: B → C → E”. In order to answer the shortest path between A and E, the model needs to infer that “A→ B → C → E” is a possible path between A and E. If there are multiple possible paths, the model has to decide which one is the shortest one using statistical information.
In the experiments, we construct a graph with 500 nodes and we randomly assign two nodes to form an edge. We split 20,000 instances for training, 10,000 instances for validation, and 10,000 instances for testing. We create the training and testing instances carefully so that the model needs to perform inference to recover the correct path. We describe the details of the graph and data construction parts in the appendix section. A sub-graph of the data is shown in Figure 2.
For the settings of the IRN, we switch the output module to a GRU decoder for a sequence generation task. We assign reward rT = 1 if all the prediction symbols are correct and 0 otherwise. We use a 64-dimensional embedding vector for input symbols, a GRU controller with 128 cells, and a GRU decoder with 128 cells. We set the maximum inference step Tmax to 5.
We compare the IRN with two baseline approaches: dynamic programming without edge-weight information and a standard sequence-to-sequence model (Sutskever et al., 2014) using a similar parameter size to our model. Without knowing the edge weights, dynamic programming only recovers 589 correct paths at test time. The sequence-to-sequence model recovers 904 correct paths. The IRN outperforms both baselines, recovering 1,319 paths. Furthermore, 76.9% of the predicted paths from IRN are valid paths, where a path is valid if the path connects the start and end node nodes of the underlying graph. In contrast, only 69.1% of the predicted paths from the sequence-to-sequence model are valid.
To further understand the inference process of the IRN, Figure 2 shows the inference process of a test instance. Interestingly, to make the correct prediction on this instance, the model has to perform a fairly complicated inference.2 We observe that the model cannot find a connected path in the first three steps. Finally, the model finds a valid path at the forth step and predict the correct shortest path sequence at the fifth step.
6 RELATED WORK
Link Prediction and Knowledge Base Completion Given that r is a relation, h is the head entity, and t is the tail entity, most of the embedding models for link prediction focus on finding the scoring function fr(h, t) that represents the implausibility of a triple. (Bordes et al., 2011; 2014; 2013; Wang et al., 2014; Ji et al., 2015; Nguyen et al., 2016). In many studies, the scoring function fr(h, t) is linear or bi-linear. For example, in TransE (Bordes et al., 2013), the function is implemented as fr(h, t) = ‖h+ r− t‖, where h, r and t are the corresponding vector representations. Recently, different studies (Guu et al., 2015; Lin et al., 2015a; Toutanova et al., 2016) demonstrate the importance for models to also learn from multi-step relations. Learning from multi-step relations injects the structured relationships between triples into the model. However, this also poses a technical challenge of considering exponential numbers of multi-step relationships. Prior approaches address this issue by designing path-mining algorithms (Lin et al., 2015a) or considering all possible paths using a dynamic programming algorithm with the restriction of using linear or bi-linear models only (Toutanova et al., 2016). Toutanova & Chen (2015) shows the effectiveness of using simple node and link features that encode structured information on FB15k and WN18. In our work, the IRN outperforms prior results and shows that similar information can be captured by the model without explicitly designing features.
2 In the example, to find the right path, the model needs to search over observed instances “215 448: 215→ 101→ 448” and “76 493: 76→ 308→ 101→ 493”, and to figure out the distance of “140→ 493” is longer than “101→ 493” (there are four shortest paths between 101→ 493 and three shortest paths between 140→ 493 in the training set).
Studies such as (Riedel et al., 2013) show that incorporating textual information can further improve the knowledge base completion tasks. It would be interesting to incorporate the information outside the knowledge bases in our model in the future.
Neural Frameworks Sequence-to-sequence models (Sutskever et al., 2014; Cho et al., 2014) have shown to be successful in many applications such as machine translation and conversation modeling (Sordoni et al., 2015). While sequence-to-sequence models are powerful, recent work has shown that the necessity of incorporating an external memory to perform inference in simple algorithmic tasks (Graves et al., 2014; 2016).
7 CONCLUSION
In this paper, we propose Implicit ReasoNets (IRNs), which perform inference over a shared memory that models large-scale structured relationships implicitly. The inference process is guided by a search controller to access the memory that is shared across instances. We demonstrate and analyze the multi-step inference capability of IRNs in the knowledge base completion tasks and a shortest path synthesis task. Our model, without using any explicit knowledge base information in the inference procedure, outperforms all prior approaches on the popular FB15k benchmark by more than 5.7%.
For future work, we aim to further extend IRNs in two ways. First, inspired from Ribeiro et al. (2016), we would like to develop techniques to exploit ways to generate human understandable reasoning interpretation from the shared memory. Second, we plan to apply IRNs to infer the relationships in unstructured data such as natural language. For example, given a natural language query such as “are rabbits animals?”, the model can infer a natural language answer implicitly in the shared memory without performing inference directly on top of huge amount of observed sentences such as “all mammals are animals” and “rabbits are animals”. We believe the ability to perform inference implicitly is crucial for modeling large-scale structured relationships.
ACKNOWLEDGMENTS
We thank Scott Wen-Tau Yih, Kristina Toutanova, Jian Tang and Zachary Lipton for their thoughtful feedback and discussions.
A DETAILS OF THE GRAPH CONSTRUCTION FOR THE SHORTEST PATH SYNTHESIS TASK
We construct the underlying graph as follows: on a three-dimensional unit-sphere, we randomly generate a set of nodes. For each node, we connect its K-nearest neighbors and use the euclidean distance between two nodes to construct a graph. We randomly sample two nodes and compute its shortest path if it is connected between these two nodes. Given the fact that all the sub-paths within a shortest path are shortest paths, we incrementally create the dataset and remove the instances which are a sub-path of previously selected paths or are super-set of previous selected paths. In this case, all the shortest paths can not be answered through directly copying from another instance. In addition, all the weights in the graph are hidden and not shown in the training data, which increases the difficulty of the tasks. We set k = 50 as a default value.
[Summary]
This paper proposes a new way for knowledge base completion which highlights: 1) adopting an implicit shared memory, which makes no assumption about its structure and is completely learned during training; 2) modeling a multi-step search process that can decide when to terminate.
The experimental results on WN18 and FB15k seem pretty good. The authors also perform an analysis on a shortest path synthetic task, and demonstrate that this model is better than standard seq2seq.
The paper is well-written and it is easy to follow.
[Major comments]
I actually do like the idea and am also impressed that this model can work well.
The main concern is that this paper presents too little analysis about how it works and whether it is sensitive to the hyper-parameters, besides that only reporting a final model on WN18 and FB15k.
One key hyper-parameter I believe is the size of shared memory (using 64 for the experiments). I don’t think that this number should be fixed for all tasks, at least it should depend on the KB scale. Could you verify this in your experiments? Would it be even possible to make a memory structure with dynamic size?
The RL setting (stochastic search process) is also one highlight of the paper, but could you demonstrate that how much it does really help? I think it is necessary to compare to the following: remove the termination gate and fix the number of inference steps and see how well the model does? Also show how the performance varies on # of steps?
I appreciate your attempts on the shortest path synthetic task. However, I think it would be much better if you can demonstrate that under a real KB setting. You can still perform the shortest path analysis, but using KB  (e.g., Freebase) entities and relations.
[Minor comments]
I am afraid that the output gate illustrated in Figure 1 is a bit confusing. There should be only one output, depending on when the search process is terminated.