mixle.stats.latent.variational_multihop_attention module

Variational multi-hop attention: a 2-hop chain over TIED latent embeddings, with prior annealing.

This combines the two hard pieces: the multi-hop chain (chained_attention) and tied latent embeddings (variational_embedding_attention). Each context position is a (key, value); a single latent embedding e_s per symbol is used in every role (query, key, value). Hop 1 attends the query embedding to the key embeddings; the attended position’s value embedding becomes the hop-2 query; hop 2 attends again; the target is emitted from the final attended value. The two hop latents are summed exactly (an N x N table); the embeddings are latent with a mean-field posterior q(e_s)=N(m_s, v_s) fit by a reparameterized-ELBO gradient step (the embedding M-step has no closed form – the softmax partition supplies the repulsion that prevents collapse, and it is not quadratic).

Because tying makes identity matching trivial, the N(0,I) prior would otherwise collapse the unused embeddings; the estimator anneals the prior weight from ~0 upward over EM iterations so the data spreads the embeddings first. Observation: (context_keys, context_values, query_symbol, target).

References: multi-hop attention = Memory Networks (Sukhbaatar et al. 2015); attention as a variational latent variable = Deng et al. 2018. The annealing is the practical face of Deterministic Annealing EM (Ueda & Nakano 1998) – tempering the objective to escape the bad (collapsed) fixed point and reach an initialization-independent solution. (We checked: principled DAEM tempering does not improve the closed-form chained head, which is already at its initialization-independent global optimum; the annealing is only load-bearing here, where the latent-embedding prior creates the collapse basin.)

class VariationalMultiHopAttentionDistribution(mean, log_var, emission, sigma2=0.3, name=None)

Bases: SequenceEncodableProbabilityDistribution

A 2-hop chain over tied latent embeddings (mean-field posterior).

Parameters:

• sigma2 (float)

• name (str | None)

density(x)

Return the probability density or mass at a single observation.

Concrete default: exponentiate log_density (the abstract method subclasses must provide). Leaves with a cheaper closed form may override this.

Return type:

float

log_density(x)

Return the log-density or log-mass at a single observation.

Return type:

float

seq_log_density(x)

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

Return type:

ndarray

predict_proba(context_keys, context_values, query)

Predictive target distribution (posterior-mean embeddings); (T,) or (n, T).

Return type:

ndarray

embeddings()

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

VariationalMultiHopAttentionSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

VariationalMultiHopAttentionEstimator

dist_to_encoder()

Return the data encoder used by this distribution for vectorized methods.

Return type:

VariationalMultiHopAttentionDataEncoder

class VariationalMultiHopAttentionSampler(dist, seed=None)

Bases: DistributionSampler

Parameters:

seed (int | None)

sample(size=None, *, batched=True)

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)

• batched (bool)

Return type:

Any

class VariationalMultiHopAttentionAccumulator(num_symbols, embed_dim, num_targets, mc, seed, keys=None, name=None)

Bases: SequenceEncodableStatisticAccumulator

seq_update(x, weights, estimate)

Return type:

None

seq_initialize(x, weights, rng)

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)

Return type:

None

initialize(x, weight, rng)

Return type:

None

combine(suff_stat)

value()

from_value(x)

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.

Return type:

None

key_replace(stats_dict)

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

Return type:

None

acc_to_encoder()

class VariationalMultiHopAttentionAccumulatorFactory(estimator, keys=None, name=None)

Bases: StatisticAccumulatorFactory

make()

class VariationalMultiHopAttentionEstimator(num_symbols, embed_dim, num_targets, *, sigma2=0.3, lr=0.05, mc=5, prior_strength=0.1, anneal_iters=100, emission_smoothing=1e-4, seed=0, name=None, keys=None)

Bases: ParameterEstimator

Variational-EM estimator with prior annealing (KL weight ramped over EM iterations).

Parameters:

• num_symbols (int)

• embed_dim (int)

• num_targets (int)

• sigma2 (float)

• lr (float)

• mc (int)

• prior_strength (float)

• anneal_iters (int)

• emission_smoothing (float)

• seed (int)

• name (str | None)

• keys (str | None)

accumulator_factory()

estimate(nobs, suff_stat)

class VariationalMultiHopAttentionDataEncoder

Bases: DataSequenceEncoder

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

x (Sequence[tuple[Any, Any, int, int]])