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 retains a bounded deterministic coreset of (z, responsibilities) rows. The coreset is selected by stable content priority, which makes it mergeable across partitioned accumulators without unbounded memory growth; a receipt reports how many rows were retained or dropped.

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 GateOptimizationReceipt(converged, steps_completed, initial_loss, final_loss, gradient_norm, termination)

Bases: object

Outcome of a bounded softmax-gate optimization.

Parameters:

• converged (bool)

• steps_completed (int)

• initial_loss (float)

• final_loss (float)

• gradient_norm (float)

• termination (str)

class GateBufferReceipt(rows_seen, rows_retained, rows_dropped, capacity, selection='deterministic_hash_top_k')

Bases: object

Bounded deterministic coreset status for gate-training rows.

Parameters:

• rows_seen (int)

• rows_retained (int)

• rows_dropped (int)

• capacity (int)

• selection (str)

class GatedMixtureStatistics(component_statistics, covariates, responsibilities, buffer_receipt, component_counts, schema_version=1)

Bases: Sequence[Any]

Versioned gated-mixture statistics with backward four-slot iteration.

Parameters:

• component_statistics (tuple[Any, ...])

• covariates (ndarray)

• responsibilities (ndarray)

• buffer_receipt (GateBufferReceipt)

• component_counts (ndarray)

• schema_version (int)

class SoftmaxGate(weight, bias, *, fit_receipt=None)

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)

• fit_receipt (GateOptimizationReceipt | None)

classmethod zeros(n_classes, n_features)

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

Parameters:

• n_classes (int)

• n_features (int)

Return type:

SoftmaxGate

log_prob_batch(z)

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

Parameters:

z (ndarray)

Return type:

ndarray

fit_with_receipt(z, resp, *, steps=200, lr=0.1, tol=1.0e-8)

Fit a softmax gate and return explicit convergence/failure status.

Parameters:

• z (ndarray)

• resp (ndarray)

• steps (int)

• lr (float)

• tol (float)

Return type:

tuple[SoftmaxGate, GateOptimizationReceipt]

fit(z, resp, *, steps=200, lr=0.1, tol=1.0e-8)

Fit a softmax gate; the returned gate carries fit_receipt.

Parameters:

• z (ndarray)

• resp (ndarray)

• steps (int)

• lr (float)

• tol (float)

Return type:

SoftmaxGate

class GatedMixtureDistribution(components, gate, name=None, keys=None, *, gate_fit_receipt=None, gate_buffer_receipt=None)

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)

• gate_fit_receipt (GateOptimizationReceipt | None)

• gate_buffer_receipt (GateBufferReceipt | None)

log_density(x)

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

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

posterior(x)

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

Parameters:

x (tuple[Any, Any])

Return type:

ndarray

sampler(seed=None)

Return a conditional sampler that requires a covariate z.

Parameters:

seed (int | None)

Return type:

GatedMixtureSampler

estimator(pseudo_count=None)

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

Parameters:

pseudo_count (float | None)

Return type:

GatedMixtureEstimator

dist_to_encoder()

Return the encoder for covariates plus expert response encodings.

Return type:

GatedMixtureDataEncoder

class GatedMixtureSampler(dist, seed=None)

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)

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

Parameters:

z (Any)

Return type:

Any

sample(size=None, *, batched=True)

Raise because unconditional sampling requires caller-supplied covariates.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class GatedMixtureDataEncoder(component_encoders, n_features=None)

Bases: DataSequenceEncoder

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

Parameters:

• component_encoders (Sequence[DataSequenceEncoder])

• n_features (int | None)

seq_encode(data)

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, *, max_buffer_rows=10_000, n_features=None)

Bases: SequenceEncodableStatisticAccumulator

Route expert statistics and retain a bounded mergeable gate-training coreset.

Parameters:

• component_accumulators (Sequence[Any])

• num_components (int)

• keys (str | None)

• max_buffer_rows (int)

• n_features (int | None)

property buffer_receipt: GateBufferReceipt

Return current bounded-buffer retention status.

seq_update(enc, weights, estimate)

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)

Initialize expert accumulators with random responsibility allocations.

Parameters:

• enc (Any)

• weights (ndarray)

• rng (RandomState)

Return type:

None

update(x, weight, estimate)

Update from one weighted (z, y) observation.

Parameters:

• x (tuple[Any, Any])

• weight (float)

• estimate (GatedMixtureDistribution | None)

Return type:

None

initialize(x, weight, rng)

Initialize from one weighted (z, y) observation.

Parameters:

• x (tuple[Any, Any])

• weight (float)

• rng (RandomState)

Return type:

None

combine(suff_stat)

Merge expert sufficient statistics and buffered gate training data.

Parameters:

suff_stat (tuple[Any, ...])

Return type:

GatedMixtureAccumulator

value()

Return owned expert statistics, bounded gate rows, and retention status.

Return type:

GatedMixtureStatistics

from_value(x)

Restore expert statistics and gate training buffers.

Parameters:

x (tuple[Any, ...])

Return type:

GatedMixtureAccumulator

key_merge(stats_dict)

Delegate keyed merges to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

Delegate keyed replacements to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Return the encoder composed from expert accumulator encoders.

Return type:

GatedMixtureDataEncoder

class GatedMixtureAccumulatorFactory(component_factories, num_components, keys=None, *, max_buffer_rows=10_000, n_features=None)

Bases: StatisticAccumulatorFactory

Create accumulators for gated-mixture EM.

Parameters:

• component_factories (Sequence[Any])

• num_components (int)

• keys (str | None)

• max_buffer_rows (int)

• n_features (int | None)

make()

Create an empty gated-mixture accumulator.

Return type:

GatedMixtureAccumulator

class GatedMixtureEstimator(component_estimators, gate, gate_steps=200, gate_lr=0.1, gate_tol=1.0e-8, max_buffer_rows=10_000, name=None, keys=None)

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)

• gate_tol (float)

• max_buffer_rows (int)

• name (str | None)

• keys (str | None)

accumulator_factory()

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

Return type:

GatedMixtureAccumulatorFactory

estimate(nobs, suff_stat)

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

Parameters:

• nobs (float | None)

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

Return type:

GatedMixtureDistribution