Holding a linear layer weight at zero through backpropagation

konkrer: <h1><a name="p-454430-how-does-one-create-a-permanently-partially-connected-linear-layer-1" class="anchor" href="#p-454430-how-does-one-create-a-permanently-partially-connected-linear-layer-1"></a>How does one create a permanently partially connected linear layer?</h1>
<p>I am trying to create a Linear layer where certain weights are held at zero and skipped during backpropagation so the weight doesn’t change from zero.</p>
<p>In the below image, imagine one of the connecting lines having a permanent weight of zero preventing communication between two neurons or between an input and neuron.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b1a2791569f6f728776eee2b30821757b74b99a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7b1a2791569f6f728776eee2b30821757b74b99a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7b1a2791569f6f728776eee2b30821757b74b99a_2_517x253.jpeg" alt="image" data-base62-sha1="hz0LpIs5ywlU6B5uB8uqmtPw1mO" width="517" height="253" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7b1a2791569f6f728776eee2b30821757b74b99a_2_517x253.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7b1a2791569f6f728776eee2b30821757b74b99a_2_775x379.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b1a2791569f6f728776eee2b30821757b74b99a.jpeg 2x" data-dominant-color="E5E5EA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">791×388 72.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’ve been trying to accomplish this by experimenting at the tensor level by setting<br>
<code>requires_gradient = False</code><br>
for the scalar tensor at some index of a containing tensor, but I can’t get it to work.</p>
<p>After seeing the error message:<br>
<code>RuntimeError: you can only change requires_grad flags of leaf variables. If you want to use a computed variable in a subgraph that doesn't require differentiation use var_no_grad = var.detach().</code></p>
<p>I’ve tried detaching the scalar tensor and re-attaching it to the containing tensor, but requires_grad changes to True when the detached tensor is reattached and I end up right where I started.</p>
<p>Is there any way to hold a linear layer weight at zero through backpropagation? (i.e. backpropagation does not propagate through the connection - as if it were not there.)</p>

pratikkorat26: <p>To ensure that gradients are zero for weights that are permanently set to zero during backpropagation, you can use a masking technique. This will prevent any updates to those weights during the optimization process.</p>
<h2><a name="p-454451-implementing-zero-gradients-for-permanently-zero-weights-1" class="anchor" href="#p-454451-implementing-zero-gradients-for-permanently-zero-weights-1"></a>Implementing Zero Gradients for Permanently Zero Weights</h2>
<p>Example</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MaskedLinear(nn.Module):
    def __init__(self, in_features, out_features, mask):
        super(MaskedLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias = nn.Parameter(torch.Tensor(out_features))
        self.mask = mask  # Binary mask to control connectivity
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in)
        nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input):
        return F.linear(input, self.weight * self.mask, self.bias)


in_features = 4
out_features = 3
mask = torch.tensor([[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 1]], dtype=torch.float32)

layer = MaskedLinear(in_features, out_features, mask)
input_data = torch.randn(2, in_features)
output = layer(input_data).sum()

output.backward()
</code></pre>
<h3><a name="p-454451-output-results-2" class="anchor" href="#p-454451-output-results-2"></a>Output results</h3>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/2/42dd1c4633b08df7d175d8a9b7252e9fb19e4863.png" alt="{A1F3FBDB-B37E-4172-AA8E-A30F13A77BCA}" data-base62-sha1="9xveSxIGRhQ40UVkoiZOjfxAqWv" width="505" height="276"></p>
<p>I still did not understand why you want to set some specific parameters to zero</p>

konkrer: <p>Awesome, thanks so much! Really appreciate it. I think this should allow me to accomplish what I’m trying to do; I’m going to implement it soon.</p>
<p>As to why - I’m experimenting with a novel network architecture, and hopefully this will allow faster processing than the method I’m currently using to accomplish my particular network connections.</p>

konkrer: <p>I suppose, yes. That’s another way to put it.</p>
<p>Is there a better way? (i.e. faster processing than the solution provided)</p>

konkrer: <p>So that MaskedLinear class allowed me to do exactly what I wanted. The only problem is that it’s quite slow. It doesn’t matter much anymore for what I’m doing though, as I discovered another shortcoming using this approach for my endeavors.</p>
<p>Here’s the class streamlined using nn.Linear inheritance:</p>
<pre><code class="lang-auto">from torch import Tensor, no_grad
from torch.nn import Linear
import torch.nn.functional as F

class MaskedLinear(Linear):
    """ Linear layer that allows for masking of weights preventing
    backpropagation through masked network connections. Masked
    weights initialized to zero.

    Uses a binary mask. i.e. -
    mask = torch.tensor([
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        ],
        dtype=torch.float32)
    """

    def __init__(self, in_features: int, out_features: int, mask: Tensor,
                 bias: bool = True, device=None, dtype=None) -&gt; None:
        super().__init__(in_features, out_features, bias, device, dtype)
        self.mask = mask
        with no_grad():
            self.weight *= self.mask

    def forward(self, input: Tensor) -&gt; Tensor:
        return F.linear(input, self.weight * self.mask, self.bias)
</code></pre>
<p>It would be nice if there were a way to make this type of thing happen without it being so slow. Theoretically, it seems like it should be possible.</p>
<p>If half the weights in a linear layer were zero and masked, that would mean half as many calculations in the forward pass (well half as many not-trivial calculations anyway) and half the backpropagation. If such a layer could be as fast as a non-masked layer of the same size that’d be good. Currently, with the MaskedLinear it’s slower.</p>
<p>It’s the <code>self.weight * self.mask</code> that makes it slow. If one could just set <code>require_gradient</code> to false for the masked weights instead, I assume it would be faster.</p>