mixle.stats.graphs.knowledge_graph module

Create, estimate, and sample from a knowledge-graph embedding distribution.

Defines KnowledgeGraphDistribution, KnowledgeGraphSampler, KnowledgeGraphAccumulatorFactory, KnowledgeGraphAccumulator, KnowledgeGraphEstimator, and KnowledgeGraphDataEncoder for use with the regular mixle estimation framework (optimize / seq_log_density / the PPL surface), like the other graph distributions.

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 the log-density or log-mass at a single observation.

Parameters:

x (Sequence[int])

Return type:

float

seq_log_density(x)

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (ndarray)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

KnowledgeGraphSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

KnowledgeGraphEstimator

dist_to_encoder()

Return the data encoder used by this distribution for 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 observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=True (the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

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)

Parameters:

• x (Sequence[int])

• weight (float)

• estimate (KnowledgeGraphDistribution | None)

Return type:

None

initialize(x, weight, rng)

Parameters:

• x (Sequence[int])

• weight (float)

• rng (RandomState | None)

Return type:

None

seq_initialize(x, weights, rng)

Parameters:

• x (ndarray)

• weights (ndarray)

• rng (RandomState | None)

Return type:

None

seq_update(x, weights, estimate)

Parameters:

• x (ndarray)

• weights (ndarray)

• estimate (KnowledgeGraphDistribution | None)

Return type:

None

combine(suff_stat)

Parameters:

suff_stat (tuple)

Return type:

KnowledgeGraphAccumulator

value()

Return type:

tuple

from_value(x)

Parameters:

x (tuple)

Return type:

KnowledgeGraphAccumulator

key_merge(stats_dict)

Pool this accumulator’s statistics into stats_dict under its merge key.

The structural default implements the common single-key pattern: store the accumulator under self.keys the first time the key is seen, else combine into the one already there. Accumulators with several named keys (e.g. an HMM’s init/trans/state keys) or a non-accumulator stats payload override this. A keys of None (the default) is a no-op.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

Replace this accumulator’s statistics from the pooled stats_dict entry (see key_merge).

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Return type:

KnowledgeGraphDataEncoder

class KnowledgeGraphAccumulatorFactory(keys=None)

Bases: StatisticAccumulatorFactory

Factory for KnowledgeGraphAccumulator.

Parameters:

keys (str | None)

make()

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 type:

KnowledgeGraphAccumulatorFactory

estimate(nobs, suff_stat)

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)

Encode the iid observation sequence x for vectorized evaluation.

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]]