Neural and LLM Models¶

Neural models enter Mixle in three different ways. Keeping them separate makes the API much easier to reason about.

This page covers the optional neural-model surface. The examples are useful for hybrid models that genuinely need a neural likelihood, but they carry more dependency, training-state, and reproducibility risk than the core mixle.stats families.

Install and document neural dependencies explicitly. A model that only works with Torch, a GPU device, or a registered module builder must record that requirement in its artifact manifest and in the deployment environment.

Surface	What it models	Use when
`mixle.models`	neural leaves inside generative/composable models	a Transformer, neural categorical, DPO model, or neural Gaussian is deliberately part of a larger distribution
`mixle.task`	local task models distilled from teachers	you want `text -> label` or `text -> fields` served cost-effectively
`mixle.reason`	uncertainty over LLM answers	you already have `generate(prompt) -> str` and need confidence, abstention, or claim reliability

This page covers the first surface. See Task Distillation and Uncertainty for the other two.

Transformer Leaves¶

TransformerLMEstimator is the direct way to put a causal Transformer into a mixle model. Its observations are (context, next_token) pairs.

from mixle.inference import optimize
from mixle.models import TransformerLMEstimator

pairs = [
    ([0, 1, 2, 3], 4),
    ([1, 2, 3, 4], 5),
]
est = TransformerLMEstimator(vocab=8000, d_model=256, n_layer=4, n_head=4, block=64)
model = optimize(pairs, est, max_its=20, out=None)

log_p_next = model.log_density((context_ids, next_id))

The fitted object is a conditional next-token distribution. It does not sample complete sequences through sampler; feed contexts and use the model’s next-token scoring or prediction behavior.

Shared Language-Model Experts¶

A practical reason to use a Transformer leaf inside Mixle is to build several language-model experts that share part of their parameterization. Each expert still models the same observation shape, (context, next_token), while the mixture learns which expert explains each example.

from mixle.models import CategoricalEmbedding, TransformerLMEstimator
from mixle.stats import MixtureEstimator

embedding = CategoricalEmbedding(num_categories=8000, dim=256, name="word")
experts = [
    TransformerLMEstimator(8000, d_model=256, n_layer=4, block=64, embedding=embedding)
    for _ in range(3)
]
est = MixtureEstimator(experts)

This is a coherent neural-composition example: the latent mixture separates token-stream regimes, while CategoricalEmbedding prevents every expert from learning an independent vocabulary table.

Neural Density Leaves¶

NeuralDensity and NeuralConditionalDensity adapt Torch modules that already expose density methods into Mixle leaves. They are useful when the neural part is a probability model, not just a feature extractor.

Surface	Models	Ready builders
`NeuralDensity`	unconditional `p(x)` with `module.log_density(x)`	`build_coupling_flow`, `build_maf`, `build_vae`
`NeuralConditionalDensity`	conditional `p(y \| x)` with `module.log_density(x, y)` and `module.sample_given(x)`	`build_mdn`, `build_conditional_flow`

build_mdn builds a mixture density network. It is useful when p(y | x) is multimodal or heteroscedastic, but its components are still diagonal Gaussians. build_conditional_flow builds the exact-density counterpart: a conditional coupling flow whose invertible transform of y is conditioned on x. Use the flow when the target dimensions have nonlinear within-y dependence and exact log-density matters.

For discrete vectors, build_autoregressive_categorical provides an exact unconditional neural density over {0, ..., C-1}^d. Its conditional sibling, build_conditional_autoregressive_categorical, provides exact p(y | x) for discrete target vectors. Use these builders when a categorical output has strong coordinate dependence and independent categorical leaves would erase the structure.

from mixle.models import NeuralConditionalDensity, build_conditional_flow

module = build_conditional_flow(x_dim=4, y_dim=2, hidden=64, layers=4)
leaf = NeuralConditionalDensity(module, m_steps=80, lr=1.0e-3)
estimator = leaf.estimator()

Observations for a conditional density leaf are (x, y) pairs. The M-step is responsibility-weighted negative log-likelihood, so the leaf can sit inside a mixture or another latent wrapper. Treat the fitted object as a Torch-backed artifact: record dependencies, training settings, validation data, and reload behavior.

Constructible Neural Density Families¶

Mixle provides direct distribution classes for common neural density families. Use them when the model tree should contain a neural density leaf without first building a Torch module and then wrapping it.

