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)[source]

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)[source]

Return the probability of one chained-attention observation.

Parameters:

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

Return type:

float

log_density(x)[source]

Return the log-probability of one context/query/target observation.

Parameters:

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

Return type:

float

seq_log_density(x)[source]

Return vectorized log-probabilities for encoded chained-attention observations.

Return type:

ndarray

predict_proba(context_keys, context_values, query)[source]

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

Parameters:
Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for synthetic chained-attention observations.

Parameters:

seed (int | None)

Return type:

ChainedAttentionSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

ChainedAttentionEstimator

dist_to_encoder()[source]

Return the encoder for context keys, values, query symbols, and targets.

Return type:

ChainedAttentionDataEncoder

class ChainedAttentionSampler(dist, seed=None)[source]

Bases: DistributionSampler

Generative sampler (uniform context + the chained gate).

Parameters:

  • dist (ChainedAttentionDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid synthetic observations.

Parameters:
Return type:

Any

class ChainedAttentionAccumulator(n_hops, num_symbols, num_targets, keys=None, name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

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

seq_update(x, weights, estimate)[source]

Update forward-backward sufficient statistics from encoded observations.

Parameters:

estimate (ChainedAttentionDistribution)

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics with random hop responsibilities.

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)[source]

Update from one weighted chained-attention observation.

Parameters:

weight (float)

Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted chained-attention observation.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge key-table, emission, likelihood, and weight statistics.

Return type:

ChainedAttentionAccumulator

value()[source]

Return key statistics, emission counts, log-likelihood, and total weight.

from_value(x)[source]

Restore accumulator state from value output.

Return type:

ChainedAttentionAccumulator

key_merge(stats_dict)[source]

Merge this accumulator into stats_dict under its configured key.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace this accumulator’s state from keyed statistics when present.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return the encoder compatible with this accumulator.

Return type:

ChainedAttentionDataEncoder

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

Bases: StatisticAccumulatorFactory

Create accumulators for chained-attention EM statistics.

make()[source]

Create an empty chained-attention accumulator.

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)[source]

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()[source]

Return a factory for chained-attention sufficient-statistic accumulators.

Return type:

ChainedAttentionAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate key tables and emissions from accumulated forward-backward statistics.

Parameters:

nobs (float | None)

Return type:

ChainedAttentionDistribution

class ChainedAttentionDataEncoder[source]

Bases: DataSequenceEncoder

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

seq_encode(x)[source]

Encode (context_keys, context_values, query, target) observations.

Parameters:

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

Return type:

tuple[ndarray, ndarray, ndarray, ndarray]