diff --git a/dwave/plugins/torch/models/boltzmann_machine.py b/dwave/plugins/torch/models/boltzmann_machine.py index f3e3e10..170e1b4 100644 --- a/dwave/plugins/torch/models/boltzmann_machine.py +++ b/dwave/plugins/torch/models/boltzmann_machine.py @@ -47,9 +47,22 @@ __all__ = ["GraphRestrictedBoltzmannMachine"] + class GraphRestrictedBoltzmannMachine(torch.nn.Module): """Creates a graph-restricted Boltzmann machine. + The initialization strategy is grounded in `Hinton's practical guide for RBM training + `_, which recommends sampling weights + from a Gaussian distribution with mean 0 and small standard deviation. The quadratic weights + are initialized with graph-connectivity-dependent standard deviations so the energy remains + extensive on sparse graphs as well as dense graphs. In particular, For edge :math:`(u, v)`, + we set the standard deviation of its J value as :math:`ß / (\deg(u)\deg(v))^{1/4}`, where + :math:`ß=2.5` is half of a representative QPU inverse sampling-temperature scale. This + initializes the GRBM in a paramagnetic regime, consistent with the `Sherrington-Kirkpatrick + model `_. + The linear biases are initialized to zero to avoid introducing any initial preference for spin + configurations. + Args: nodes (Iterable[Hashable]): List of nodes. edges (Iterable[tuple[Hashable, Hashable]]): List of edges. @@ -60,6 +73,13 @@ class GraphRestrictedBoltzmannMachine(torch.nn.Module): quadratic (dict[tuple[Hashable, Hashable], float]): A dictionary mapping from edges of the model to its corresponding quadratic bias. """ + # QPU beta has been measured to be 5-8 (in inverse units of programmed J) + # Considering the higher temperature within this range, to sample from a beta=1 + # Boltzmann distribution, a prefactor of 5 has to multiply the initial Hamiltonian. + # To keep the energy scale of the initial Hamiltonian below the effective thermal + # energy, we multiply the Hamiltonian weights by an even smaller prefactor so + # that the prepared distribution is that of a paramagnet. + _INIT_INVERSE_TEMP = 2.5 def __init__( self, @@ -83,11 +103,26 @@ def __init__( self._idx_to_edge = {i: e for i, e in enumerate(self._edges)} self._edge_to_idx = {e: i for i, e in self._idx_to_edge.items()} - self._linear = torch.nn.Parameter(0.05 * (2 * torch.rand(self._n_nodes) - 1)) - self._quadratic = torch.nn.Parameter(5.0 * (2 * torch.rand(self._n_edges) - 1)) + edge_idx_i = torch.tensor([self._node_to_idx[i] for i, _ in self._edges], dtype=torch.long) + edge_idx_j = torch.tensor( + [self._node_to_idx[j] for _, j in self._edges], dtype=torch.long + ) + + degrees = torch.zeros(self._n_nodes) + for i, j in zip(edge_idx_i, edge_idx_j): + degrees[i] += 1 + degrees[j] += 1 + + if self._n_edges: + quadratic_std = self._INIT_INVERSE_TEMP / ( + degrees[edge_idx_i] * degrees[edge_idx_j] + )**0.25 + quadratic_init = torch.randn(self._n_edges) * quadratic_std + else: + quadratic_init = torch.empty(0) - edge_idx_i = torch.tensor([self._node_to_idx[i] for i, _ in self._edges]) - edge_idx_j = torch.tensor([self._node_to_idx[j] for _, j in self._edges]) + self._linear = torch.nn.Parameter(torch.zeros(self._n_nodes)) + self._quadratic = torch.nn.Parameter(quadratic_init) if (edge_idx_i == edge_idx_j).any(): loop_indices = edge_idx_i[(edge_idx_i == edge_idx_j).argwhere()].tolist() diff --git a/releasenotes/notes/gaussian-rbm-init-28fd4d295ef86d77.yaml b/releasenotes/notes/gaussian-rbm-init-28fd4d295ef86d77.yaml new file mode 100644 index 0000000..f0824c1 --- /dev/null +++ b/releasenotes/notes/gaussian-rbm-init-28fd4d295ef86d77.yaml @@ -0,0 +1,16 @@ +--- +upgrade: + - | + Initialize ``GraphRestrictedBoltzmannMachine`` weights using Gaussian \ + random variables with graph-connectivity-dependent standard deviations. \ + For an edge :math:`(u, v)`, the default standard deviation is \ + :math:`2.5 / (\deg(u)\deg(v))^{1/4}`. \ + The weight-initialization strategy is grounded in `Hinton's practical \ + guide for RBM training \ + `_, \ + which recommends sampling weights from a Gaussian distribution with mean 0 and standard \ + deviation 0.01 (for zero-one-valued RBMs). The connectivity scaling keeps \ + the energy functional extensive on sparse graphs, while the temperature factor initializes \ + the GRBM deep in a paramagnetic regime for QPU-backed sampling, \ + consistent with the `Sherrington-Kirkpatrick model \ + `_. diff --git a/tests/test_boltzmann_machine.py b/tests/test_boltzmann_machine.py index afeedde..b7d0fae 100644 --- a/tests/test_boltzmann_machine.py +++ b/tests/test_boltzmann_machine.py @@ -70,6 +70,36 @@ def test_constructor(self): self.assertAlmostEqual(bm.linear[2].item(), w1, 2) self.assertAlmostEqual(bm.quadratic[3].item(), w2, 2) + def test_default_quadratic_initialization_uses_connectivity(self): + nodes = list("abcd") + edges = [("a", "b"), ("a", "c"), ("a", "d"), ("b", "c")] + degrees = torch.tensor([3.0, 2.0, 2.0, 1.0]) + edge_idx_i = torch.tensor([0, 0, 0, 1]) + edge_idx_j = torch.tensor([1, 2, 3, 2]) + expected_std = 2.5 / (degrees[edge_idx_i] * degrees[edge_idx_j])**0.25 + + torch.manual_seed(1234) + expected_quadratic = torch.randn(len(edges)) * expected_std + + torch.manual_seed(1234) + bm = GRBM(nodes, edges) + + torch.testing.assert_close(bm.linear, torch.zeros(len(nodes))) + torch.testing.assert_close(bm.quadratic, expected_quadratic) + + def test_default_quadratic_initialization_edgeless(self): + bm = GRBM([0, 1, 2], []) + + torch.testing.assert_close(bm.linear, torch.zeros(3)) + self.assertEqual(0, bm.quadratic.numel()) + + def test_custom_quadratic_overrides_default_initialization(self): + bm = GRBM( + ["a", "b", "c"], [("a", "b"), ("b", "c")], quadratic={("b", "c"): 1.25} + ) + + self.assertAlmostEqual(1.25, bm.quadratic[1].item()) + def test_selfloop(self): # Create a triangle graph with an additional dangling vertex # a-SELF-LOOP diff --git a/tests/test_dvae_winci2020.py b/tests/test_dvae_winci2020.py index 38dfff7..e22cd39 100644 --- a/tests/test_dvae_winci2020.py +++ b/tests/test_dvae_winci2020.py @@ -78,12 +78,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # are the models themselves latent_dims_list = [1, 2] self.encoders = {i: Encoder(i) for i in latent_dims_list} - # self.decoders is independent of number of latent dims, but we also create a dict to separate - # them + # self.decoders is independent of number of latent dims, but we also create a dict to + # separate them self.decoders = {i: Decoder(latent_features, input_features) for i in latent_dims_list} - - # self.dvaes is a dict whose keys are the numbers of latent dims and the values are the models - # themselves + # self.dvaes is a dict whose keys are the numbers of latent dims and the values are the + # models themselves self.dvaes = {i: DVAE(self.encoders[i], self.decoders[i]) for i in latent_dims_list} @@ -248,19 +247,22 @@ def test_latent_to_discrete(self, n_samples, expected): @parameterized.expand([(i, j) for i in range(1, 3) for j in [0, 1, 5, 1000]]) def test_forward(self, n_latent_dims, n_samples): """Test the forward method.""" + torch.manual_seed(1234) # Set seed for reproducibility of latent_to_discrete sampling expected_latents = self.encoders[n_latent_dims](self.data) expected_discretes = self.dvaes[n_latent_dims].latent_to_discrete( expected_latents, n_samples ) expected_reconstructed_x = self.decoders[n_latent_dims](expected_discretes) + torch.manual_seed(1234) # Set seed again to ensure that the sampling in the forward method + # is the same as in the expected_discretes latents, discretes, reconstructed_x = self.dvaes[n_latent_dims].forward( x=self.data, n_samples=n_samples ) + torch.testing.assert_close(latents, expected_latents) + torch.testing.assert_close(discretes, expected_discretes) + torch.testing.assert_close(reconstructed_x, expected_reconstructed_x) - assert torch.equal(reconstructed_x, expected_reconstructed_x) - assert torch.equal(discretes, expected_discretes) - assert torch.equal(latents, expected_latents) if __name__ == "__main__":