ProjE: Embedding Projection for Knowledge Graph Completion

In plain words

A knowledge graph stores facts as triples, such as (Springfield, located in, Illinois). Many facts are missing, so the task is to guess them. ProjE handles a half-known fact, say "Illinois, CityOf, ?", by asking which entities fit.

It blends the two known pieces into one target point in a vector space, then scores every entity by how well it lines up with that point. The entities are ranked so the true answer comes first. The model is small, and it learns only from the graph itself, without pretrained embeddings or multi-step paths. The authors report that it ranks answers better than earlier methods on a standard Freebase benchmark, and that it can also check whether single statements are true.

TL;DR

ProjE is a two-layer network for knowledge-graph completion that ranks all candidate entities for a partial triple, not one score at a time. Its combination step is cheap: diagonal operators on the entity and relation embeddings, then a tanh. A shared projection onto the entity table gives the scores. The authors report better filtered Hits@10 than the listed baselines on FB15K with fewer parameters. The code is MIT-licensed. Its entity-prediction path runs, but relation prediction and the fact-checking experiments have no code here.

Problem

Models that score each triple one at a time (TransE and its successors) use one margin loss per triple §1 eq. 1 eq. 2. Richer models add per-relation matrices, neural layers or multi-hop paths, which raises the parameter count or needs pretrained inputs §1 tab. 1. The paper argues that the ranking objective matters more than extra features: scoring all candidates together, with a simple architecture, should beat those models §1. Entity prediction is the main task: given (h, r, ?) or (?, r, t), rank all entities [def. 1] §1. The head and tail tasks are symmetric, so the same model can serve either §1.

Parameter sizes for the compared models, as the paper lists them tab. 1:

Here n_e, n_r and n_w are the numbers of entities, relations and words, k is the embedding size, and z is the hidden-layer size tab. 1. The "5k" count for ProjE is the paper's own figure; see Paper vs code for what the code has.

Method

Scope. ProjE is self-contained: it trains only on direct, length-1 edges and needs no pretrained embeddings or pre-selected paths §1 §4.

Combination. Entity and relation embeddings are merged by two diagonal matrices and a bias, so no full matrix is learned §2.1 eq. 4: e ⊕ r = D_e e + D_r r + b_c.

Projection. The merged vector is passed through f (tanh), then projected onto candidate entity rows with a bias b_p, and g scores each candidate §2.1 eq. 5: h(e, r) = g(W^c f(e ⊕ r) + b_p). The candidate matrix W^c is made of rows of the entity table, so it adds no new variables §2.1. Figure 1 shows the architecture for the example (?, CityOf, Illinois) as a two-layer network with a combination layer and a projection layer fig. 1.

Three variants. They differ only in g and in the loss §2.2:

• ProjE_pointwise: g is sigmoid, and the loss is a binary cross-entropy over the positive candidates and sampled negatives §2.2 eq. 6 eq. 7.
• ProjE_listwise: g is softmax, and each positive gets target probability 1 / (number of positives for the query) §2.2 eq. 8 eq. 9.
• ProjE_wlistwise: as listwise, with the 1/(number of positives) factor removed. The paper calls this an instance-level weight Σ_i 1(y_i = 1) §3.2. Its training procedure is Algorithm 1 app. A.1 alg. 1. It is the variant the paper reports as best app. A.2 tab. 2.

Candidate sampling. Each non-positive candidate is kept with probability p_y, and the positives are always kept §2.3. The paper states that p_y = 25% performs well and that 5% falls below the top baselines app. A.3 fig. 2.

Training (Algorithm 1). For every epoch alg. 1:

1. For each training triple, pick one of its two entities at random. The other entity's side is the one predicted, and the query is (known entity, r). The positives are all known answers for that query in the training set. Negatives are sampled with probability p_y.
2. Split the instances into mini-batches. For each instance, compute the softmax over its candidates, sum -log p over the positives, and add an L1 penalty on the entity and relation embeddings and the combination operators. Update all parameters with the sum, scaled by the L1 weight α.

Dropout with probability p_d is applied in training only alg. 1 §3.1.

# ProjE_wlistwise: one epoch of instance construction  [alg. 1] [§2.3]
# Each triple gives one instance; negatives are drawn once per epoch
def build_epoch(triples, entities, p_y):
    instances = []
    for head, relation, tail in triples:
        # Pick the known side at random; the other side is predicted  [alg. 1]
        if random_choice([head, tail]) == head:
            query = (head, relation)
            positives = known_tails(head, relation)    # training set only
            direction = "tail"
        else:
            query = (tail, relation)
            positives = known_heads(relation, tail)    # training set only
            direction = "head"
        # Each non-positive entity is kept with probability p_y  [§2.3]
        negatives = [e for e in entities
                     if e not in positives and random() < p_y]
        candidates = positives + negatives
        # 1 for positives, 0 for sampled negatives  [alg. 1]
        positive_mask = [1.0] * len(positives) + [0.0] * len(negatives)
        instances.append((query, direction, candidates, positive_mask))
    return instances

