mixle.models.softmax_leaf module

A neural classifier as a Mixle conditional-density leaf: p(y | x) = softmax(module(x)).

The discriminative sibling of NeuralGaussian. NeuralCategorical(module) wraps a Torch module that emits k logits as a mixle distribution over observations (x, y) with y an integer class index. It implements the full SequenceEncodableProbabilityDistribution contract, so it drops into MixtureDistribution / CompositeDistribution / HMM emissions like any leaf – and its EM M-step is a responsibility-weighted cross-entropy gradient step on the module (warm-started across EM iterations => generalized EM). The model’s seq_log_density IS -cross_entropy(module(x), y): the objective is the leaf’s log-density, never a user-supplied loss closure.

This is the leaf that the declarative Categorical(logits=Net(...)) PPL slot lowers to, and the component that makes a Mix([Categorical(logits=Net(...)), ...]) a mixture of neural classifiers fit by ordinary EM.

Requires torch. The leaf is conditional: predict(x) and sampler().sample_given(x) work; sample() raises because the model has no marginal p(x). This is the same conditional contract used by NeuralGaussian and RandomForestConditional.

class NeuralCategorical(module, m_steps=40, lr=0.01, name=None, batch_size=None, device='cpu')[source]

Bases: SequenceEncodableProbabilityDistribution

p(y | x) = softmax(module(x)) as a mixle leaf. Observation is the pair (x, y), y an int class.

batch_size (None = full batch) makes the M-step minibatch SGD over m_steps passes – needed to train a real conv net on a large image set; device (e.g. "mps"/"cuda") runs it on the GPU.

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • name (str | None)

  • batch_size (int | None)

  • device (str)

log_density(xy)[source]

Return log p(y | x) for one feature/class observation pair.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

Return per-row categorical conditional log probabilities for encoded pairs.

Parameters:

enc (Any)

Return type:

ndarray

predict(x)[source]

Return maximum-probability class predictions for one or more inputs.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler over labels given features.

Parameters:

seed (int | None)

Return type:

NeuralCategoricalSampler

estimator(pseudo_count=None)[source]

Return the generalized-EM estimator for weighted cross-entropy training.

Parameters:

pseudo_count (float | None)

Return type:

NeuralCategoricalEstimator

dist_to_encoder()[source]

Return the encoder for (x, class) observation pairs.

Return type:

NeuralCategoricalEncoder

to_dict()[source]

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Rebuild a NeuralCategorical from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralCategorical

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

Bases: DistributionSampler

Conditional sampler over class labels for NeuralCategorical.

Parameters:

  • dist (NeuralCategorical)

  • seed (int | None)

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

Raise because the leaf defines p(y | x) and has no marginal p(x).

Parameters:
Return type:

Any

sample_given(x)[source]

Draw one class label from p(y | x).

Parameters:

x (Any)

Return type:

int

class NeuralCategoricalEncoder[source]

Bases: DataSequenceEncoder

Encode feature/class pairs for neural-categorical scoring and fitting.

seq_encode(data)[source]

Convert (x, class) pairs into batched feature and integer-label arrays.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class NeuralCategoricalAccumulator[source]

Bases: SequenceEncodableStatisticAccumulator

Buffer weighted feature/class batches for the neural-categorical M-step.

update(xy, weight, estimate)[source]

Add one weighted feature/class pair to the accumulator.

Parameters:
Return type:

None

seq_update(enc, weights, estimate)[source]

Add an encoded batch and responsibility weights to the accumulator.

Parameters:
Return type:

None

initialize(xy, weight, rng)[source]

Initialize from one observation using the ordinary update path.

Parameters:
Return type:

None

seq_initialize(enc, weights, rng)[source]

Initialize from an encoded batch using the ordinary batch update path.

Parameters:
Return type:

None

combine(other)[source]

Merge the value tuple from another categorical accumulator.

Parameters:

other (Any)

Return type:

NeuralCategoricalAccumulator

value()[source]

Return contiguous (x, class, weights) arrays for the M-step.

Return type:

tuple

from_value(value)[source]

Restore accumulator buffers from a value tuple.

Parameters:

value (tuple)

Return type:

NeuralCategoricalAccumulator

acc_to_encoder()[source]

Return the encoder expected by this accumulator.

Return type:

NeuralCategoricalEncoder

class NeuralCategoricalAccumulatorFactory[source]

Bases: StatisticAccumulatorFactory

Factory for neural-categorical accumulators.

make()[source]

Create a fresh accumulator.

Return type:

NeuralCategoricalAccumulator

class NeuralCategoricalEstimator(module, m_steps=40, lr=0.01, name=None, batch_size=None, device='cpu', ewc=None)[source]

Bases: ParameterEstimator

EM estimator for a NeuralCategorical: the M-step is m_steps of responsibility-weighted cross-entropy gradient on the module (the module is warm-started across EM iterations => generalized EM).

The weighted CE is normalized by the responsibility mass sum(w) so the M-step is scale-invariant to the responsibility magnitude (the easy bug: an unnormalized weighted loss makes the step size track cluster size).

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • name (str | None)

  • batch_size (int | None)

  • device (str)

  • ewc (Any)

accumulator_factory()[source]

Return an accumulator factory for weighted classification batches.

Return type:

NeuralCategoricalAccumulatorFactory

estimate(nobs, suff_stat)[source]

Run the weighted cross-entropy M-step and return the updated leaf.

Parameters:
Return type:

NeuralCategorical

SoftmaxNeuralLeaf

alias of NeuralCategorical

SoftmaxNeuralLeafEstimator

alias of NeuralCategoricalEstimator