probabilistic circuits in torch

Author: jjcmoon

src/klay/__init__.py

@@ -4,16 +4,20 @@
 from collections.abc import Sequence
 
 
-def to_torch_module(self: Circuit, semiring: str = "log"):
+def to_torch_module(self: Circuit, semiring: str = "log", probabilistic: bool = False):
     """
     Convert the circuit into a PyTorch module.
 
     :param semiring:
         The semiring in which the circuit should be evaluated. Supported options are ("log", "real", "mpe", "godel").
+    :param probabilistic:
+        If true, construct a probabilistic circuit instead of an arithmetic circuit.
+        This means the inputs to a sum node are multiplied by a probability, and
+        we can interpret sum nodes as latent Categorical variables.
     """
     from .backends import torch_backend
     indices = self._get_indices()
-    return torch_backend.KnowledgeModule(*indices, semiring=semiring)
+    return torch_backend.KnowledgeModule(*indices, semiring=semiring, probabilistic=probabilistic)
 
 
 def to_jax_function(self: Circuit, semiring: str = "log"):

src/klay/backends/torch_backend.py

@@ -1,6 +1,7 @@
 import math
 
 import torch
+from torch import nn
 
 CUTOFF = -math.log(2)
 
@@ -33,11 +34,11 @@ def unroll_csr(csr):
     return ixs.repeat_interleave(repeats=deltas)
 
 
-class KnowledgeModule(torch.nn.Module):
-    def __init__(self, pointers, csrs, semiring='real'):
+class KnowledgeModule(nn.Module):
+    def __init__(self, pointers, csrs, semiring='real', probabilistic=False):
         super(KnowledgeModule, self).__init__()
         layers = []
-        sum_layer, prod_layer, self.zero, self.one, self.negate = get_semiring(semiring)
+        sum_layer, prod_layer, self.zero, self.one, self.negate = get_semiring(semiring, probabilistic)
         for i, (ptrs, csr) in enumerate(zip(pointers, csrs)):
             ptrs = torch.as_tensor(ptrs)
             csr = torch.as_tensor(csr, dtype=torch.long)
@@ -46,85 +47,118 @@ def __init__(self, pointers, csrs, semiring='real'):
                 layers.append(prod_layer(ptrs, csr))
             else:
                 layers.append(sum_layer(ptrs, csr))
-        self.layers = torch.nn.Sequential(*layers)
+        self.layers = nn.Sequential(*layers)
 
     def forward(self, weights, neg_weights=None, eps=0):
         if neg_weights is None:
             neg_weights = self.negate(weights, eps)
         x = encode_input(weights, neg_weights, self.zero, self.one)
         return self.layers(x)
 
-    def sparsity(self, nb_vars):
+    def sparsity(self, nb_vars: int) -> float:
         sparse_params = sum(len(l.csr) for l in self.layers)
         layer_widths = [nb_vars] + [l.out_shape[0] for l in self.layers]
         dense_params  = sum(layer_widths[i] * layer_widths[i+1] for i in range(len(layer_widths) - 1))
         return sparse_params / dense_params
 
 
-class KnowledgeLayer(torch.nn.Module):
+class KnowledgeLayer(nn.Module):
     def __init__(self, ptrs, csr):
         super().__init__()
         self.register_buffer('ptrs', ptrs)
         self.register_buffer('csr', csr)
         self.out_shape = (self.csr[-1].item() + 1,)
 
+    def _scatter_reduce(self, src: torch.Tensor, reduce: str):
+        output = torch.empty(self.out_shape, dtype=src.dtype, device=src.device)
+        output = torch.scatter_reduce(output, 0, index=self.csr, src=src, reduce=reduce, include_self=False)
+        return output
+
+    def _logsumexp_scatter_reduce(self, x: torch.Tensor, epsilon: float):
+        with torch.no_grad():
+            max_output = torch.empty(self.out_shape, dtype=x.dtype, device=x.device)
+            max_output = torch.scatter_reduce(max_output, 0, index=self.csr, src=x, reduce="amax", include_self=False)
+        x = x - max_output[self.csr]
+        x.nan_to_num_(nan=0., posinf=float('inf'), neginf=float('-inf'))
+        x = torch.exp(x)
+
+        output = torch.full(self.out_shape, epsilon, dtype=x.dtype, device=x.device)
+        output = torch.scatter_add(output, 0, index=self.csr, src=x)
+        output = torch.log(output) + max_output
+        return output
+
+
+
+class ProbabilisticKnowledgeLayer(KnowledgeLayer):
+    def __init__(self, ptrs, csr):
+        super().__init__(ptrs, csr)
+        self.weights = nn.Parameter(torch.randn_like(ptrs))
+
 
 class SumLayer(KnowledgeLayer):
     def forward(self, x):
