mixle.stats.graphs.knowledge_graph module

Knowledge-graph embedding distributions for integer triple observations.

Data type: a triple (h, r, t) of integer indices – head entity, relation, tail entity. The model embeds each entity and relation in dim dimensions and scores a triple by the DistMult bilinear form

score(h, r, t) = sum_k E[h, k] * R[r, k] * E[t, k] = (E[h] * R[r]) . E[t],

and defines the conditional tail distribution by a softmax over all entities,

p(t | h, r) = softmax_t score(h, r, t), log p(h, r, t) = score(h, r, t) - logsumexp_a score(h, r, a).

This is the standard tail-prediction likelihood; maximizing it over observed triples is the model’s MLE. It has no closed form, so – exactly like the Plackett-Luce minorization-maximization estimator in this package – each fit / optimize iteration performs one full-batch gradient-ascent step on the embeddings, evaluated at the previous estimate (a random seeded init seeds the first pass). The threaded estimate carries the embeddings between passes, so no parameter state lives outside the framework.

class KnowledgeGraphDistribution(entity_embeddings, relation_embeddings, name=None, keys=None)

Bases: SequenceEncodableProbabilityDistribution

DistMult knowledge-graph embedding distribution over triples (h, r, t).

entity_embeddings is (num_entities, dim) and relation_embeddings is (num_relations, dim). log_density((h, r, t)) is the conditional tail log-probability log p(t | h, r) under the entity softmax.

Parameters:

• entity_embeddings (Any)

• relation_embeddings (Any)

• name (str | None)

• keys (str | None)

score(h, r, t)

DistMult score of a single triple (higher is more plausible).

Parameters:

• h (int)

• r (int)

• t (int)

Return type:

float

tail_log_posterior(h, r)

Length-num_entities vector of log p(t | h, r) over all tail candidates.

Parameters:

• h (int)

• r (int)

Return type:

ndarray

head_log_posterior(r, t)

Length-num_entities vector of log p(h | r, t) over all head candidates.

Parameters:

• r (int)

• t (int)

Return type:

ndarray

relation_log_posterior(h, t)

Length-num_relations vector of log p(r | h, t) over all relation candidates.

Parameters:

• h (int)

• t (int)

Return type:

ndarray

complete(h=None, r=None, t=None)

Log-posterior over candidates for the single missing slot of a query.

Exactly one of h, r, t must be None; the returned vector is over entities (for a missing head or tail) or relations (for a missing relation).

Parameters:

• h (int | None)

• r (int | None)

• t (int | None)

Return type:

ndarray

rank(h=None, r=None, t=None, exclude=(), top_n=None)

Rank candidates for the missing slot by log-probability, dropping exclude candidates.

Returns [(candidate, log_prob), ...] highest first (the most plausible completions).

Parameters:

• h (int | None)

• r (int | None)

• t (int | None)

• exclude (Any)

• top_n (int | None)

Return type:

list[tuple[int, float]]

recommend(known, top_n=10)

Recommend the most plausible missing tail facts for the (h, r) contexts in known.

known is a sequence of observed (h, r, t) triples; for each distinct (h, r) the already-present tails are excluded, the remaining tails are ranked by log p(t | h, r), and the global top top_n new facts are returned as [(h, r, t, log_prob), ...].

Parameters:

• known (Any)

• top_n (int)

Return type:

list[tuple[int, int, int, float]]

recommend_subgraph(node, known, top_n=5)

Recommend plausible new edges incident to node (both (node, r, ?) and (?, r, node)).

Excludes edges already in known and returns the top top_n by log-probability as [(h, r, t, log_prob), ...], the suggested missing subgraph around the node.

Parameters:

• node (int)

• known (Any)

• top_n (int)

Return type:

list[tuple[int, int, int, float]]

pattern(pattern, candidates=None, known=None, beam=64)

A subgraph-pattern query over this model for flexible enumeration of missing parts.

pattern is a list of triples whose slots are either fixed integer ids or named variables (strings starting with '?'), variables shared across edges (e.g. [(alice, friend, '?x'), ('?x', lives_in, '?c')]). The returned KnowledgeGraphPattern enumerates the variable bindings (completed subgraphs) in descending joint plausibility, restricts variables to candidates if given, drops groundings that add nothing new when known is given, and plugs into ConformalStructure for a calibrated set of completed subgraphs.

