mixle.stats.latent.responsibility_attention module

Responsibility attention: one attention head as an EM-able mixture over context positions.

A single attention head, written as a generative latent-variable model. The latent z is which context position you attend to; the “gate” is generative (the query is generated from the attended token’s key) so the attention weights are posterior responsibilities rather than a discriminative dot-product softmax. That one change is what makes the head fully EM-able: every M-step is a closed-form, additively-mergeable sufficient-statistic update (a GMM mean for the keys, a weighted categorical for the emission, counts for the position prior), so the head trains by expectation- maximization and composes with the rest of mixle (drop it into a mixture, an HMM emission, a composite) instead of needing gradient descent.

Generative story for one observation (context, query, target) with context tokens c_1..c_N:

z ~ Categorical(position_prior) query ~ Normal(key_means[c_z], sigma2 I) # query generated from the attended token’s key target ~ Categorical(emission[c_z, :]) # target emitted from the attended token

so p(query, target | context) = sum_i position_prior_i N(query; key_means[c_i]) emission[c_i, target] and the attention weight on position i is the posterior p(z = i | query, target, context).

The observation is the triple (context, query, target); the context is conditioned on (it is the covariate that defines the positions), so the density is the conditional p(query, target | context) that EM maximizes. This is the single-hop leaf; deeper / stackable multi-hop versions compose several of these latents (a chain is forward–backward, an HME-style tree is nested responsibilities).

This is not a new idea – it is the EM-able / generative reading of attention. Closely related prior work: latent-alignment EM (Brown et al. 1993, IBM models; Vogel & Ney 1996, HMM word alignment), attention-as-kernel-smoothing (Tsai et al. 2019; Ramsauer et al. 2020, modern Hopfield), hierarchical mixtures of experts (Jordan & Jacobs 1994), and attention-as-a-latent-variable with variational inference (Deng, Kim, Chiu, Guo & Rush 2018, “Latent Alignment and Variational Attention”). The contribution here is packaging it as a composable, EM-fit mixle Distribution primitive – which is exactly the “compose attention with probabilistic models / do posterior inference” gap Deng et al. identify – and the empirical confirmation that exact marginalization (our closed-form E-step) is the right call, matching their finding that exact latent models beat soft attention.

class ResponsibilityAttentionDistribution(key_means, emission, position_prior=None, sigma2=1.0, name=None)

Bases: SequenceEncodableProbabilityDistribution

One generative-gate attention head; a mixture over context positions (EM-able).

Parameters:

• key_means (np.ndarray)

• emission (np.ndarray)

• position_prior (np.ndarray | None)

• sigma2 (float)

• name (str | None)

density(x)

Density p(query, target | context) at one observation.

Parameters:

x (tuple[Any, Any, int])

Return type:

float

log_density(x)

Log conditional density log p(query, target | context) at one observation (ctx, y, t).

Parameters:

x (tuple[Any, Any, int])

Return type:

float

seq_log_density(x)

Vectorized log p(query, target | context) over an encoded batch -> (n,).

Parameters:

x (tuple[ndarray, ndarray, ndarray])

Return type:

ndarray

predict_proba(context, query)

Predictive target distribution p(target | context, query) (target marginalized over z).

Attention here does not see the target: a_i ∝ pi_i N(query; K[c_i]), then p(target) = sum_i a_i emission[c_i, :]. Accepts a single example or a batch.

Returns:

(T,) for a single example, or (n, T) for a batch.

Parameters:

• context (ndarray)

• query (ndarray)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

ResponsibilityAttentionSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

ResponsibilityAttentionEstimator

dist_to_encoder()

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

Return type:

ResponsibilityAttentionDataEncoder

class ResponsibilityAttentionSampler(dist, seed=None)

Bases: DistributionSampler

Joint generative sampler (adds a uniform-distinct context prior so it is a full generator).

Parameters:

• dist (ResponsibilityAttentionDistribution)

• 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 ResponsibilityAttentionAccumulator(num_symbols, context_length, query_dim, num_targets, keys=None, name=None)

Bases: SequenceEncodableStatisticAccumulator

Accumulates the additive sufficient statistics for the responsibility-attention M-step.

Parameters:

• num_symbols (int)

• context_length (int)

• query_dim (int)

• num_targets (int)

• keys (str | None)

• name (str | None)

update(x, weight, estimate)

Parameters:

• weight (float)

• estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

seq_update(x, weights, estimate)

Parameters:

• weights (ndarray)

• estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

initialize(x, weight, rng)

Parameters:

• weight (float)

• rng (RandomState)

Return type:

None

seq_initialize(x, weights, rng)

Parameters:

• weights (ndarray)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Return type:

ResponsibilityAttentionAccumulator

value()

from_value(x)

Return type:

ResponsibilityAttentionAccumulator

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:

ResponsibilityAttentionDataEncoder

class ResponsibilityAttentionAccumulatorFactory(num_symbols, context_length, query_dim, num_targets, keys=None, name=None)

Bases: StatisticAccumulatorFactory

Parameters:

• num_symbols (int)

• context_length (int)

• query_dim (int)

• num_targets (int)

• keys (str | None)

• name (str | None)

make()

Return type:

ResponsibilityAttentionAccumulator

class ResponsibilityAttentionEstimator(num_symbols, context_length, query_dim, num_targets, *, sigma2=1.0, estimate_sigma2=False, min_sigma2=1e-6, emission_smoothing=1e-6, pseudo_count=None, name=None, keys=None)

Bases: ParameterEstimator

Estimate a ResponsibilityAttentionDistribution by closed-form EM M-steps.

Parameters:

• num_symbols (int)

• context_length (int)

• query_dim (int)

• num_targets (int)

• sigma2 (float)

• estimate_sigma2 (bool)

• min_sigma2 (float)

• emission_smoothing (float)

• pseudo_count (float | None)

• name (str | None)

• keys (str | None)

accumulator_factory()

Return type:

ResponsibilityAttentionAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

nobs (float | None)

Return type:

ResponsibilityAttentionDistribution

class ResponsibilityAttentionDataEncoder

Bases: DataSequenceEncoder

Encodes a sequence of (context, query, target) triples into stacked arrays.

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

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

Return type:

tuple[ndarray, ndarray, ndarray]

sequence_to_triples(tokens, context_length, *, num_symbols=None, embeddings=None)

Unroll a token sequence into the (context, query, target) triples this head consumes.

For each prediction point p (with a full window behind it and a next token ahead), emits context = tokens[p-N+1 : p+1] (the N most recent tokens, including the current one), query derived from the current token tokens[p] (its embedding, or a one-hot), and target = tokens[p+1]. This is the bridge from a real sequence to the iid-triple leaf.

Honest scope: with the single-hop leaf this yields an attention-weighted next-token model whose predictive is a function of the current token (an attention-flavoured bigram) – because the emission is keyed by the attended token’s identity, the prediction depends only on which token the query selects. In-context copy / induction (example-specific values) is what the multi-hop / stackable extension adds; this helper is the plumbing both share.

Parameters:

• tokens (Sequence[int]) – a 1-D sequence of integer token ids.

• context_length (int) – the attention window N.

• num_symbols (int | None) – vocabulary size (for the one-hot query dim); inferred from tokens if omitted.

• embeddings (ndarray | None) – optional (S, D) embedding table; if given the query is embeddings[token], otherwise a one-hot of dimension num_symbols.

Returns:

A list of (context, query, target) triples ready for mixle.inference.optimize().

Return type:

list[tuple[ndarray, ndarray, int]]