mixle.stats.latent.variational_embedding_attention module

Variational-EM responsibility attention with latent, tied token embeddings.

The plain responsibility_attention head learns a key vector per symbol from an observed query, by closed-form EM. To learn a genuinely tied embedding – one shared latent vector e_s used in both the query and the key role – the embedding has to become a latent variable, because as a parameter its M-step is non-closed-form (it appears on both sides of the gate). Its posterior is intractable, so it is approximated variationally with a mean-field q(e_s) = N(m_s, v_s). The discrete attention stays an exact-EM latent.

Generative model for one observation (context_symbols, query_symbol, target):

e_s ~ N(0, I) # latent tied embedding per symbol a_i ∝ pi_i exp(-||e_q - e_{c_i}||^2 / 2 sigma2) # generative gate: attend where embeddings match t ~ Categorical( sum_i a_i emission[c_i, :] ) # target read from the attended symbol

Fitting is variational EM, run through the ordinary mixle.inference.optimize() loop:

• the variational E-step takes one reparameterized-ELBO gradient step on the embedding posterior (m, v) – the per-observation ELBO-gradient is additive, so it accumulates like any sufficient statistic; the Adam optimizer state lives on the estimator (which persists across EM iterations), keeping the accumulator purely additive;

• the M-step updates the emission and position prior in closed form.

Because tying ties the query and key roles, a symbol’s embedding learned from its appearances as a key (in context) also serves as its query – so the model transfers a representation across roles, which a lookup/one-hot query cannot. Honest caveats: the objective is an ELBO (a bound, monotone in the bound, not the exact likelihood); the embedding posterior is a global latent fit by an inner gradient step, so this estimator is single-process for the embedding update (the discrete-attention / emission / prior parts remain ordinary additive EM).

References: variational attention as a latent variable (Deng, Kim, Chiu, Guo & Rush 2018); latent coordinates / embeddings via variational inference (Titsias & Lawrence 2010, Bayesian GP-LVM, which likewise makes nonlinearly-appearing latents tractable variationally). The reparameterized-ELBO E-step follows their recommendation of low-variance variational gradients over REINFORCE-style hard attention.

class VariationalEmbeddingAttentionDistribution(mean, log_var, emission, position_prior, sigma2=0.5, name=None)

Bases: SequenceEncodableProbabilityDistribution

Responsibility-attention head over tied latent embeddings (mean-field posterior).

Parameters:

• mean (np.ndarray)

• log_var (np.ndarray)

• emission (np.ndarray)

• position_prior (np.ndarray)

• 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.

Parameters:

x (tuple[Any, int, int])

Return type:

float

log_density(x)

Plug-in (posterior-mean) log conditional density log p(target | context, query).

Parameters:

x (tuple[Any, int, int])

Return type:

float

seq_log_density(x)

Posterior-mean log density over an encoded batch -> (n,) (uses e = m).

Parameters:

x (tuple[ndarray, ndarray, ndarray])

Return type:

ndarray

predict_proba(context, query)

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

Parameters:

• context (ndarray)

• query (ndarray)

Return type:

ndarray

embeddings()

The learned (posterior-mean) tied embedding table (S, D).

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

VariationalEmbeddingAttentionSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

VariationalEmbeddingAttentionEstimator

dist_to_encoder()

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

Return type:

VariationalEmbeddingAttentionDataEncoder

class VariationalEmbeddingAttentionSampler(dist, seed=None)

Bases: DistributionSampler

Joint generative sampler (draws embeddings from the posterior + a uniform-distinct context).

Parameters:

• dist (VariationalEmbeddingAttentionDistribution)

• 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 VariationalEmbeddingAttentionAccumulator(num_symbols, context_length, embed_dim, num_targets, mc, seed, keys=None, name=None)

Bases: SequenceEncodableStatisticAccumulator

Accumulates the additive ELBO-gradient + emission/position sufficient statistics.

seq_update(x, weights, estimate)

Parameters:

estimate (VariationalEmbeddingAttentionDistribution)

Return type:

None

update(x, weight, estimate)

Parameters:

weight (float)

Return type:

None

seq_initialize(x, weights, rng)

Parameters:

rng (RandomState)

Return type:

None

initialize(x, weight, rng)

Parameters:

• weight (float)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Return type:

VariationalEmbeddingAttentionAccumulator

value()

from_value(x)

Return type:

VariationalEmbeddingAttentionAccumulator

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:

VariationalEmbeddingAttentionDataEncoder

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

Bases: StatisticAccumulatorFactory

Parameters:

estimator (VariationalEmbeddingAttentionEstimator)

make()

Return type:

VariationalEmbeddingAttentionAccumulator

class VariationalEmbeddingAttentionEstimator(num_symbols, context_length, embed_dim, num_targets, *, sigma2=0.5, lr=0.05, mc=6, prior_strength=1.0, emission_smoothing=1e-4, seed=0, name=None, keys=None)

Bases: ParameterEstimator

Variational-EM estimator; holds the embedding posterior + Adam state across EM iterations.

Parameters:

• num_symbols (int)

• context_length (int)

• embed_dim (int)

• num_targets (int)

• sigma2 (float)

• lr (float)

• mc (int)

• prior_strength (float)

• emission_smoothing (float)

• seed (int)

• name (str | None)

• keys (str | None)

accumulator_factory()

Return type:

VariationalEmbeddingAttentionAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

nobs (float | None)

Return type:

VariationalEmbeddingAttentionDistribution

class VariationalEmbeddingAttentionDataEncoder

Bases: DataSequenceEncoder

Encodes (context_symbols, query_symbol, target) triples into stacked integer arrays.

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

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

Return type:

tuple[ndarray, ndarray, ndarray]