from mixle.models import Flow, VAE
from mixle.stats import MixtureDistribution

prior_shape = MixtureDistribution(
    [Flow(dim=4, hidden=64, layers=4), VAE(dim=4, latent=2)],
    [0.5, 0.5],
)
estimator = prior_shape.estimator()

Available constructible families are:

Flow: Exact continuous density via a RealNVP-style coupling flow.
MAF: Exact continuous density via a masked autoregressive flow.
VAE: Latent-variable density with an ELBO-style lower-bound score. Compare it against other bounded neural leaves carefully; it is not an exact likelihood like a flow.
DiscreteAR: Exact normalized autoregressive density over fixed-length discrete vectors.

These classes still use the same NeuralDensityEstimator route underneath, so EM responsibilities and sample weights reach the neural M-step.

Energy Models¶

EnergyModel wraps a Torch module whose energy(x) method returns a scalar compatibility score. Unlike normalizing flows, an energy model has an intractable normalizer, so Mixle trains it with noise-contrastive estimation and reports an approximately normalized log_density.

from mixle.inference import optimize
from mixle.models import EnergyModel, build_energy_net

leaf = EnergyModel(build_energy_net(dim=2, hidden=64, layers=3))
fitted = optimize(points, leaf.estimator(), max_its=10, out=None)
score = fitted.log_density([0.2, -0.5])

Use this when compatibility is a better inductive bias than an invertible flow or an autoregressive factorization. The sampler uses Langevin dynamics, and the normalization caveat should be recorded when mixing an energy leaf with exact density leaves.

Streaming Transformer¶

StreamingTransformer is the lower-level adapter when you already have a Torch module. Its accumulator owns a persistent optimizer and trains on streamed micro-batches, so it does not materialize the whole corpus as sufficient statistics.

from mixle.models import LM, StreamingTransformer

leaf = StreamingTransformer(
    LM(vocab=K, d_model=96, n_layer=3, n_head=4, block=B).module
)
est = leaf.estimator()

Use TransformerLMEstimator first. Reach for StreamingTransformer when you need to bring your own module, control streaming behavior, or share a live module across a larger fitting loop.

Streaming Transformer accumulation preserves sample weights. When the leaf sits below a mixture or HMM, EM responsibilities are passed into the streaming update instead of being discarded. That makes the streaming adapter consistent with the other neural leaves for latent-model M-steps.

Shared Embeddings¶

Mixtures of language-model experts can waste memory if every expert learns its own token embedding. CategoricalEmbedding declares one table and hands it to each expert.

from mixle.models import CategoricalEmbedding, TransformerLMEstimator
from mixle.stats import MixtureEstimator

emb = CategoricalEmbedding(num_categories=8000, dim=256, name="word")
experts = [
    TransformerLMEstimator(8000, d_model=256, n_layer=4, block=64, embedding=emb)
    for _ in range(3)
]
estimator = MixtureEstimator(experts)

All three experts share token vectors while learning separate expert dynamics. Run examples/shared_embedding_example.py to see the parameter accounting.

Direct LM Helper¶

LM is a small direct language-model helper:

from mixle.models import LM

lm = LM(vocab=5000, d_model=128, n_layer=2, n_head=4, block=64)
lm.fit(token_ids, epochs=10, batch_size=128, lr=3e-4)
out = lm.generate(prompt_ids, n=20, temperature=0.8, seed=0)
nll = lm.nll(token_ids)

Use LM directly when the language model is the artifact. Use TransformerLMEstimator when the language model is one leaf in a larger distribution.

Neural PPL Predictors¶

The PPL can place neural predictors inside distribution parameters:

from mixle.ppl import Categorical, Transformer

model = Categorical(
    logits=Transformer(out=vocab, d_model=64, n_layer=2, n_head=4)
).fit(
    next_tokens,
    given={"x": contexts},
    epochs=40,
    batch_size=128,
    lr=0.003,
)

The same contract supports MLP and convolutional predictors. The output is still a probabilistic model with scoring and diagnostics, not just a neural network.

Preference Optimization¶

DPOModel trains from preference triples (prompt, chosen, rejected):

from mixle.inference import estimate
from mixle.models import DPOModel

leaf = DPOModel(policy, reference_policy, beta=0.1, m_steps=20, lr=1e-4, device="cuda")
model = estimate(preference_triples, leaf.estimator())