# ProjE_wlistwise: loss for one instance  [eq. 4] [eq. 5] [alg. 1]
# Input: entity table E (N, k); relation table R (M, k)
# Output: scalar loss before the L1 term
def instance_loss(query, direction, candidates, positive_mask, E, R, ops):
    entity, relation = query
    # Diagonal operators and bias, chosen by prediction direction  [eq. 4]
    D_e, D_r, b_c = ops[direction]              # each (k,)
    combined = D_e * E[entity] + D_r * R[relation] + b_c
    # Dropout sits inside the tanh, as in the algorithm  [alg. 1]
    hidden = tanh(dropout(combined, p_drop))    # (k,)
    # Score only the candidate rows, add projection bias  [alg. 1] [eq. 9]
    logits = E[candidates] @ hidden + b_p[candidates]   # (C,)
    log_probs = log_softmax(logits)             # (C,)
    # Negative log-likelihood over positives, no 1/|positives| factor  [alg. 1]
    return -sum(positive_mask * log_probs)

# ProjE_wlistwise: mini-batch update  [alg. 1]
# The L1 term covers embeddings and all four operators, not biases  [alg. 1]
def update(batch, E, R, ops, alpha):
    loss = 0.0
    for instance in batch:
        loss += instance_loss(*instance, E, R, ops)
    # Penalty is added to the summed loss, not averaged  [alg. 1]
    penalty = l1_norm(E) + l1_norm(R) + l1_norm(ops)
    total = loss + alpha * penalty
    # Adam step on all trainable parameters  [§3.1] [alg. 1]
    adam_step(total)

Test time. Score all entities for the partial triple and sort them in descending order. Filtered metrics remove the other known true answers from the ranking §3.2. Hits@k is the share of queries with the true answer in the top k.

Fact checking (not in the code). This task uses the predicate paths between two entities as input, through an extra input layer that turns paths into entity-like vectors. It uses the ProjE_wlistwise loss §3.3. Table 4 reports AUC on DBpedia and SemMedDB tab. 4.

Hyperparameters

Results

Author-reported. Setup: FB15K, a 15,000-entity subset of Freebase, with the standard TransE and PTransE protocols §3 §3.2. Mean rank and Hits@k are reported as raw and filtered §3.2. Table 2 is entity prediction tab. 2.

Relation prediction on FB15K tab. 3. Mean rank (MR) and HITS@1 (%), raw and filtered. This table has no code in the repository (see Paper vs code).

Fact checking, AUC on DBpedia and SemMedDB tab. 4. Only the ProjE row and the best baseline per column are shown here; the full table is in the paper.

The paper says ProjE beats all but one question type in fact checking §3.3. The exception is SemMedDB Disease, where PredPath is higher tab. 4.

Paper vs code

Official code: bxshi/ProjE, the first author's repository, MIT-licensed (copyright Baoxu "Dash" Shi). The paper names nddsg/ProjE §3. That repository is a GitHub fork of bxshi/ProjE, and its last push is dated 2018-01-22. I did not compare its tree with the one I read. Everything below cites the cloned bxshi/ProjE at 8d8d6b4. Code citations use ProjE_softmax.py unless a file is named.

