Natural Gradient Descent

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
   "h5py>=3.14.0",
   "numpy>=2.0.0",
   "matplotlib>=3.8.0",
-  "torch>=2.10.0",
+  "torch>=2.8.0",
   "tqdm>=4.65.0",
 ]
 

rbms/bernoulli_bernoulli/classes.py

@@ -53,6 +53,7 @@ def __init__(
         self.hbias = hbias.to(device=self.device, dtype=self.dtype)
         self.name = "BBRBM"
         self.flags = []
+        self.start = None
 
     def __add__(self, other):
         return BBRBM(
@@ -100,9 +101,9 @@ def compute_energy_hiddens(self, h: Tensor) -> Tensor:
             weight_matrix=self.weight_matrix,
         )
 
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         return _compute_energy_visibles(
-            v=v,
+            v=chains["visible"],
             vbias=self.vbias,
             hbias=self.hbias,
             weight_matrix=self.weight_matrix,

rbms/bernoulli_gaussian/classes.py

@@ -112,9 +112,9 @@ def compute_energy_hiddens(self, h: Tensor) -> Tensor:
             h=h, vbias=self.vbias, hbias=self.hbias, weight_matrix=self.weight_matrix
         )
 
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         return _compute_energy_visibles(
-            v=v,
+            v=chains["visible"],
             vbias=self.vbias,
             hbias=self.hbias,
             weight_matrix=self.weight_matrix,

rbms/bm/potts.py

@@ -74,7 +74,7 @@ def sample_visibles(
         chains["visible"] = one_hot_gen.argmax(-1).to(self.dtype)
         return chains
 
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         """Returns the marginalized energy of the model computed on the visible configurations
 
         Args:
@@ -83,6 +83,7 @@ def compute_energy_visibles(self, v: Tensor) -> Tensor:
         Returns:
             Tensor: The computed energy.
         """
+        v = chains["visible"]
         L, q = self.bias.shape
         batch_size = v.shape[0]
         x_flat = (

rbms/classes.py

@@ -56,7 +56,7 @@ def sample_visibles(
         ...
 
     @abstractmethod
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         """Returns the marginalized energy of the model computed on the visible configurations
 
         Args:
@@ -371,3 +371,6 @@ def get_metrics_display(
 
     @abstractmethod
     def get_metrics_save(self) -> dict[str, np.ndarray] | None: ...
+
+    @abstractmethod
+    def get_curr_conf(self) -> dict[str, Tensor]: ...

rbms/dataset/__init__.py

@@ -19,6 +19,7 @@ def load_dataset(
     shuffle: bool = False,
     device: torch.device | str = "cpu",
     dtype: torch.dtype = torch.float32,
+    verbose: bool = True,
 ) -> tuple[RBMDataset, RBMDataset | None]:
     return_datasets = []
     for dset_name in [dataset_name, test_dataset_name]:
@@ -29,7 +30,8 @@ def load_dataset(
 
         if dset_name is not None:
             dset_name = Path(dset_name)
-            print(f"Reading dataset from {str(dset_name)}...")
+            if verbose:
+                print(f"Reading dataset from {str(dset_name)}...")
             match dset_name.suffix:
                 case ".h5":
                     data, labels, variable_type, weights = load_HDF5(

rbms/dataset/dataset_class.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 import gzip
 import textwrap
 from typing import Union
@@ -128,9 +129,9 @@ def get_gzip_entropy(self, mean_size: int = 50, num_samples: int = 100):
             )
         return np.mean(en)
 
-    def match_model_variable_type(self, visible_type: str):
+    def match_model_variable_type(self, visible_type: str, verbose: bool = True):
         self.data = convert_data[self.variable_type][visible_type](self.data)
-        if self.variable_type != visible_type:
+        if self.variable_type != visible_type and verbose:
             print(f"Converting from '{self.variable_type}' to '{visible_type}'")
             print(self.data)
         self.variable_type = visible_type

rbms/ising_gaussian/classes.py

@@ -7,7 +7,7 @@
 from torch import Tensor
 
 from rbms.classes import RBM
-from rbms.custom_fn import check_keys_dict, log2cosh
+from rbms.custom_fn import log2cosh
 from rbms.ising_gaussian.implement import (
     _compute_energy,
     _compute_energy_hiddens,
@@ -105,9 +105,9 @@ def compute_energy_hiddens(self, h: Tensor) -> Tensor:
             h=h, vbias=self.vbias, hbias=self.hbias, weight_matrix=self.weight_matrix
         )
 
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         return _compute_energy_visibles(
-            v=v,
+            v=chains["visible"],
             vbias=self.vbias,
             hbias=self.hbias,
             weight_matrix=self.weight_matrix,

rbms/ising_gaussian/implement.py

@@ -9,8 +9,8 @@
 def _sample_hiddens(
     v: Tensor, weight_matrix: Tensor, hbias: Tensor, beta: float = 1.0
 ) -> Tuple[Tensor, Tensor]:
-    mh = hbias + (v @ weight_matrix)
-    h = (
+    mh = beta*(hbias + (v @ weight_matrix)) / weight_matrix.shape[0]
+    h = ( 
         torch.randn_like(mh) / torch.sqrt(torch.ones_like(mh) * weight_matrix.shape[0])
         + mh
     )
@@ -49,7 +49,7 @@ def _compute_energy_visibles(
     field = v @ vbias
     t = hbias + (v @ weight_matrix)
     quad_term = 0.5 * (t * t).sum(1) / float(weight_matrix.shape[0])
-    return -field - quad_term + const
+    return -field - quad_term - const
 
 
 def _compute_energy_hiddens(

rbms/ising_ising/classes.py

@@ -10,6 +10,7 @@
     _compute_energy,
     _compute_energy_hiddens,
     _compute_energy_visibles,
+    _compute_energy_visibles_gradient,
     _compute_gradient,
     _init_chains,
     _init_parameters,
@@ -100,9 +101,9 @@ def compute_energy_hiddens(self, h: Tensor) -> Tensor:
             weight_matrix=self.weight_matrix,
         )
 
-    def compute_energy_visibles(self, v: Tensor) -> Tensor:
+    def compute_energy_visibles(self, chains: dict[str, Tensor]) -> Tensor:
         return _compute_energy_visibles(
-            v=v,
+            v=chains["visible"],
             vbias=self.vbias,
             hbias=self.hbias,
             weight_matrix=self.weight_matrix,
@@ -122,6 +123,14 @@ def compute_gradient(self, data, chains, centered=True):
             centered=centered,
         )
 
+    def compute_energy_visible_gradient(self, v: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        return _compute_energy_visibles_gradient(
+            v=v,
+            vbias=self.vbias,
+            hbias=self.hbias,
+            weight_matrix=self.weight_matrix,
+        )
+
     def independent_model(self):
         return IIRBM(
             weight_matrix=torch.zeros_like(self.weight_matrix),
@@ -149,7 +158,7 @@ def init_chains(self, num_samples, weights=None, start_v=None):
         )
 
     @staticmethod
-    def init_parameters(num_hiddens, dataset, device, dtype, var_init=0.0001):
+    def init_parameters(num_hiddens, dataset, device, dtype, var_init=0.001):
         data = dataset.data
         # Convert to torch Tensor if necessary
         if isinstance(data, np.ndarray):

rbms/ising_ising/implement.py

@@ -159,3 +159,16 @@ def _init_parameters(
     vbias = torch.atanh(frequencies).to(device=device, dtype=dtype)
     hbias = torch.zeros(num_hiddens, device=device, dtype=dtype)
     return vbias, hbias, weight_matrix
+
+
+def _compute_energy_visibles_gradient(
+    v: Tensor, vbias: Tensor, hbias: Tensor, weight_matrix: Tensor
+) -> tuple[Tensor, Tensor, Tensor]:
+    local_field = hbias + (v @ weight_matrix)
+    tanh_term = torch.tanh(local_field)
+
+    grad_vbias = -v
+    grad_hbias = -tanh_term
+    grad_weight_matrix = -v.T @ tanh_term
+    # grad_weight_matrix = None
+    return grad_vbias, grad_hbias, grad_weight_matrix