Implementing a discontinuous NN layer in `torch`

Federico_Ottomano: <p>I’m trying to implement the following <em>discontinuous layer</em> proposed by a recent interesting <a href="https://www.sciencedirect.com/science/article/pii/S0377042722003430" rel="noopener nofollow ugc">paper</a> ,and I’m wondering whether I’m doing it correctly.</p>
<p>Layer:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/2/d215f55f5256b2823b6d4e910b3970cd9551f983.png" data-download-href="https://discuss.pytorch.org/uploads/default/d215f55f5256b2823b6d4e910b3970cd9551f983" title="discontinuous_nn"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d215f55f5256b2823b6d4e910b3970cd9551f983_2_690x77.png" alt="discontinuous_nn" data-base62-sha1="tYvjvrRotApgPKGIIAqBo57VIGv" width="690" height="77" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d215f55f5256b2823b6d4e910b3970cd9551f983_2_690x77.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d215f55f5256b2823b6d4e910b3970cd9551f983_2_1035x115.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d215f55f5256b2823b6d4e910b3970cd9551f983.png 2x" data-dominant-color="EAEAEA"><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">discontinuous_nn</span><span class="informations">1138×127 73.9 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>My first attempt was the following:</p>
<pre><code class="lang-auto">class DiscontLayer(nn.Module):
    
    def __init__(self, input_dim, output_dim):
        
        super().__init__()
        
        self.linear = nn.Linear(input_dim, output_dim)
        self.epsilon = nn.Parameter(torch.randn(output_dim))
                
    def forward(self, x):
        
        cont_part = nn.ReLU()(self.linear(x))
            
        discont_part = self.epsilon * torch.heaviside(self.linear(x),torch.ones_like(self.linear(x)))
        
        return cont_part + discont_part
</code></pre>
<p>But trying to run it, I encountered an error due to backpropagating through the Heaviside activation function. The derivative of this function is should be 0 everywhere, except at the discontinuity, but this seems not being implemented. Then I found this post, suggesting to protect the <code>torch.heaviside</code> call with <code>torch.no_grad()</code>. So I tried in the following way:</p>
<p>Then:</p>
<pre><code class="lang-auto">class DiscontLayer(nn.Module):
    
    def __init__(self, input_dim, output_dim):
        
        super().__init__()
        
        self.linear = nn.Linear(input_dim, output_dim)
        self.epsilon = nn.Parameter(torch.randn(output_dim))
                
    def forward(self, x):
        
        cont_part = nn.ReLU()(self.linear(x))
        
        discont_part = self.linear(x)
        
        with torch.no_grad():
            
            discont_part = torch.heaviside(discont_part,torch.ones_like(discont_part))
            
        discont_part = self.epsilon * discont_part
        
        return cont_part + discont_part
</code></pre>
<p>But I’m not yet 100% sure that I have done that correctly.</p>
<p>Many thanks for your help,</p>
<p>Fede</p>

joao_romana: <p>Hey Fede, were you able to solve this problem? I was also interested in the effect of using a discontinuous activation function.</p>