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 is carried explicitly by the model and sufficient-statistic artifact, so distributed and resumed updates have the same semantics;

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

• optimizer_state (AttentionOptimizerState | dict[str, Any] | 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)

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

Parameters:

seed (int | None)

Return type:

VariationalEmbeddingAttentionSampler

posterior_predictive_sampler(seed=None)

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

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, *, posterior_predictive=False)

Bases: DistributionSampler

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

Parameters:

• dist (VariationalEmbeddingAttentionDistribution)

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

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

Stateless variational-EM estimator consuming explicit model-carried Adam state.

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(num_symbols=None, context_length=None, num_targets=None)

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]

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