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 part of the model contract.

• 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 because it is only approximately normalized it can bias mixture weights against an exact leaf.

• 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)

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)

Return the approximate normalized log density for one observation.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return approximate normalized log densities for a batch of observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a Langevin sampler for the learned energy model.

Parameters:

seed (int | None)

Return type:

EnergyModelSampler

estimator(pseudo_count=None)

Return the NCE estimator used as the model’s M-step.

Parameters:

pseudo_count (float | None)

Return type:

EnergyModelEstimator

dist_to_encoder()

Return the encoder for vectorized energy-model scoring.

Return type:

EnergyModelEncoder

to_dict()

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

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild an EnergyModel from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

EnergyModel

class EnergyModelSampler(dist, seed=None)

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)

Draw approximate samples with unadjusted Langevin dynamics.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class EnergyModelEncoder

Bases: DataSequenceEncoder

Encode observations for vectorized energy-model scoring and fitting.

seq_encode(data)

Convert observations to a two-dimensional float array.

Parameters:

data (list)

Return type:

ndarray

class EnergyModelAccumulator

Bases: SequenceEncodableStatisticAccumulator

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

update(x, weight, estimate)

Add one weighted observation to the NCE accumulator.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Add an encoded batch and responsibility weights to the accumulator.

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Initialize from one observation using the ordinary update path.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

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

Parameters:

• enc (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(other)

Merge the value tuple from another energy-model accumulator.

Parameters:

other (Any)

Return type:

EnergyModelAccumulator

value()

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

Return type:

tuple

from_value(v)

Restore accumulator buffers from a value tuple.

Parameters:

v (tuple)

Return type:

EnergyModelAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

EnergyModelEncoder

class EnergyModelAccumulatorFactory

Bases: StatisticAccumulatorFactory

Factory for energy-model accumulators.

make()

Create a fresh accumulator.

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)

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()

Return an accumulator factory for weighted NCE batches.

Return type:

EnergyModelAccumulatorFactory

estimate(nobs, suff_stat)

Run the weighted NCE M-step and return the updated energy leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

EnergyModel

build_energy_net(dim, *, hidden=64, layers=3)

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:

• dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_convex_energy_net(dim, *, hidden=64, layers=3)

An input-convex MLP energy E(x): R^dim -> R, convex in x BY CONSTRUCTION – ready to wrap in an EnergyModel exactly like build_energy_net(). A convex energy gives Langevin sampling (EnergyModelSampler) a unimodal target with no spurious local minima to get stuck in, and gives any consumer of the fitted energy a certified-convex scalar-valued potential (e.g. a verified optimum for a downstream mixle.doe search over -E(x)). See _convex_energy_net_class() for the construction.

Parameters:

• dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_product_energy_net(experts)

Combine expert energy modules multiplicatively: one module with energy(x) = sum_k experts[k].energy(x).

A product of experts, p(x) ∝ prod_k p_k(x) – a conjunction (each expert a soft constraint, the product their intersection), as opposed to a mixture’s disjunction. This is the ENERGY-BASED, arbitrary-density complement to mixle.ops.product_of_experts(), which pools tractable families (Categorical, Gaussian) in closed form but deliberately raises on the general continuous case because the product normalizer is then intractable. That intractable normalizer is exactly what the energy stack already handles: wrap the result here in an EnergyModel to fit the shared log_norm by NCE and sample by Langevin, no new machinery.

Each expert must expose energy(x) -> (n,) and a dim attribute (e.g. any build_energy_net() / build_convex_energy_net() module, all sharing one input dim), and stays individually inspectable via the built module’s .experts / .expert_energies(x). Fit it in one line:

model = EnergyModel(build_product_energy_net([expert_a, expert_b]), m_steps=250)

Parameters:

experts (Any)

Return type:

Any