Code for Convolutional Logic Gate Networks

[OmiKakadiya]

When you will be made the code base available for Convolutional Logic Gate Networks because many months passed since paper has been published
and we are waiting for it

[SuperrMurlocc]

I know that my code is probably no match for what Felix Petersen and his group did, but since I needed Convolutional Logic Gate Networks for my own research, I share the code that I used for it:

The first code block is for regular Logic layer, but I post it with changes from the new paper (residual).
The second code block is LogicTree and the third one - convolutional layer.

It is surely unoptimized! But perhaps someone will need it for their research, just as I did.

logic_gates = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]


def apply_logic_gate(a: Tensor, b: Tensor, logic_gate: int):
    return {
        0:  torch.zeros_like(a),
        1:  a * b,
        2:  a - a * b,
        3:  a,
        4:  b - a * b,
        5:  b,
        6:  a + b - 2 * a * b,
        7:  a + b - a * b,
        8:  1 - (a + b - a * b),
        9:  1 - (a + b - 2 * a * b),
        10:  1 - b,
        11:  1 - b + a * b,
        12:  1 - a,
        13:  1 - a + a * b,
        14:  1 - a * b,
        15:  torch.ones_like(a),
    }[logic_gate]


class Logic(pl.LightningModule):
    def __init__(self,
                 in_dim: int,
                 out_dim: int,
                 initialization_type: InitializationType = 'residual',

                 *args,
                 **kwargs
                 ):
        super().__init__(*args, **kwargs)

        self.in_dim = in_dim
        self.out_dim = out_dim
        self.initialization_type = initialization_type

        a, b = self.get_connections()
        self.register_buffer('a', a)
        self.register_buffer('b', b)

        weights = torch.randn(out_dim, len(logic_gates), device=self.device)
        if self.initialization_type == 'residual':
            weights[:, :] = 0
            weights[:, 3] = 5

        self.weights = torch.nn.parameter.Parameter(weights)

    def forward(self, x: Tensor):
        a, b = x[:, self.a, ...], x[:, self.b, ...]

        if self.training:
            normalized_weights = torch.nn.functional.softmax(self.weights, dim=-1).to(x.dtype).to(self.device)

            r = torch.zeros_like(a).to(x.dtype).to(self.device)

            for logic_gate in logic_gates:
                if len(a.shape) > 2:
                    nw = einops.repeat(normalized_weights[..., logic_gate], 'weights -> weights depth', depth=a.shape[-1])
                else:
                    nw = normalized_weights[..., logic_gate]
                r = r + nw * apply_logic_gate(a, b, logic_gate)
            return r
        else:
            one_hot_weights = torch.nn.functional.one_hot(self.weights.argmax(-1), len(logic_gates)).to(torch.float32).to(self.device)

            with torch.no_grad():
                r = torch.zeros_like(a).to(x.dtype).to(self.device)

                for logic_gate in logic_gates:
                    if len(a.shape) > 2:
                        ohw = einops.repeat(one_hot_weights[..., logic_gate], 'weights -> weights depth', depth=a.shape[-1])
                    else:
                        ohw = one_hot_weights[..., logic_gate]

                    r = r + ohw * apply_logic_gate(a, b, logic_gate)
                return r

    def get_connections(self):
        connections = torch.randperm(2 * self.out_dim) % self.in_dim
        connections = torch.randperm(self.in_dim)[connections]
        connections = connections.reshape(2, self.out_dim)
        a, b = connections[0], connections[1]
        a, b = a.to(torch.int64), b.to(torch.int64)
        a, b = a.to(self.device), b.to(self.device)
        return a, b

class LogicTree(pl.LightningModule):
    def __init__(self,
                 in_dim: int,
                 depth: int = 3,
                 initialization_type: InitializationType = 'residual',

                 *args,
                 **kwargs
                 ):
        super().__init__(*args, **kwargs)

        self.tree = nn.Sequential(
            Logic(in_dim, int(2 ** (depth - 1)), initialization_type=initialization_type),
            *[
                Logic(int(2 ** (depth - 1 - i)), int(2 ** (depth - 1 - i - 1)), initialization_type=initialization_type) for i in range(0, depth - 1, 1)
            ]
        )

    def forward(self, x: Tensor):
        return self.tree(x)

