mixle.models.energy module

EnergyModel – an energy-based density p(x) ∝ exp(-E(x)) as a composable mixle leaf.

The one neural density whose normalizer is intractable: p(x) = exp(-E(x)) / Z with Z = ∫ exp(-E(x)) dx unavailable in closed form. So unlike the flows (exact) it is trained and scored approximately, and this is stated plainly – it is the energy-model analogue of the VAE’s ELBO caveat.

  • Training is Noise-Contrastive Estimation (Gutmann & Hyvärinen 2010), not maximum likelihood: the model learns to tell data from samples of a known noise distribution, and in doing so learns a scalar log-normalizer c alongside the energy net. NCE is consistent – as data grow, c -> log Z and -E(x) + c -> log p(x) – so log_density(x) = -E(x) + c is an approximately normalized log-density, usable directly (no per-evaluation partition estimate). It composes in a mixture, but being only approximately normalized it can bias mixture weights against an exact leaf (same honesty caveat as the VAE).

  • Sampling is unnormalized-density MCMC: a few steps of Langevin dynamics x <- x - s ∇E(x) + sqrt(2s) ε.

Its value over the flows is the inductive bias: an energy net imposes no ordering and no invertibility – it scores compatibility, so it captures undirected/symmetric structure a coupling or autoregressive flow parameterizes awkwardly. build_energy_net() is a ready MLP energy to wrap.

class EnergyModel(module, *, m_steps=200, lr=5e-3, noise_ratio=1, langevin_steps=40, langevin_step=0.05, device='cpu', name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

log p(x) ≈ -E(x) + c for an energy module (module.energy(x) -> (n,) and a learned scalar log_norm).

Approximately normalized (trained by NCE); log_density returns -E(x) + c. Composes like any leaf.

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • noise_ratio (int)

  • langevin_steps (int)

  • langevin_step (float)

  • device (str)

  • name (str | None)

log_density(x)[source]

Return the log-density or log-mass at a single observation.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

EnergyModelSampler

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

EnergyModelEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

EnergyModelEncoder

to_dict()[source]

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

EnergyModel

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

Bases: DistributionSampler

Langevin dynamics on the (unnormalized) energy: x <- x - s ∇E(x) + sqrt(2 s) ε.

Parameters:

  • dist (EnergyModel)

  • seed (int | None)

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

Draw observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=True (the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

Parameters:
Return type:

Any

class EnergyModelEncoder[source]

Bases: DataSequenceEncoder

seq_encode(data)[source]

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

ndarray

class EnergyModelAccumulator[source]

Bases: SequenceEncodableStatisticAccumulator

Buffers responsibility-weighted data for the NCE M-step (weights = the E-step soft counts).

update(x, weight, estimate)[source]
Parameters:
Return type:

None

seq_update(enc, weights, estimate)[source]
Parameters:
Return type:

None

initialize(x, weight, rng)[source]
Parameters:
Return type:

None

seq_initialize(enc, weights, rng)[source]
Parameters:
Return type:

None

combine(other)[source]
Parameters:

other (Any)

Return type:

EnergyModelAccumulator

value()[source]
Return type:

tuple

from_value(v)[source]
Parameters:

v (tuple)

Return type:

EnergyModelAccumulator

acc_to_encoder()[source]
Return type:

EnergyModelEncoder

class EnergyModelAccumulatorFactory[source]

Bases: StatisticAccumulatorFactory

make()[source]
Return type:

EnergyModelAccumulator

class EnergyModelEstimator(module, *, m_steps=200, lr=5e-3, noise_ratio=1, langevin_steps=40, langevin_step=0.05, device='cpu', name=None)[source]

Bases: ParameterEstimator

M-step: Noise-Contrastive Estimation against a Gaussian noise fit to the (weighted) data.

Learns the energy net and the scalar log-normalizer log_norm by logistic discrimination of data from noise – so the resulting -E(x) + log_norm is a consistent, approximately-normalized log-density.

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • noise_ratio (int)

  • langevin_steps (int)

  • langevin_step (float)

  • device (str)

  • name (str | None)

accumulator_factory()[source]
Return type:

EnergyModelAccumulatorFactory

estimate(nobs, suff_stat)[source]
Parameters:
Return type:

EnergyModel

build_energy_net(dim, *, hidden=64, layers=3)[source]

An MLP energy E(x): R^dim -> R (plus a learned scalar log_norm) – ready to wrap in an EnergyModel.

Lower energy = higher (unnormalized) density. log_norm is the NCE-learned normalizer, so the paired EnergyModel scores -E(x) + log_norm. Swap in any module exposing energy(x) -> (n,), a log_norm parameter and a dim attribute.

Parameters:
Return type:

Any