How to implement permanent weight dropout in a layer?

Uncategorized
1

Hi! I am trying to implement a weight dropout mechanism, that drops weights from a feed forward layer and keeps them zero during training. I am wondering, what is the best way to do this?

class SparseLinear(torch.nn.Module):
    """
    Linear layer that uses sparse dense weights.
    """

    def __init__(self, in_features, out_features, sparsity):
        super().__init__()

        w = torch.empty(out_features, in_features)
        torch.nn.init.sparse_(w, sparsity=sparsity)
        self.w = torch.nn.Parameter(w)

        b = torch.empty(1, out_features)
        torch.nn.init.sparse_(b, sparsity=sparsity)
        self.b = torch.nn.Parameter(b)

        self.w_mask = self.w == 0
        self.b_mask = self.b == 0

    def forward(self, inputs):
        return torch.nn.functional.linear(
            inputs,
            torch.masked_fill(self.w, self.w_mask, 0.0),
            bias=torch.masked_fill(self.b, self.b_mask, 0.0))

This is my current solution, however I don’t like that it performs masking on every forward call, I would rather prefer a solution that prevent the optimizer from modifying them in the first place.

2

I think you would have to mask these parameters at some point (either directly or their gradients to prevent an update).
The optimizer cannot take parts of tensors, so I’m not aware of any other clean approach. :confused:

3

May be something like this could work:

from torch import nn
from torch.nn import functional as F

class ConvDropW(nn.Module):
    def __init__(self, *args  ):
        super().__init__()
        self.conv = nn.Conv2d( *args )
        self.drop_rate = 0.5
        self.args = args[3:]

    def forward(self, x):
        W, B = self.conv.parameters()
        Wd = F.dropout( W, self.drop_rate )
        return F.conv2d( x, Wd, B, *self.args )

In my test this module not work great, but not break things either.

4

A common approach is to mask the gradients instead of the weights. Register a gradient hook on the parameter (or zero the gradients after backward()) so masked weights never receive updates. Then apply the mask once after each optimizer step (weight.data.mul_(mask)) to keep those entries at zero. This avoids doing masked_fill on every forward pass while ensuring pruned weights stay zero.