class Conv(pl.LightningModule):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 depth: int = 3,
                 kernel_size: int = 3,
                 stride: int = 1,
                 padding: int = 1,
                 initialization_type: InitializationType = 'residual',

                 *args,
                 ** kwargs
                 ):
        super().__init__(*args, **kwargs)

        self.kernel_size = kernel_size
        self.in_channels = in_channels
        self.out_channels = out_channels

        self.filters = nn.ModuleList([LogicTree(in_dim=kernel_size ** 2 * in_channels, depth=depth, initialization_type=initialization_type) for _ in range(out_channels)])

        self.stride = stride
        self.padding = padding

    def forward(self, x: Tensor) -> Tensor:
        batch_size, _, height, width = x.shape
        x = F.pad(x, (self.padding, self.padding, self.padding, self.padding), mode='constant', value=0)

        out_height = (height + 2 * self.padding - self.kernel_size) // self.stride + 1
        out_width = (width + 2 * self.padding - self.kernel_size) // self.stride + 1

        patches = F.unfold(x, kernel_size=self.kernel_size, stride=self.stride)

        outputs = []
        for filter in self.filters:
            out = filter(patches)  # Shape: (batch_size, 1, out_height * out_width)
            out = einops.rearrange(out, 'b 1 (h w) -> b (h w)', h=out_height, w=out_width)
            outputs.append(out)

        output_tensor = torch.stack(outputs, dim=1)  # Shape: (batch_size, out_channels, out_height * out_width)
        output_tensor = einops.rearrange(output_tensor, 'b c (h w) -> b c h w', h=out_height, w=out_width)

        return output_tensor

[SampadChowdhury]

Hey @SuperrMurlocc , thanks for the headstart by sharing the code.
I will go through the code today, but I was just wondering were you able to understand/code the part mentioned in Page 6 (Section 3.4), where the author talks about "splitting the model in k/8 groups and each tree observing only 2 input channels"
As per my understanding, this is something quite different from normal convolutions, where one kernel convolves each input channel and ultimately does a sum of them to get one output feature.
I'm trying to understand where did the numbers 2 and 8 come from.

[SuperrMurlocc]

