Hi guys!
I’m deploying a sparse autoencoder. Its structure is:
Encoder input dimension: 20
Encoder hidden layers dimension: {80,40}
Encoder output dimension: 4 ( == Decoder input dimension)
Decoder hidden layers dimension: {40,80}
Decoder output dimension: 20 ( == Encoder input dimension)

and the corresponding code is:
context_dimension = 20
output_neurons = 4
class SparseAutoencoder(nn.Module):
    def __init__(self):
        super(SparseAutoencoder, self).__init__()
 
        # encoder
        self.enc1 = nn.Linear(in_features=context_dimension, out_features=80)
        self.enc2 = nn.Linear(in_features=80, out_features=40)
        self.enc3 = nn.Linear(in_features=40, out_features=output_neurons)
 
        # decoder 
        self.dec1 = nn.Linear(in_features=output_neurons, out_features=40)
        self.dec2 = nn.Linear(in_features=40, out_features=80)
        self.dec3 = nn.Linear(in_features=80, out_features=context_dimension)

    def forward(self, x):
        
        # encoding
        x = F.relu(self.enc1(x))
        x = F.relu(self.enc2(x))
        x = F.relu(self.enc3(x))
        
        y = x # encoded context
 
        # decoding
        x = F.relu(self.dec1(x))
        x = F.relu(self.dec2(x))
        x = F.relu(self.dec3(x))

        return y,x

I want to add a sparsity penalty at the bottleneck layer. In particolar, I would like to compute the activation frequency of each output neuron over an entire batch, and try to minimize it.
The idea is the following:
def kl_divergence(p, q):
    s1 = torch.sum(p * torch.log(p / q))
    s2 = torch.sum((1 - p) * torch.log((1 - p) / (1 - q)))
    return s1 + s2

encoded,outputs = model(data)  # encoded is a tensor of dimension [batch_size, output_neurons]
[...] # The MSE_loss is computed based on data and outputs
if ADD_SPARSITY:
      rho_hat = (encoded > 0).sum(dim=0) # given the output neuron, sum of non-zeros outputs
      rho_hat = torch.div(rho_hat, dataloader.batch_size/context_dimension) # averaging over the batch
      # This for and the nested conditions are needed in order to avoid kl_loss=inf or kl_loss=nan
      for index in range(rho_hat.shape[0]):
          if rho_hat[index]==0:
              rho_hat[index] = 0.001
          if rho_hat[index]==1:
              rho_hat[index] = 0.999
      rho = torch.FloatTensor([RHO for _ in range(output_neurons)]).unsqueeze(0) # a vector of RHOs
      kl_loss = BETA * kl_divergence(rho, rho_hat)
      loss = MSE_loss + kl_loss

By applying this procedure, the NN learns how to minimize the MSE_loss, without taking into account the sparsity term.
I guess that the problem about this approach is that the kl_loss is not differentiable because of the for cycle.
The only solution that comes into my mind is to avoid the kl_loss to be inf by applying an initialization procedure to the NN weights. Namely, knowing that the input data is bounded into [0,1], make the outputs be partially negative and partially positive for every input. But I know that this idea is hardly implementable.
That is why I’m open to any suggestion.
Thank you all in advance