mixle.stats.latent.gated_mixture module

Gated mixture (mixture of experts): mixing weights are a learned function of a covariate, not constants.

MixtureDistribution mixes K components with FIXED weights w_k. A gated mixture replaces those constants with a gate p(k | z) – a function of a per-observation covariate z – so the mixture that explains y shifts smoothly as z moves. That is the classic mixture-of-experts (Jacobs et al. 1991): each component is an “expert” over y, the gate routes probability mass among them by z.

An observation is a pair (z, y): z drives the gate, y is scored by the experts. The density is

p(y | z) = sum_k gate_k(z) * f_k(y), gate_k(z) = softmax over experts of the gate’s logits at z.

EM is the same responsibility loop as a plain mixture, with the gate in place of the constant prior: the E-step forms r_nk ∝ gate_k(z_n) f_k(y_n); the M-step (a) refits each expert on y weighted by its responsibilities (exactly as a plain mixture does) and (b) refits the gate to predict r_nk from z_n (a soft-target multinomial regression). Unlike a plain mixture’s closed-form weight update, the gate step is an optimization, so the accumulator buffers (z, responsibilities) – the same buffer-the-rows pattern the neural leaves and CopulaDistribution use.

The gate is pluggable (any object implementing the small Gate protocol below). The default SoftmaxGate is a torch-free multinomial logistic regression, so a gated mixture needs no torch; a NeuralCategorical-backed gate can be substituted for a deep gate.

Reference: Jacobs, Jordan, Nowlan & Hinton, “Adaptive Mixtures of Local Experts” (Neural Computation, 1991).

class SoftmaxGate(weight, bias)[source]

Bases: object

A torch-free multinomial-logistic gate p(k | z) = softmax(W z + b)_k, fit on soft targets.

fit(Z, R) minimizes the soft cross-entropy -sum_{n,k} R_{n,k} log p(k | z_n) by gradient descent – R are the responsibilities (rows need not sum to 1; the sample weight is folded in).

Parameters:

  • weight (np.ndarray)

  • bias (np.ndarray)

classmethod zeros(n_classes, n_features)[source]

Create a zero-logit gate with uniform initial class probabilities.

Parameters:

  • n_classes (int)

  • n_features (int)

Return type:

SoftmaxGate

log_prob_batch(z)[source]

(n, K) log-gate log p(k | z_n) for each row of z (shape (n, p)).

Parameters:

z (ndarray)

Return type:

ndarray

fit(z, resp, *, steps=200, lr=0.1)[source]

Fit a softmax gate to responsibility-weighted soft targets.

Parameters:
Return type:

SoftmaxGate

class GatedMixtureDistribution(components, gate, name=None, keys=None)[source]

Bases: SequenceEncodableProbabilityDistribution

A mixture whose weights are a gate p(k | z); observations are (z, y) pairs.

Parameters:

  • components (Sequence[SequenceEncodableProbabilityDistribution])

  • gate (Any)

  • name (str | None)

  • keys (str | None)

log_density(x)[source]

Return log p(y | z) for one covariate/response pair.

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(enc)[source]

Return vectorized conditional log-densities for encoded (z, y) pairs.

Parameters:

enc (Any)

Return type:

ndarray

posterior(x)[source]

Return posterior expert responsibilities for one (z, y) observation.

Parameters:

x (tuple[Any, Any])

Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler that requires a covariate z.

Parameters:

seed (int | None)

Return type:

GatedMixtureSampler

estimator(pseudo_count=None)[source]

Return an EM estimator for experts and the covariate-dependent gate.

Parameters:

pseudo_count (float | None)

Return type:

GatedMixtureEstimator

dist_to_encoder()[source]

Return the encoder for covariates plus expert response encodings.

Return type:

GatedMixtureDataEncoder

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

Bases: DistributionSampler

Sample y given a supplied z: draw a component from gate(z), then sample that expert.

Parameters:

  • dist (GatedMixtureDistribution)

  • seed (int | None)

sample_given(z)[source]

Sample a response from the gated mixture conditional on covariate z.

Parameters:

z (Any)

Return type:

Any

sample(size=None)[source]

Raise because unconditional sampling requires caller-supplied covariates.

Parameters:

size (int | None)

Return type:

Any

class GatedMixtureDataEncoder(component_encoders)[source]

Bases: DataSequenceEncoder

Encode [(z, y), ...] as (z array (n, p), per-expert encodings of the y column).

Parameters:

component_encoders (Sequence[DataSequenceEncoder])

seq_encode(data)[source]

Encode covariates as a dense matrix and responses for every expert.

Parameters:

data (Sequence[tuple[Any, Any]])

Return type:

tuple[ndarray, tuple[Any, …]]

class GatedMixtureAccumulator(component_accumulators, num_components, keys=None)[source]

Bases: SequenceEncodableStatisticAccumulator

E-step responsibilities route weight to expert sub-accumulators; buffer (z, resp) for the gate.

Parameters:

  • component_accumulators (Sequence[Any])

  • num_components (int)

  • keys (str | None)

seq_update(enc, weights, estimate)[source]

Update expert accumulators and gate buffers from encoded observations.

Parameters:

  • enc (Any)

  • weights (ndarray)

  • estimate (GatedMixtureDistribution | None)

Return type:

None

seq_initialize(enc, weights, rng)[source]

Initialize expert accumulators with random responsibility allocations.

Parameters:
Return type:

None

update(x, weight, estimate)[source]

Update from one weighted (z, y) observation.

Parameters:
Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted (z, y) observation.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge expert sufficient statistics and buffered gate training data.

Parameters:

suff_stat (tuple[tuple[Any, ...], ndarray, ndarray])

Return type:

GatedMixtureAccumulator

value()[source]

Return expert statistics, buffered covariates, and responsibility targets.

Return type:

tuple[tuple[Any, …], ndarray, ndarray]

from_value(x)[source]

Restore expert statistics and gate training buffers.

Parameters:

x (tuple[tuple[Any, ...], ndarray, ndarray])

Return type:

GatedMixtureAccumulator

key_merge(stats_dict)[source]

Delegate keyed merges to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Delegate keyed replacements to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return the encoder composed from expert accumulator encoders.

Return type:

GatedMixtureDataEncoder

class GatedMixtureAccumulatorFactory(component_factories, num_components, keys=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for gated-mixture EM.

Parameters:

  • component_factories (Sequence[Any])

  • num_components (int)

  • keys (str | None)

make()[source]

Create an empty gated-mixture accumulator.

Return type:

GatedMixtureAccumulator

class GatedMixtureEstimator(component_estimators, gate, gate_steps=200, gate_lr=0.1, name=None, keys=None)[source]

Bases: ParameterEstimator

M-step: refit each expert from its responsibility-weighted stats, refit the gate on (z, resp).

Parameters:

  • component_estimators (Sequence[ParameterEstimator])

  • gate (Any)

  • gate_steps (int)

  • gate_lr (float)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return a factory for gated-mixture sufficient-statistic accumulators.

Return type:

GatedMixtureAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate experts from responsibility-weighted stats and refit the gate.

Parameters:
Return type:

GatedMixtureDistribution