Parameters:

• pattern (Any)

• candidates (Any)

• known (Any)

• beam (int)

Return type:

KnowledgeGraphPattern

log_density(x)

Return log p(t | h, r) for one integer triple.

Parameters:

x (Sequence[int])

Return type:

float

seq_log_density(x)

Return vectorized tail log-probabilities for encoded triples.

Parameters:

x (ndarray)

Return type:

ndarray

sampler(seed=None)

Return a sampler for observed triples.

Parameters:

seed (int | None)

Return type:

KnowledgeGraphSampler

estimator(pseudo_count=None)

Return a DistMult embedding estimator for this entity/relation shape.

Parameters:

pseudo_count (float | None)

Return type:

KnowledgeGraphEstimator

dist_to_encoder()

Return the triple encoder used by vectorized methods.

Return type:

KnowledgeGraphDataEncoder

class KnowledgeGraphSampler(dist, seed=None)

Bases: DistributionSampler

Draw triples: head and relation uniformly, tail from the conditional softmax p(t | h, r).

Parameters:

• dist (KnowledgeGraphDistribution)

• seed (int | None)

sample(size=None)

Draw one triple or size iid triples.

Parameters:

size (int | None)

Return type:

Any

class KnowledgeGraphAccumulator(keys=None)

Bases: SequenceEncodableStatisticAccumulator

Collect the observed triples (and weights) for the estimator to train on.

A DistMult embedding model has no finite sufficient statistic, so – like other non-exponential-family models in this package – the accumulator retains the data: it concatenates the (h, r, t) triples seen across the (possibly distributed) partitions. The estimator then runs the gradient training in KnowledgeGraphEstimator.estimate().

Parameters:

keys (str | None)

update(x, weight, estimate)

Store one weighted triple for embedding training.

Parameters:

• x (Sequence[int])

• weight (float)

• estimate (KnowledgeGraphDistribution | None)

Return type:

None

initialize(x, weight, rng)

Store one weighted triple during initialization.

Parameters:

• x (Sequence[int])

• weight (float)

• rng (RandomState | None)

Return type:

None

seq_initialize(x, weights, rng)

Store encoded triples during initialization.

Parameters:

• x (ndarray)

• weights (ndarray)

• rng (RandomState | None)

Return type:

None

seq_update(x, weights, estimate)

Store encoded triples and weights for embedding training.

Parameters:

• x (ndarray)

• weights (ndarray)

• estimate (KnowledgeGraphDistribution | None)

Return type:

None

combine(suff_stat)

Merge stored triples and weights from another accumulator value.

Parameters:

suff_stat (tuple)

Return type:

KnowledgeGraphAccumulator

value()

Return total weight, stacked triples, and stacked weights.

Return type:

tuple

from_value(x)

Restore stored triples and weights from value output.

Parameters:

x (tuple)

Return type:

KnowledgeGraphAccumulator

key_merge(stats_dict)

Merge this accumulator into stats_dict under its configured key.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

Replace this accumulator’s state from keyed statistics when present.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Return the encoder compatible with stored triples.

Return type:

KnowledgeGraphDataEncoder

class KnowledgeGraphAccumulatorFactory(keys=None)

Bases: StatisticAccumulatorFactory

Factory for KnowledgeGraphAccumulator.

Parameters:

keys (str | None)

make()

Create an empty knowledge-graph accumulator.

Return type:

KnowledgeGraphAccumulator

class KnowledgeGraphEstimator(num_entities, num_relations, dim=16, lr=0.5, epochs=100, batch_size=256, weight_decay=1.0e-4, init_scale=0.3, max_norm=1.0, directions=('tail', 'head', 'relation'), negatives=None, seed=1, pseudo_count=None, name=None, keys=None)

Bases: ParameterEstimator

Train DistMult knowledge-graph embeddings by maximizing the tail-softmax log-likelihood.

estimate runs vectorized mini-batch gradient ascent (epochs passes, batch size batch_size, step lr with L2 weight_decay) from a deterministic seeded init, projecting each entity embedding back to the unit ball every epoch so the scale – hence the step size – stays well behaved. One optimize / fit iteration (max_its=1) trains the model; the data is supplied through the accumulator like any other estimator.

Parameters:

• num_entities (int)

• num_relations (int)

• dim (int)

• lr (float)

• epochs (int)

• batch_size (int)

