THRML

Discrete energy-based models

Discrete EBM building blocks for spin and categorical variables, with square-tensor specializations and their matching Gibbs conditionals.

DiscreteEBMFactorclass
DiscreteEBMFactor(spin_node_groups: list[Block], categorical_node_groups: list[Block], weights: Array)

Implements batches of energy function terms of the form s_1 * ... * s_M * W[c_1, ..., c_N], where the s_i are spin variables and the c_i are categorical variables.

No variable should show up twice in any given interaction. If this happens, the result of sampling from a model that includes the bad factor might not agree with the Boltzmann distribution. For example, the interaction w * s_1 * s_1 * s_2 would violate this rule because s_1 shows up twice. To allow you to do something weird if you want to, this condition has not been enforced in the code.

Attributes:

  • spin_node_groups: the node groups involved in the batch of factors that represent spin-valued random variables.
  • categorical_node_groups: the node groups involved in the batch of factors that represent categorical-valued random variables.
  • weights: the batch of weight tensors W associated with the factors we are implementing. weights should have leading dimension b, where b is number of nodes in each element of spin_node_groups and categorical_node_groups. This tensor has shape [b, x_1, ..., x_N] where b is the number of nodes in each block and N is the length of categorical_node_groups.
  • is_spin: a map that indicates if a given node type represents a spin-valued random variable or not.
spin_node_groupsattribute
spin_node_groups: list[Block]
categorical_node_groupsattribute
categorical_node_groups: list[Block]
weightsattribute
weights: Array
is_spinattribute
is_spin: dict[Type[AbstractNode], bool]
node_groupsattribute
node_groups: list[Block]
to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Produce interaction groups that implement this factor.

In this case, we have to treat the spin and categorical node groups slightly differently.

energymethod
energy(global_state: list[Array], block_spec: BlockSpec)

Compute the energy associated with this factor.

In this case, that is the sum of terms like s_1 * ... * s_M * W[c_1, ..., c_N].

DiscreteEBMInteractionclass
DiscreteEBMInteraction(n_spin: int, weights: Array)

An interaction that shows up when sampling from discrete-variable EBMs.

Attributes:

  • n_spin: the number of spin states involved in the interaction.
  • weights: the weight tensor associated with this interaction.
n_spinattribute
n_spin: int
weightsattribute
weights: Array
SquareDiscreteEBMFactorclass
SquareDiscreteEBMFactor(spin_node_groups: list[Block], categorical_node_groups: list[Block], weights: Array)

A discrete factor with a square interaction weight tensor (shape [b, x, x, ..., x]).

If a discrete factor is square, the interaction groups corresponding to different choices of the head node blocks can be merged. This could yield smaller XLA programs and improved runtime performance via more efficient use of accelerators.

spin_node_groupsattribute
spin_node_groups: list[Block]
categorical_node_groupsattribute
categorical_node_groups: list[Block]
weightsattribute
weights: Array
is_spinattribute
is_spin: dict[Type[AbstractNode], bool]
node_groupsattribute
node_groups: list[Block]
energymethod
energy(global_state: list[Array], block_spec: BlockSpec)

Compute the energy associated with this factor.

In this case, that is the sum of terms like s_1 * ... * s_M * W[c_1, ..., c_N].

to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Call the parent class to_interaction_groups, and merge the results.

SpinEBMFactorclass
SpinEBMFactor(node_groups: list[Block], weights: Array)

A DiscreteEBMFactor that involves only spin variables.

spin_node_groupsattribute
spin_node_groups: list[Block]
categorical_node_groupsattribute
categorical_node_groups: list[Block]
weightsattribute
weights: Array
is_spinattribute
is_spin: dict[Type[AbstractNode], bool]
node_groupsattribute
node_groups: list[Block]
energymethod
energy(global_state: list[Array], block_spec: BlockSpec)

Compute the energy associated with this factor.

In this case, that is the sum of terms like s_1 * ... * s_M * W[c_1, ..., c_N].

to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Call the parent class to_interaction_groups, and merge the results.

CategoricalEBMFactorclass
CategoricalEBMFactor(node_groups: list[Block], weights: Array)

A DiscreteEBMFactor that involves only categorical variables.

spin_node_groupsattribute
spin_node_groups: list[Block]
categorical_node_groupsattribute
categorical_node_groups: list[Block]
weightsattribute
weights: Array
is_spinattribute
is_spin: dict[Type[AbstractNode], bool]
node_groupsattribute
node_groups: list[Block]
to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Produce interaction groups that implement this factor.

