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

Parameters:

seed (int | None)

Return type:

ResponsibilityAttentionSampler

estimator(pseudo_count=None)

Return a closed-form EM estimator for this attention head.

Parameters:

pseudo_count (float | None)

Return type:

ResponsibilityAttentionEstimator

dist_to_encoder()

Return the encoder for contexts, query vectors, and targets.

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

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)

Update sufficient statistics from one weighted observation.

Parameters:

• weight (float)

• estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

seq_update(x, weights, estimate)

Update sufficient statistics from encoded observations.

Parameters:

• weights (ndarray)

• estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

initialize(x, weight, rng)

Initialize sufficient statistics from one weighted observation.

Parameters:

• weight (float)

• rng (RandomState)

Return type:

None

seq_initialize(x, weights, rng)

Initialize sufficient statistics with random attention responsibilities.

Parameters:

• weights (ndarray)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Merge key, emission, position, and query-scatter statistics.

Return type:

ResponsibilityAttentionAccumulator

value()

Return a snapshot of additive attention sufficient statistics.

from_value(x)

Restore attention sufficient statistics from value output.

Return type:

ResponsibilityAttentionAccumulator

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:

ResponsibilityAttentionDataEncoder

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

Bases: StatisticAccumulatorFactory

Create accumulators for responsibility-attention EM statistics.

Parameters:

• num_symbols (int)

• context_length (int)

• query_dim (int)

• num_targets (int)

• keys (str | None)

• name (str | None)

make()

Create an empty responsibility-attention accumulator.

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 a factory for responsibility-attention sufficient-statistic accumulators.

Return type:

ResponsibilityAttentionAccumulatorFactory

estimate(nobs, suff_stat)

Estimate key means, emissions, position prior, and optional gate variance.

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

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.

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