• weight_decay (float)

• init_scale (float)

• max_norm (float)

• directions (tuple)

• negatives (int | None)

• seed (int)

• pseudo_count (float | None)

• name (str | None)

• keys (str | None)

accumulator_factory()

Return a factory for stored-triple accumulators.

Return type:

KnowledgeGraphAccumulatorFactory

estimate(nobs, suff_stat)

Fit DistMult embeddings from stored triples and weights.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

KnowledgeGraphDistribution

class KnowledgeGraphDataEncoder

Bases: DataSequenceEncoder

Encode a sequence of (h, r, t) triples into an (N, 3) integer array.

seq_encode(x)

Validate and encode triples as an (N, 3) integer array.

Parameters:

x (Sequence[Sequence[int]])

Return type:

ndarray

class KnowledgeGraphEnsemble(members)

Bases: object

An ensemble of independently fit KnowledgeGraphDistribution models, for epistemic (model) uncertainty over completions.

The members share the entity and relation index spaces but are fit from different random seeds, so where the data pins the answer down they agree and where it does not they disagree. The mean tail posterior averages p(t | h, r) across members; the epistemic uncertainty is the mutual information (BALD) H(mean) - mean_m H(member_m) – the part of the predictive entropy that comes from disagreement among members rather than from genuine ambiguity.

Parameters:

members (list[KnowledgeGraphDistribution])

mean_tail_posterior(h, r)

The ensemble-averaged p(t | h, r) over all tail candidates.

Parameters:

• h (int)

• r (int)

Return type:

ndarray

epistemic_tail_uncertainty(h, r)

Mutual-information (BALD) epistemic uncertainty of the tail completion (nats); 0 if members agree.

Thin wrapper over the general mixle.inference.uncertainty.decompose_entropy() – the tail posteriors p(t | h, r) per member are exactly the categorical predictives it splits.

Parameters:

• h (int)

• r (int)

Return type:

float

fit_knowledge_graph_ensemble(triples, num_entities, num_relations, dim=16, members=5, bootstrap=False, rng=None, **estimator_kwargs)

Fit members knowledge-graph models and wrap them in an ensemble.

Members differ by their random seed; with bootstrap=True each is also fit on a bootstrap resample of the triples (bagging), which spreads the members further apart where the data is thin and so sharpens the epistemic-uncertainty estimate.

Parameters:

• triples (Sequence[Sequence[int]])

• num_entities (int)

• num_relations (int)

• dim (int)

• members (int)

• bootstrap (bool)

• rng (Any)

• estimator_kwargs (Any)

Return type:

KnowledgeGraphEnsemble

class KnowledgeGraphPattern(kg, pattern, candidates=None, known=None, beam=64)

Bases: object

A subgraph-pattern query over a fitted KnowledgeGraphDistribution.

A pattern is a list of triples whose slots are fixed integer ids or named variables (strings starting with '?'); a variable may recur across edges (shared join), and a variable in the relation slot ranges over relations, otherwise over entities. A binding assigns every variable a value; its joint score is the sum over edges of log p(tail | head, relation).

enumerate returns the most plausible completed subgraphs, and enumerator yields them lazily in descending score (a best-first beam of width beam), so the object also satisfies the structure-distribution interface (log_density + enumerator) and can be handed to ConformalStructure for a calibrated set of completed subgraphs. A binding is represented as a tuple of values in the canonical (sorted) variable order; binding() builds one from a dict and triples() grounds it to edges.

Parameters:

• kg (KnowledgeGraphDistribution)

• pattern (Any)

• candidates (Any)

• known (Any)

• beam (int)

binding(assignment)

Canonical binding tuple (sorted-variable order) from a {variable: value} dict.

Parameters:

assignment (dict)

Return type:

tuple

triples(binding)

Ground a binding tuple to the list of completed (h, r, t) edges.

Parameters:

binding (tuple)

Return type:

list[tuple]

log_density(binding)

Joint log-probability of a complete binding (sum of edge tail-conditional log-probs).

Parameters:

binding (tuple)

Return type:

float

enumerator()

Yield (binding, joint_log_prob) over completed subgraphs in descending score (beam-limited).

enumerate(top_n=10)

Top completed subgraphs as [({variable: value}, [edges], joint_log_prob), ...].

Parameters:

top_n (int | None)

Return type:

list[tuple[dict, list[tuple], float]]