-        output = torch.zeros(self.out_shape, dtype=x.dtype, device=x.device)
-        output = torch.scatter_add(output, 0, index=self.csr, src=x[self.ptrs])
-        return output
+        return self._scatter_reduce(x[self.ptrs], "sum")
 
 
 class ProdLayer(KnowledgeLayer):
     def forward(self, x):
-        output = torch.empty(self.out_shape, dtype=x.dtype, device=x.device)
-        output = torch.scatter_reduce(output, 0, index=self.csr, src=x[self.ptrs], reduce="prod", include_self=False)
-        return output
+        return self._scatter_reduce(x[self.ptrs], "prod")
 
 
 class MinLayer(KnowledgeLayer):
     def forward(self, x):
-        output = torch.empty(self.out_shape, dtype=x.dtype, device=x.device)
-        output = torch.scatter_reduce(output, 0, index=self.csr, src=x[self.ptrs], reduce="amin", include_self=False)
-        return output
+        return self._scatter_reduce(x[self.ptrs], "amin")
 
 
 class MaxLayer(KnowledgeLayer):
     def forward(self, x):
-        output = torch.empty(self.out_shape, dtype=x.dtype, device=x.device)
-        output = torch.scatter_reduce(output, 0, index=self.csr, src=x[self.ptrs], reduce="amax", include_self=False)
-        return output
+        return self._scatter_reduce(x[self.ptrs], "amax")
 
 
 class LogSumLayer(KnowledgeLayer):
     def forward(self, x, epsilon=10e-16):
-        x = x[self.ptrs]
-        with torch.no_grad():
-            max_output = torch.empty(self.out_shape, dtype=x.dtype, device=x.device)
-            max_output = torch.scatter_reduce(max_output, 0, index=self.csr, src=x, reduce="amax", include_self=False)
-        x = x - max_output[self.csr]
-        x.nan_to_num_(nan=0., posinf=float('inf'), neginf=float('-inf'))
-        x = torch.exp(x)
+        return self._logsumexp_scatter_reduce(x[self.ptrs], epsilon)
 
-        output = torch.full(self.out_shape, epsilon, dtype=x.dtype, device=x.device)
-        output = torch.scatter_add(output, 0, index=self.csr, src=x)
-        output = torch.log(output) + max_output
-        return output
+
+class ProbabilisticSumLayer(ProbabilisticKnowledgeLayer):
+    def forward(self, x):
+        x = self.get_edge_weights() * x[self.ptrs]
+        return self._scatter_reduce(x, "sum")
+
+    def get_edge_weights(self):
+        exp_weights = torch.exp(self.weights)
+        norm = self._scatter_reduce(exp_weights, "sum")
+        return exp_weights / norm
+
+
+class ProbabilisticLogSumLayer(ProbabilisticKnowledgeLayer):
+    def forward(self, x, epsilon=10e-16):
+        x = self.get_edge_weights(epsilon) + x[self.ptrs]
+        return self._logsumexp_scatter_reduce(x, epsilon)
+
+    def get_edge_weights(self, epsilon):
+        norm = self._logsumexp_scatter_reduce(self.weights, epsilon)
+        return self.weights - norm
 
 
-def get_semiring(name: str):
+def get_semiring(name: str, probabilistic: bool):
     """
     For a given semiring, returns the sum and product layer,
     the zero and one elements, and a negation function.
     """
-    if name == "real":
-        return SumLayer, ProdLayer, 0, 1, negate_real
-    elif name == "log":
-        return LogSumLayer, SumLayer, float('-inf'), 0, log1mexp
-    elif name == "mpe":
-        return MaxLayer, ProdLayer, 0, 1, negate_real
-    elif name == "godel":
-        return MaxLayer, MinLayer, 0, 1, negate_real
+    if probabilistic:
+        if name == "real":
+            return ProbabilisticSumLayer, ProdLayer, 0, 1, negate_real
+        if name == "log":
+            return ProbabilisticLogSumLayer, SumLayer, float('-inf'), 0, log1mexp
+        raise ValueError(f"Unknown probabilistic semiring {name}")
     else:
+        if name == "real":
+            return SumLayer, ProdLayer, 0, 1, negate_real
+        elif name == "log":
+            return LogSumLayer, SumLayer, float('-inf'), 0, log1mexp
+        elif name == "mpe":
+            return MaxLayer, ProdLayer, 0, 1, negate_real
+        elif name == "godel":
+            return MaxLayer, MinLayer, 0, 1, negate_real
         raise ValueError(f"Unknown semiring {name}")