In this case, we have to treat the spin and categorical node groups slightly differently.

energymethod
energy(global_state: list[Array], block_spec: BlockSpec)

Compute the energy associated with this factor.

In this case, that is the sum of terms like s_1 * ... * s_M * W[c_1, ..., c_N].

SquareCategoricalEBMFactorclass
SquareCategoricalEBMFactor(node_groups: list[Block], weights: Array)

A DiscreteEBMFactor that involves only categorical variables that also has a square weight tensor.

spin_node_groupsattribute
spin_node_groups: list[Block]
categorical_node_groupsattribute
categorical_node_groups: list[Block]
weightsattribute
weights: Array
is_spinattribute
is_spin: dict[Type[AbstractNode], bool]
node_groupsattribute
node_groups: list[Block]
energymethod
energy(global_state: list[Array], block_spec: BlockSpec)

Compute the energy associated with this factor.

In this case, that is the sum of terms like s_1 * ... * s_M * W[c_1, ..., c_N].

to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Call the parent class to_interaction_groups, and merge the results.

SpinGibbsConditionalclass
SpinGibbsConditional()

A conditional update for spin-valued random variables that will perform a Gibbs sampling update given one or more DiscreteEBMInteractions.

This function can be extended to handle a broader class of interactions via inheritance. Specifically, a child class can override the compute_parameters method defined here, compute contributions to $\gamma$ from other types of interactions, and then call this method to take into account the contributions from DiscreteEBMInteractions.

initmethod
init()

Initialize the sampler state before sampling begins.

This is called before the first iteration of block sampling, after which the return of this method is superseded by the return from sample.

Returns: the initial sampler state to use for the first iteration of block sampling.

samplemethod
sample(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]], sampler_state: ~_SamplerState, output_sd: PyTree[jax.ShapeDtypeStruct]) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], 'State'], ~_SamplerState]

Sample from the distribution by first computing the parameters and then generating a sample based off of them.

sample_given_parametersmethod
sample_given_parameters(key: Key, parameters: PyTree, sampler_state: None, output_sd: PyTree[jax.ShapeDtypeStruct]) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], 'State'], None]

Sample from a spin-valued bernoulli distribution given the parameter $\gamma$. In THRML, 1 is represented by the boolean value True and -1 is represented by False.

compute_parametersmethod
compute_parameters(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']]], sampler_state: None, output_sd: PyTree[jax.ShapeDtypeStruct]) -> PyTree

Compute the parameter $\gamma$ of a spin-valued Bernoulli distribution given DiscreteEBMInteractions:

$$\gamma = \sum_i s_1^i \dots s_K^i \: W^i[x_1^i, \dots, x_M^i]$$

where the sum over $i$ is over all the DiscreteEBMInteractions seen by this function.

CategoricalGibbsConditionalclass
CategoricalGibbsConditional(n_categories: int)

A conditional update for categorical random variables that will perform a Gibbs sampling update given one or more DiscreteEBMInteractions.

This function can be extended to handle other interactions in the same way as [thrml.models.SpinGibbsConditional][].

Attributes:

  • n_categories: how many categories are involved in the softmax distribution this sampler will sample from.
n_categoriesattribute
n_categories: int
initmethod
init()

Initialize the sampler state before sampling begins.

This is called before the first iteration of block sampling, after which the return of this method is superseded by the return from sample.

Returns: the initial sampler state to use for the first iteration of block sampling.

samplemethod
sample(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]], sampler_state: ~_SamplerState, output_sd: PyTree[jax.ShapeDtypeStruct]) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], 'State'], ~_SamplerState]

Sample from the distribution by first computing the parameters and then generating a sample based off of them.

sample_given_parametersmethod
sample_given_parameters(key: Key, parameters: PyTree, sampler_state: None, output_sd: PyTree[jax.ShapeDtypeStruct]) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], 'State'], None]

Sample from a softmax distribution given the parameter vector $\theta$.

compute_parametersmethod
compute_parameters(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']]], sampler_state: None, output_sd: PyTree[jax.ShapeDtypeStruct]) -> PyTree

Compute the parameter $\theta$ of a softmax distribution given DiscreteEBMInteractions:

$$\theta = \sum_i s_1^i \dots s_K^i \: W^i[:, x_1^i, \dots, x_M^i]$$

where the sum over $i$ is over all the DiscreteEBMInteractions seen by this function.