mixle.stats.latent.chained_attention module

Chained (multi-hop) attention: an L-hop stack of responsibility-attention, EM-able via forward-backward.

A single attention head answers one content lookup. Chaining L of them answers a multi-step lookup – “find b that a points to, then find c that b points to” – which one hop provably cannot do. Each position in the context carries a (key_symbol, value_symbol) pair. The initial query is the one-hot of the query symbol; at hop l it attends over positions by matching the current query to the hop’s key table K^(l), and the attended position’s value (as a one-hot) becomes the query for hop l+1. After L hops the target is emitted from the final attended value.

The hop latents z_1..z_L (which position each hop lands on) form a time-inhomogeneous chain over the N context positions – a content-addressed HMM with an emission only at the last hop – so forward-backward gives exact responsibilities in O(L N^2) and every M-step stays closed-form and additively mergeable. Crucially the queries are one-hot (observed symbols / retrieved values), so they anchor the key tables and the closed-form M-step does not collapse (unlike a tied latent embedding, which needs the variational treatment).

Observation: (context_keys, context_values, query_symbol, target). n_hops = 1 recovers a single responsibility-attention head; n_hops = 2 does transitive (a->b->c) lookup. The gate variance sigma2 is the attention temperature – keep it small (the one-hot query/key separation is only 2, so a large sigma2 blurs the chained prediction).

References: multi-hop content-addressed attention is the End-to-End Memory Network (Sukhbaatar, Szlam, Weston & Fergus 2015); the transitive lookup is a bAbI-style reasoning task and the 2-hop copy is the induction-head circuit (Olsson et al. 2022). Exact forward-backward over the alignment chain mirrors HMM-based word alignment (Vogel & Ney 1996), which – like here – keeps the marginal likelihood exact.

class ChainedAttentionDistribution(keys, emission, sigma2=0.1, name=None)

Bases: SequenceEncodableProbabilityDistribution

An L-hop stack of responsibility-attention heads (chained, content-addressed).

Parameters:

• keys (np.ndarray)

• emission (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, Any, int, int])

Return type:

float

log_density(x)

Return the log-density or log-mass at a single observation.

Parameters:

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

Return type:

float

seq_log_density(x)

Return vectorized log-density values for sequence-encoded observations.

Return type:

ndarray

predict_proba(context_keys, context_values, query)

Predictive target distribution (target marginalized); (T,) or (n, T).

Parameters:

• context_keys (ndarray)

• context_values (ndarray)

• query (ndarray)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

ChainedAttentionSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

ChainedAttentionEstimator

dist_to_encoder()

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

Return type:

ChainedAttentionDataEncoder

class ChainedAttentionSampler(dist, seed=None)

Bases: DistributionSampler

Generative sampler (uniform context + the chained gate).

Parameters:

• dist (ChainedAttentionDistribution)

• 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 ChainedAttentionAccumulator(n_hops, num_symbols, num_targets, keys=None, name=None)

Bases: SequenceEncodableStatisticAccumulator

Forward-backward sufficient statistics (additive): per-hop key numerators + emission counts.

seq_update(x, weights, estimate)

Parameters:

estimate (ChainedAttentionDistribution)

Return type:

None

seq_initialize(x, weights, rng)

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)

Parameters:

weight (float)

Return type:

None

initialize(x, weight, rng)

Parameters:

• weight (float)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Return type:

ChainedAttentionAccumulator

value()

from_value(x)

Return type:

ChainedAttentionAccumulator

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:

ChainedAttentionDataEncoder

class ChainedAttentionAccumulatorFactory(n_hops, num_symbols, num_targets, keys=None, name=None)

Bases: StatisticAccumulatorFactory

make()

Return type:

ChainedAttentionAccumulator

class ChainedAttentionEstimator(n_hops, num_symbols, num_targets, *, sigma2=0.1, emission_smoothing=1e-4, pseudo_count=None, name=None, keys=None)

Bases: ParameterEstimator

Closed-form EM estimator: per-hop key tables (GMM means of one-hot queries) + emission counts.

Parameters:

• n_hops (int)

• num_symbols (int)

• num_targets (int)

• sigma2 (float)

• emission_smoothing (float)

• pseudo_count (float | None)

• name (str | None)

• keys (str | None)

accumulator_factory()

Return type:

ChainedAttentionAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

nobs (float | None)

Return type:

ChainedAttentionDistribution

class ChainedAttentionDataEncoder

Bases: DataSequenceEncoder

Encodes (context_keys, context_values, query_symbol, target) into stacked integer arrays.

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

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

Return type:

tuple[ndarray, ndarray, ndarray, ndarray]