Ising Models

Ising-specific model, program, and training utilities. IsingEBM is the most common entry point — it constructs the factor graph from a list of nodes, edges, biases, and coupling weights.

hamon.models.IsingEBM

An EBM with the energy function,

\[\mathcal{E}(s) = -\beta \left( \sum_{i \in S_1} b_i s_i + \sum_{(i, j) \in S_2} J_{ij} s_i s_j \right)\]

where \(S_1\) and \(S_2\) are the sets of biases and weights that make up the model, respectively. \(b_i\) represents the bias associated with the spin \(s_i\) and \(J_{ij}\) is a weight that couples \(s_i\) and \(s_j\). \(\beta\) is the usual temperature parameter.

Attributes:

• nodes: the nodes that have an associated bias (i.e \(S_1\))
• biases: the bias associated with each node in nodes.
• edges: the edges that have an associated weight (i.e \(S_2\))
• weights: the weight associated with each pair of nodes in edges.
• beta: the scalar temperature parameter for the model.

__init__(nodes: list[hamon.pgm.AbstractNode], edges: list[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], biases: Array, weights: Array, beta: Array)

hamon.models.IsingSamplingProgram

A very thin wrapper on FactorSamplingProgram that specializes it to the case of an Ising Model.

__init__(ebm: IsingEBM, free_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], clamped_blocks: list[hamon.block_management.Block])

hamon.models.IsingTrainingSpec

Contains a complete specification of an Ising EBM that can be trained using sampling-based gradients.

Defines sampling programs and schedules that allow for collection of the positive and negative phase samples required for Monte Carlo estimation of the gradient of the KL-divergence between the model and a data distribution.

__init__(ebm: IsingEBM, data_blocks: list[hamon.block_management.Block], conditioning_blocks: list[hamon.block_management.Block], positive_sampling_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], negative_sampling_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], schedule_positive: SamplingSchedule, schedule_negative: SamplingSchedule)

hamon.models.hinton_init(key: Key[Array, ''], model: IsingEBM, blocks: list[hamon.block_management.Block[hamon.pgm.AbstractNode]], batch_shape: tuple[int, ...]) -> list[Bool[Array, 'batch_size block_size']]

Initialize the blocks according to the marginal bias.

Each binary unit \(i\) in a block is sampled independently as

\[\mathbb{P}(S_i = 1) = \sigma(\beta h_i) = \frac{1}{1 + e^{-\beta h_i}}\]

where \(h_i\) is the bias of unit i and \(\beta\) is the inverse-temperature scaling factor. See Hinton (2012) for a discussion of this initialization heuristic.

All blocks are sampled in parallel via vmap over a stacked index array, so the number of XLA ops is O(1) in the number of blocks rather than O(n_blocks). Blocks must all have the same size; empty blocks are rejected by BlockSpec upstream.

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the JAX PRNG key to use	required
model	IsingEBM	the Ising model to initialize for	required
blocks	list[Block[AbstractNode]]	the blocks that are to be initialized	required
batch_shape	tuple[int, ...]	the pre-pended batch dimension	required

Returns:

Type	Description
list[Bool[Array, 'batch_size block_size']]	the initialized blocks as a list of bool arrays, one per block

hamon.models.estimate_moments(key: Key[Array, ''], first_moment_nodes: list[hamon.pgm.AbstractNode], second_moment_edges: list[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], program: BlockSamplingProgram, schedule: SamplingSchedule, init_state: list[Array], clamped_data: list[Array])

Estimates the first and second moments of an Ising model Boltzmann distribution via sampling.

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the jax PRNG key	required
first_moment_nodes	list[AbstractNode]	the nodes that represent the variables we want to estimate the first moments of	required
second_moment_edges	list[tuple[AbstractNode, AbstractNode]]	the edges that connect the variables we want to estimate the second moments of	required
program	BlockSamplingProgram	the BlockSamplingProgram to be used for sampling	required
schedule	SamplingSchedule	the schedule to use for sampling	required
init_state	list[Array]	the variable values to use to initialize the sampling	required
clamped_data	list[Array]	the variable values to assign to the clamped nodes	required

Returns: the first and second moment data

hamon.models.estimate_kl_grad(key: Key[Array, ''], training_spec: IsingTrainingSpec, bias_nodes: list[hamon.pgm.AbstractNode], weight_edges: list[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], data: list[Array], conditioning_values: list[Array], init_state_positive: list[Array], init_state_negative: list[Array]) -> tuple

Estimate the KL-gradients of an Ising model with respect to its weights and biases.

Uses the standard two-term Monte Carlo estimator of the gradient of the KL-divergence between an Ising model and a data distribution.

The gradients are:

\[\Delta W = -\beta (\langle s_i s_j \rangle_{+} - \langle s_i s_j \rangle_{-})\]

\[\Delta b = -\beta (\langle s_i \rangle_{+} - \langle s_i \rangle_{-})\]

Here, \(\langle\cdot\rangle_{+}\) denotes an expectation under the positive phase (data-clamped Boltzmann distribution) and \(\langle\cdot\rangle_{-}\) under the negative phase (model distribution).

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the JAX PRNG key	required
training_spec	IsingTrainingSpec	the Ising EBM for which to estimate the gradients	required
bias_nodes	list[AbstractNode]	the nodes for which to estimate the bias gradients	required
weight_edges	list[tuple[AbstractNode, AbstractNode]]	the edges for which to estimate the weight gradients	required
data	list[Array]	The data values to use for the positive phase of the gradient estimate. Each array has shape [batch nodes]	required
conditioning_values	list[Array]	values to assign to the nodes that the model is conditioned on. Each array has shape [nodes]	required
init_state_positive	list[Array]	initial state for the positive sampling chain. Each array has shape [n_chains_pos batch nodes]	required
init_state_negative	list[Array]	initial state for the negative sampling chain. Each array has shape [n_chains_neg nodes]	required

Returns: the weight gradients and the bias gradients