Use it when the learning signal is comparative preference rather than a categorical label. DPO accumulation preserves per-pair weights, so responsibilities, streaming decay, or sample weights affect the DPO loss instead of being dropped.

Serialization and Artifacts¶

The neural surface now supports more durable round trips:

LM objects can serialize trained state and guard edge cases such as empty sequences.
StreamingTransformer and DPOModel expose to_dict/to_json style state.
NeuralGaussian, NeuralCategorical, NeuralDensity, NeuralConditionalDensity, EnergyModel, and related density leaves can round-trip through pickle and JSON-style state where the module builder is registered.
The constructible families VAE, Flow, MAF, and DiscreteAR are registered with the serialization layer.

Treat serialization as an artifact boundary, not a quality guarantee. Reload the model in a fresh process and rerun a small scoring or prediction check before relying on it in a service.

Device and Dependency Checks¶

Neural artifacts should not assume the training environment is still present at serve time. Before promoting a neural or LLM-centered model, verify:

the artifact loads on the intended CPU or GPU target;
missing optional dependencies produce a clear error or a documented fallback;
seeds and training settings are recorded;
a held-out prediction or log-density matches the pre-save result within the expected numerical tolerance;
the artifact does not claim unconditional sampling when the leaf is purely conditional.

These checks are especially important for streaming Transformer leaves, DPO models, and custom neural-density modules whose builder registration is part of the artifact contract.

How Neural Leaves Compose with Latents¶

A neural leaf can sit under a latent wrapper. For a mixture of Transformer experts, EM assigns responsibilities to examples, and each expert’s M-step is gradient fitting under those weights. For an HMM, the neural emission can train against expected state occupancy. The parent latent model sees an estimator; the child decides how to do its M-step.

Practical Checklist¶

Install mixle[torch].
Start with TransformerLMEstimator unless you already own the module.
Keep one observation as (context, target) for next-token leaves.
Use neural leaves when the neural likelihood is the natural model for that field or latent expert.
Use CategoricalEmbedding to tie embeddings across experts.
Prefer Flow, MAF, VAE, or DiscreteAR when a common neural density family is the target distribution.
Treat EnergyModel scores as approximately normalized; validate them before comparing directly with exact-density leaves.
Verify serialized neural artifacts after load, especially when optional Torch or device state is involved.
Fix seeds, record training settings, and evaluate on held-out data.
Use mixle.describe(model) after fitting to see what query capabilities the resulting object supports.

The older names StreamingTransformerLeaf, NeuralLeaf, SoftmaxNeuralLeaf, and DPOLeaf remain as compatibility aliases. Prefer StreamingTransformer, NeuralGaussian, NeuralCategorical, and DPOModel in new code.

Composition Grid¶

The composition grid makes the advertised behavior checkable. It lists which latent or structured contexts are validated and where a JSON or pickle round trip is part of that validation rather than assumed.

Leaf	Mixture component	Composite field	HMM emission	Serializes in each context
`NeuralGaussian` (conditional, `p(y\|x)`)	tested	tested	tested	tested in all three
`NeuralCategorical` (conditional, `p(y\|x)`)	tested	tested	tested	tested in all three
`NeuralDensity` / `Flow` / `VAE` / `MAF` / `DiscreteAR` (unconditional)	tested	tested	tested	tested in all three
`NeuralConditionalDensity` (`mixture_density`, conditional)	tested	tested	tested	tested in all three
`EnergyModel` (unconditional, approximately normalized)	tested	tested	not currently validated	tested (Mixture, Composite)

A purely conditional leaf such as NeuralCategorical, NeuralGaussian, or NeuralConditionalDensity has no marginal p(x) to draw from. A Mixture or HMM whose every component or state is conditional therefore cannot use the model’s own sampler(); NeuralCategorical.sample() raises by design. Fitting still works because optimize needs log_density and the accumulator contract, not unconditional sampling, so a conditionally emitting HMM trains normally on externally supplied (x, y) sequences.

EnergyModel as an HMM emission is outside the current validation grid. Do not document that route as supported until it has the same fitting, scoring, and serialization evidence as the validated neural leaves.

Examples and Validation¶

Validation for neural and LLM-centered models should cover:

shared embeddings across language-model mixture experts;
neural PPL predictors, streaming Transformer leaves, EWC, DPO, and the direct LM surface;
constructible neural-density families; and
neural and LM artifact round trips in a fresh process.