mixle.models.neural_density module¶
NeuralDensity – the adapter that turns ANY torch density module into a composable mixle distribution.
The point is not a specific architecture; it is the wrapper. NeuralGaussian already adapts a conditional net
(p(y | x)); this is its unconditional sibling: give it any torch module that exposes log_density(x) -> (n,)
(and, to draw samples, sample(n) -> (n, d)) and you get a full five-piece mixle Distribution – so a
flexible neural density drops into a MixtureDistribution component, an HMM emission, or a
CompositeDistribution field, and is fit jointly with classical families by EM. Its M-step is a
responsibility-weighted maximum-likelihood gradient ascent on the module, warm-started across EM iterations.
That is the thing no NN library offers: “a mixture of a normalizing flow and a Gamma”, “an HMM whose emissions
are flows”. Ready instances ship: build_coupling_flow() (a RealNVP-style flow, exact), build_maf() (a masked
autoregressive flow, exact, richer autoregressive dependence), build_vae() (a variational autoencoder, a
latent-variable density whose log_density is the ELBO lower bound) and build_autoregressive_categorical()
(an exact autoregressive density over discrete vectors) – structurally different families, continuous and
discrete, behind one adapter. Any other density (a normalized energy model, …) plugs in the same way.
- class NeuralDensity(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)[source]
Bases:
GradLeafWrap a torch density
module(module.log_density(x) -> (n,)) as a composable mixle distribution.A thin named subclass of
GradLeaf– the generic bridge owns the manufactured contract (buffer accumulator, array encoder, gradient M-step, sampler); this class owns only its name, its JSON payload, and its ready-module builders below.loss/optimizerhooks pass through (see the grad_leaf module docstring for the control story).- log_density(x)[source]
Return
log p(x)for one observation under the wrapped density module.
- seq_log_density(x)[source]
Return per-row log densities for encoded observations.
- sampler(seed=None)[source]
Return a sampler delegating to the wrapped module’s
samplemethod.- Parameters:
seed (int | None)
- Return type:
NeuralDensitySampler
- estimator(pseudo_count=None)[source]
Return the generalized-EM estimator for weighted neural-density training.
- Parameters:
pseudo_count (float | None)
- Return type:
NeuralDensityEstimator
- dist_to_encoder()[source]
Return the encoder for vectorized neural-density scoring and fitting.
- Return type:
NeuralDensityEncoder
- to_dict()[source]
Serialize hyperparameters and module bytes for registry-based round trips.
- class NeuralDensitySampler(dist, seed=None)[source]
Bases:
DistributionSamplerSampler for wrapped neural density modules exposing
sample(n).- Parameters:
dist (NeuralDensity)
seed (int | None)
- class NeuralDensityEncoder[source]
Bases:
DataSequenceEncoderEncode observations for vectorized neural-density scoring and fitting.
- class NeuralDensityAccumulator[source]
Bases:
SequenceEncodableStatisticAccumulatorBuffers the (responsibility-weighted) data for the M-step – the weights are the E-step’s soft counts.
- update(x, weight, estimate)[source]
Add one weighted observation to the accumulator.
- seq_update(enc, weights, estimate)[source]
Add an encoded batch and responsibility weights to the accumulator.
- initialize(x, weight, rng)[source]
Initialize from one observation using the ordinary update path.
- seq_initialize(enc, weights, rng)[source]
Initialize from an encoded batch using the ordinary batch update path.
- combine(other)[source]
Merge the value tuple from another neural-density accumulator.
- Parameters:
other (Any)
- Return type:
NeuralDensityAccumulator
- from_value(v)[source]
Restore accumulator buffers from a value tuple.
- Parameters:
v (tuple)
- Return type:
NeuralDensityAccumulator
- acc_to_encoder()[source]
Return the encoder expected by this accumulator.
- Return type:
NeuralDensityEncoder
- class NeuralDensityAccumulatorFactory[source]
Bases:
StatisticAccumulatorFactoryFactory for neural-density accumulators.
- make()[source]
Create a fresh accumulator.
- Return type:
NeuralDensityAccumulator
- class NeuralDensityEstimator(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)[source]
Bases:
ParameterEstimatorM-step: responsibility-weighted MLE –
max sum_i w_i log p(x_i)by gradient ascent on the module (warm).- accumulator_factory()[source]
Return an accumulator factory for weighted neural-density batches.
- Return type:
NeuralDensityAccumulatorFactory
- build_coupling_flow(dim, *, hidden=32, layers=4)[source]
A RealNVP coupling flow over
R^dimwith an exactlog_density(x)andsample(n)– ready to wrap.Alternating affine-coupling layers map data to a standard-normal base;
log_densityis the base log-prob plus the log-determinant of the (triangular) Jacobian. A minimal, correct instance of the density module aNeuralDensityadapts – swap in any other module with the same two methods.
- build_vae(dim, *, latent=2, hidden=32)[source]
Build a variational autoencoder over
R^dim.An amortized encoder
q(z | x)and a decoderp(x | z)(diagonal-Gaussian, learned observation scale) are trained by the ELBO with the reparameterization trick. This is a different family from a flow: structure is represented through a low-dimensional latent rather than an invertible map, while the sameNeuralDensityadapter can still use it because it exposes the same two methods.log_density(x)returns the ELBO, a lower bound onlog p(x), not the exact value. Compare VAE leaves with other bounded leaves whenever possible. Mixing a VAE with an exact-density leaf, such as a Gaussian or flow, compares a bound against an exact value and can under-weight the VAE.log_densityis deterministic: it evaluates the ELBO at the encoder meanz = mu(x)(norandnresample), so repeated scoring of the samexis bit-identical and an EM log-likelihood stays monotone. Training still uses the reparameterized sample (training=True) for an unbiased gradient.
- build_maf(dim, *, hidden=64, blocks=3)[source]
A masked autoregressive flow over
R^dim– an exact density that factorizesp(x)by the chain rule, eachp(x_i | x_{<i})an affine map with autoregressive (MADE-masked) mean and log-scale.Unlike the coupling flow it conditions every coordinate on all earlier ones (a richer autoregressive dependence), and unlike the VAE its
log_densityis exact. It can therefore be compared directly with a Gaussian, a flow, or another exact-density leaf. Sampling is the sequential inverse (one coordinate at a time). Another ready module forNeuralDensity; the adapter is unchanged.
- build_autoregressive_categorical(dim, n_categories, *, hidden=64)[source]
An autoregressive neural density over discrete vectors
x in {0..C-1}^dim– exact, normalizedp(x).The continuous flows/VAE above model
R^d; heterogeneous data is also categorical. This factorizesp(x) = prod_i p(x_i | x_{<i})with a MADE-masked network whose per-coordinate softmax is each conditional, so the density is exactly normalized (sums to 1 over the finite space) and can be compared directly with count/categorical families.log_densitysums the picked log-softmax logits;samplefills the vector one coordinate at a time. Another ready module forNeuralDensity; the adapter is unchanged.