• Projection bias is missing. The paper's projection has a bias b_p eq. 5 eq. 7 eq. 9 alg. 1. The code's projection has none: the scores are the tanh output times the entity table ProjE_softmax.py ProjE_softmax.py. The same holds in ProjE_sigmoid.py.
• Dropout placement differs. Algorithm 1 puts dropout inside the tanh, before it alg. 1. The code applies dropout to the tanh output ProjE_softmax.py ProjE_softmax.py. The text only says dropout is on top of the combination operator §2.1.
• Operators per direction. Equation 4 has one D_e and one D_r eq. 4. Algorithm 1 has four diagonal operators, one set per prediction direction alg. 1. The code also keeps two sets, each with a weight vector of 2k and a bias of k ProjE_softmax.py. Counting those, the code has 6k non-embedding parameters per model. Table 1 and §2.1 say 5k (or 5k + 1) tab. 1 §2.1.
• Biases are initialized to zero. §3.1 says all parameters are drawn from the uniform range §3.1. The code uses zeros for the biases [code: ProjE_softmax.py#L210-L213]. Algorithm 1 does not state the bias initialization alg. 1.
• The two weight vectors share a seed. The head and tail combination vectors use the same seed, 445, so they start with identical values ProjE_softmax.py. The paper does not state this.
• Relation prediction is not in the code. The test function scores only entities ProjE_softmax.py. The p_y = 0.75 and k = 100 of the relation task §3.1 cannot be reproduced from this repository.
• Fact checking is not in the code. The predicate-path input layer §3.3 is absent from all three scripts.
• Candidate sampling does not cut computation. The paper motivates sampling by the cost of projecting onto many candidates §2.3. The code projects onto all entity rows and masks the unsampled ones with weight 0 ProjE_softmax.py ProjE_softmax.py. The per-instance cost stays proportional to the number of entities.
• Loss variants. The wlistwise loss (no 1/N factor) matches Algorithm 1 and is in ProjE_softmax.py ProjE_softmax.py. The listwise loss of eq. 8, which divides by the positive count, is in ProjE_softmax_noweight.py ProjE_softmax_noweight.py. The file names do not say this; the mapping is my reading of the formulas. The pointwise sigmoid loss is in ProjE_sigmoid.py ProjE_sigmoid.py.
• Mean rank is 0-based. Ranks count from zero, so a top-1 answer has rank 0 ProjE_softmax.py. Hits@10 counts ranks below 10, which is correct for 0-based ranks ProjE_softmax.py. The paper does not say which convention it uses, so compare mean ranks with care.
• Entity vectors are not normalized. The code assigns the entity table to a variable named as normalized ProjE_softmax.py, but nothing is normalized. The paper does not mention normalization.
• Negative sampling in the code. Each non-positive candidate gets a weight of 0 or -1, sampled with probability p_y and independent per candidate ProjE_softmax.py. This matches §2.3.
• Filtered evaluation. The filter uses known triples from the train, valid and test splits ProjE_softmax.py ProjE_softmax.py.
• Run command. The readme sets dimension, batch, iterations and generator count, but not the learning rate, keep probability or negative rate. Those take the defaults above readme.md.
• Mini-batch default. ProjE_softmax_noweight.py defaults to 100, not 200 ProjE_softmax_noweight.py.
• No reproduction output. The repository has no logs, checkpoints or seeds for the data split. Table 2 cannot be checked from the code alone.

Pitfalls

• Missing projection bias. The code scores with no b_p. A reimplementation that follows the equations will differ from the released code and from Algorithm 1. Whether the paper's reported numbers used b_p is not stated.
• Dropout location. Algorithm 1 applies dropout before the tanh; the code applies it after. Both use p = 0.5. Pick one and state which.
• Keep versus drop probability. The code's flag is a keep probability ProjE_softmax.py. The paper calls its value a dropout probability alg. 1 §3.1. At 0.5 the two agree, but a different value would not.
• Shape of the projection bias. Equation 5 gives b_p the length s, the number of candidates §2.1, so it cannot be one fixed vector over all entities without a per-entity index. The paper does not say how it is indexed. The code has no such bias at all (see Paper vs code).
• Loss scale. The loss is summed over the mini-batch and over positives, not averaged, and the L1 term is added to that sum with α = 1e-5 ProjE_softmax.py. The effective strength of the penalty therefore depends on batch size and on the number of positives per query. The paper does not say whether the sum or the mean was intended.
• Cost per instance. Each instance computes a score for every entity, so memory grows with batch size times entity count. The sampling step saves negatives in the loss, not in the matrix product (see Paper vs code).
• Early stopping and selection. The code prints validation and test metrics after each epoch and saves checkpoints; it does not select a model by validation ProjE_softmax.py. The paper does not state the stopping rule or the epoch for the reported numbers.
• Sampling rate. Good results at p_y = 25% are reported only for the wlistwise variant app. A.3. At 5%, the appendix reports results below the top three baselines app. A.3 fig. 2.
• Pretrained initialization. §3.1 says ProjE can start from pretrained embeddings §3.1. The code has no such option.
• Unseeded training data. Triple order and negatives use NumPy's global random state ProjE_softmax.py. Only the initialization has fixed seeds ProjE_softmax.py. Repeated runs will differ.
• Appendix not stated in the main text. The algorithm and the sampling sweep are in the appendix app. A.1 app. A.3. Use Algorithm 1 for the wlistwise procedure, not the main text.
• Dataset. The readme points to the FB15k and WN18 files from the KB2E release, not to the paper's data split readme.md. The paper gives only "15,000-entity subset of Freebase" §3.

Lineage

• Builds on: TransE, which models a relation as a translation between two entity embeddings §1 eq. 1. The ranking loss is the margin-based pairwise loss used by TransE-family models eq. 2. ProjE replaces it with a listwise objective over all candidates.
• Compared against: TransE, TransH, TransR, TransA, PTransE, DKRL, HolE, SME, RESCAL, SE, LFM, NTN, CVSM, RTransE tab. 1 tab. 2 tab. 3.
• Relationship to later work: the readme points to the authors' ConMask work on open-world completion readme.md. This wiki does not hold a card for that work, so no claim about it is made here.