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 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, optimizer_state=None)

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

• sigma2 (float)

• name (str | None)

• optimizer_state (AttentionOptimizerState | dict[str, Any] | None)

density(x)

Return the probability of one context/query/target observation.

Return type:

float

log_density(x)

Return the log-probability of one context/query/target observation.

Return type:

float

seq_log_density(x)

Return vectorized log-probabilities for encoded two-hop attention 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 posterior mean embeddings for the tied latent symbols.

Return type:

ndarray

sampler(seed=None)

Sample iid observations from the posterior-mean plug-in law used by scoring.

Parameters:

seed (int | None)

Return type:

VariationalMultiHopAttentionSampler

posterior_predictive_sampler(seed=None)

Sample iid observations, independently drawing an embedding table for every event.

Parameters:

seed (int | None)

Return type:

VariationalMultiHopAttentionSampler

estimator(pseudo_count=None)

Return a variational EM estimator initialized with this model’s dimensions.

Parameters:

pseudo_count (float | None)

Return type:

VariationalMultiHopAttentionEstimator

dist_to_encoder()

Return the encoder for context keys, values, query symbols, and targets.

Return type:

VariationalMultiHopAttentionDataEncoder

class VariationalMultiHopAttentionSampler(dist, seed=None, *, posterior_predictive=False)

Bases: DistributionSampler

Sampler for either the plug-in or explicitly requested posterior-predictive law.

Parameters:

• seed (int | None)

• posterior_predictive (bool)

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

Draw one observation or size iid synthetic observations.

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

Accumulate Monte-Carlo ELBO gradients and emission responsibilities for variational EM.

seq_update(x, weights, estimate)

Update ELBO gradients and emission counts from encoded observations.

Return type:

None

seq_initialize(x, weights, rng)

Initialize emission counts with random final-hop responsibilities.

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)

Update from one weighted two-hop attention observation.

Return type:

None

initialize(x, weight, rng)

Initialize from one weighted two-hop attention observation.

Return type:

None

combine(suff_stat)

Merge variational gradients, emission counts, log-likelihood, and weight totals.

value()

Return accumulated gradients, emission counts, log-likelihood, and total weight.

from_value(x)

Restore accumulator state from value output.

acc_to_encoder()

Return the encoder compatible with this attention accumulator.

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

Bases: StatisticAccumulatorFactory

Create accumulators for variational multi-hop attention EM steps.

make()

Create an accumulator with a deterministic per-iteration Monte-Carlo seed.

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

Stateless variational-EM estimator with model-carried Adam state and prior annealing.

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()

Return a factory for variational multi-hop attention accumulators.

estimate(nobs, suff_stat)

Apply one variational EM update and return the updated attention distribution.

class VariationalMultiHopAttentionDataEncoder(num_symbols=None, num_targets=None)

Bases: DataSequenceEncoder

Encode context keys, context values, query symbols, and targets as integer arrays.

seq_encode(x)

Encode (context_keys, context_values, query, target) observations.

Parameters:

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

row_count(x)

Return the number of observations represented by x.

The default covers simple array/list encodings and structurally aligned tuples/dicts. Encoders with flattened, ragged, or otherwise specialized layouts must override this method so metadata wrappers can verify counts instead of trusting a caller assertion.

Return type:

int