Hi! As far as I understood, this part was only used for easier conversion to physical chip. Since I only used the network on the computer (I didn't build the actual FPGA resulting from it) I skipped this part completely. However, please let me know if this is something you would need, I believe change wouldn't be particularly hard and I could post an updated version within a week.

[SampadChowdhury]

@SuperrMurlocc
That would be actually great if that part code is done, thanks
I'm actually trying to understand how the splitting is done. I can't justify the numbers 2 and 8 here.
Were u able to understand whats actually happeining here?
in that case even before the code, i wud like to understand that part first.

[SuperrMurlocc]

Sure thing!

Imagine that you have a convolutional layer with the following parameters:

in_channels = 16
out_channels = 32
kernel_size = 3x3
padding = 1
stride = 1

So effectively, you create 32 filters. And then moving each filter through 16 x 3 x 3 patches. Each filter is in form of a LogicTree, of (let's say) depth = 3. That means that the tree, would be in a form of:

     x
   /    \
  x     x
 /  \   /  \
x  x  x  x
/\ /\  /\  /\

and this were the part "rather than up to eight" comes from. Because without any restrictions, we can choose any 8 inputs of a patch, given that out patch is 16 x 3 x 3 the indices could be for example [0, 2, 2] [1,2,1] [2,0,1] [3,2,1] [4, 1, 2] and so on...

Now that could be hard to build in real world, because routing circuits in such a "weird" way could mean potential chaos. Because of that, they restrict the possible indices of a given filter.

Firstly, we split the model in k/8 groups. That is, for example if k = 32, we create 4 groups, that is we divide our channels into groups, 16/4 = 4, so we still have a patch 16 x 3 x 3 but now we can think of it as (4 x 4) x 3 x 3. And then we split the "filters" too. 32/4 = 8. So first 8. filters would only choose from channels 0-3, second 8 filters would only choose from channels 4-7 and so on.

Secondly, about "each tree observing only 2 input channels", we restrict the number of channels a single LogicTree observes. So for example, 2nd filter is limited to channels 0-3 and we additionally restrict that this very filter will only look at channels number 1 and 3. That means, the final 8 connections in our LogicTree of depth three, can only choose indices that would choose 1st and 3rd channel.

Let me know if you need any further explanation.

[SampadChowdhury]

That makes a lot more sense now. This is roughly what grouped convolution does it seems. Thanks @SuperrMurlocc
Do u have any plan to code this? I'll be trying to implement this anyways. It's always better to have a second set of implementation.

[SampadChowdhury]

@SuperrMurlocc Do you implement Maxpool(ORPool) layer as well in your implementation? In my implementation i see something weird. With just convolution blocks, i can achieve good accuracy, but when i introduce or pool blocks in between my accuracy drops to 10%(basically unstable training). In the paper, the author mentions something about saturation of 1's due to orpool blocks too.

[SuperrMurlocc]

@SampadChowdhury
Hi, sorry for my absence, unfortunately due to current duties i didn't find time to implement grouped convolutions :(
However, when it comes to implementing ORPooling, I simply used regular MaxPooling instead. Please find an example network below:

class MNISTArchitecture(pl.LightningModule):
    def __init__(self, *,
                 model_scale: int = 16,
                 temperature: float = 6.5,
                 lr: float = 0.01):
        super().__init__()

        self.output_gate_factor = 2 if model_scale <= 64 else 1
        self.temperature = temperature

        self.k = model_scale
        self.lr = lr

        self.save_hyperparameters()

        self.model = nn.Sequential(
            Conv(in_channels=1, out_channels=self.k, kernel_size=5, padding=0, stride=1, depth=3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            Conv(in_channels=self.k, out_channels=3 * self.k, kernel_size=2, padding=1, stride=1, depth=3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            Conv(in_channels=3 * self.k, out_channels=9 * self.k, kernel_size=2, padding=1, stride=1, depth=3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            Logic(in_dim=81 * self.k, out_dim=1280 * self.k * self.output_gate_factor),
            Logic(in_dim=1280 * self.k * self.output_gate_factor, out_dim=640 * self.k * self.output_gate_factor),
            Logic(in_dim=640 * self.k * self.output_gate_factor, out_dim=320 * self.k * self.output_gate_factor),
            Aggregate(num_classes=10, temperature=self.temperature)
        )

        self.loss_function = torch.nn.CrossEntropyLoss()

        self.validation_accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=10)

    def forward(self, x):
        return self.model(x)

    def logits2preds(self, logits):
        return torch.argmax(logits, dim=1)

    def predict(self, x):
        return self.logits2preds(self(x))

    def prepare_images(self, images):
        images = torch.where(images != 0, torch.ones_like(images), torch.zeros_like(images))

        return images

    def training_step(self, batch, batch_idx):
        images, labels = batch
        images = self.prepare_images(images)

        x = self.forward(images)
        loss = self.loss_function(x, labels)

        self.log(f'train_loss', loss)

        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        images = self.prepare_images(images)

        x = self.forward(images)
        loss = self.loss_function(x, labels)

        predictions = self.logits2preds(x)
        self.validation_accuracy.update(predictions, labels)

        self.log(f'val_loss', loss)

        return loss

    def on_validation_epoch_end(self):
        acc = self.validation_accuracy.compute()
        self.log('val_accuracy', acc, prog_bar=True)
        self.validation_accuracy.reset()

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        return optimizer