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. 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 posterior-mean probability of one attention observation.

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 synthetic latent-embedding attention observations.

Parameters:

seed (int | None)

Return type:

VariationalEmbeddingAttentionSampler

estimator(pseudo_count=None)

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

Parameters:

pseudo_count (float | None)

Return type:

VariationalEmbeddingAttentionEstimator

dist_to_encoder()

Return the encoder for contexts, queries, and targets.

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 one observation or size iid synthetic observations.

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)

Update ELBO gradients and closed-form count statistics from encoded observations.

Parameters:

estimate (VariationalEmbeddingAttentionDistribution)

Return type:

None

update(x, weight, estimate)

Update from one weighted latent-embedding attention observation.

Parameters:

weight (float)

Return type:

None

seq_initialize(x, weights, rng)

Initialize emission and position counts with random attention responsibilities.

Parameters:

rng (RandomState)

Return type:

None

initialize(x, weight, rng)

Initialize from one weighted latent-embedding attention observation.

Parameters:

• weight (float)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Merge ELBO gradients, emission counts, position counts, and scalar totals.

Return type:

VariationalEmbeddingAttentionAccumulator

value()

Return accumulated gradients, count statistics, log-likelihood, and total weight.

from_value(x)

Restore accumulator state from value output.

Return type:

VariationalEmbeddingAttentionAccumulator

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 this accumulator.

Return type:

VariationalEmbeddingAttentionDataEncoder

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

Bases: StatisticAccumulatorFactory

Create accumulators for variational embedding attention EM steps.

Parameters:

estimator (VariationalEmbeddingAttentionEstimator)

make()

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

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 a factory for variational embedding attention accumulators.

Return type:

VariationalEmbeddingAttentionAccumulatorFactory

estimate(nobs, suff_stat)

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

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 (context_symbols, query, target) observations.

Parameters:

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

Return type:

tuple[ndarray, ndarray, ndarray]