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 an incubating part of the project. The examples are useful for experiments and 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.

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 cheaply

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, but it should be treated as an incubating Torch-backed surface with the usual training and reproducibility checks.

Constructible Neural Density Families

Version 0.6.2 adds 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.

In 0.6.2, 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,
)

Tests also cover 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. In v0.6.2, 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.

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, not just to demonstrate that different APIs can be forced together.

  • Use CategoricalEmbedding to tie embeddings across experts.

  • Prefer Flow, MAF, VAE, or DiscreteAR when a common neural density family is the intended 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.

Examples And Tests

  • examples/shared_embedding_example.py: shared embeddings across LM mixture experts.

  • mixle/tests/neural_ppl_test.py: neural PPL, streaming Transformer leaves, EWC, DPO, and the direct LM surface.

  • mixle/tests/neural_families_test.py: constructible neural-density families.

  • mixle/tests/neural_leaf_serialization_test.py and mixle/tests/lm_serialization_test.py: neural and LM artifact round trips.