gaussian grbm initialization

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
+    <https://www.cs.toronto.edu/~hinton/absps/guideTR.pdf>`_, 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 <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.35.1792>`_.
+    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()

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 \
+    <https://www.cs.toronto.edu/~hinton/absps/guideTR.pdf>`_, \
+    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 \
+    <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.35.1792>`_.

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

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__":