Variational dropout?

hughperkins: <p>Hi, what is the standard-ish way to do variational dropout in PyTorch?</p>
<p>(Edit: I just need something that works, and can plug in; don’t need to understand how it works, just how to use it <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"> )</p>
<p>(edit2: though one or two sentences of intuition behind how ti works / what it is doing would be very welcome <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"> )</p>

justusschock: <p>These implementations seem pretty similar and straight forward:</p>
<p><a href="https://www.google.de/url?sa=t&amp;source=web&amp;rct=j&amp;url=https://github.com/j-min/Dropouts&amp;ved=2ahUKEwj2iuDuhurcAhVDUlAKHfBiA1sQFjAAegQIARAB&amp;usg=AOvVaw2apbAB14XAPtcq2mOZHPgM" rel="nofollow noopener">Number One</a><br>
<a href="https://www.google.de/url?sa=t&amp;source=web&amp;rct=j&amp;url=https://github.com/kefirski/variational_dropout&amp;ved=2ahUKEwj2iuDuhurcAhVDUlAKHfBiA1sQFjABegQIBxAB&amp;usg=AOvVaw24HeGDVfGdSatm_KfnFecm" rel="nofollow noopener">Number two</a></p>
<p>Regarding your second edit: I haven’t even tried understanding it, since the paper is still a part of my (always growing) reading list.</p>

hughperkins: <p>Thank you <a class="mention" href="/u/justusschock">@justusschock</a> It looks like these are both dropout. The first looks like ZoneOut? And the second looks like standard dropout? Am I mis-reading? (I’m looking specifically for ‘variational’ dropout)</p>

justusschock: <p>In the first repo there is a jupyter notebook containing several variations of dropout including this:</p>
<pre><code class="lang-auto">class VariationalDropout(nn.Module):
    def __init__(self, alpha=1.0, dim=None):
        super(VariationalDropout, self).__init__()
        
        self.dim = dim
        self.max_alpha = alpha
        # Initial alpha
        log_alpha = (torch.ones(dim) * alpha).log()
        self.log_alpha = nn.Parameter(log_alpha)
        
    def kl(self):
        c1 = 1.16145124
        c2 = -1.50204118
        c3 = 0.58629921
        
        alpha = self.log_alpha.exp()
        
        negative_kl = 0.5 * self.log_alpha + c1 * alpha + c2 * alpha**2 + c3 * alpha**3
        
        kl = -negative_kl
        
        return kl.mean()
    
    def forward(self, x):
        """
        Sample noise   e ~ N(1, alpha)
        Multiply noise h = h_ * e
        """
        if self.train():
            # N(0,1)
            epsilon = Variable(torch.randn(x.size()))
            if x.is_cuda:
                epsilon = epsilon.cuda()

            # Clip alpha
            self.log_alpha.data = torch.clamp(self.log_alpha.data, max=self.max_alpha)
            alpha = self.log_alpha.exp()

            # N(1, alpha)
            epsilon = epsilon * alpha

            return x * epsilon
        else:
            return x
</code></pre>
<p>And the second repo contains this implementation:</p>
<pre><code class="lang-auto">
class VariationalDropout(nn.Module):
    def __init__(self, input_size, out_size, log_sigma2=-10, threshold=3):
        """
        :param input_size: An int of input size
        :param log_sigma2: Initial value of log sigma ^ 2.
               It is crusial for training since it determines initial value of alpha
        :param threshold: Value for thresholding of validation. If log_alpha &gt; threshold, then weight is zeroed
        :param out_size: An int of output size
        """
        super(VariationalDropout, self).__init__()
 
        self.input_size = input_size
        self.out_size = out_size
 
        self.theta = Parameter(t.FloatTensor(input_size, out_size))
        self.bias = Parameter(t.Tensor(out_size))
 
        self.log_sigma2 = Parameter(t.FloatTensor(input_size, out_size).fill_(log_sigma2))
 
        self.reset_parameters()
 
        self.k = [0.63576, 1.87320, 1.48695]
 
        self.threshold = threshold
 
    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.out_size)
 
        self.theta.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)
 
    @staticmethod
    def clip(input, to=8):
        input = input.masked_fill(input &lt; -to, -to)
        input = input.masked_fill(input &gt; to, to)
 
        return input
 
    def kld(self, log_alpha):
 
        first_term = self.k[0] * F.sigmoid(self.k[1] + self.k[2] * log_alpha)
        second_term = 0.5 * t.log(1 + t.exp(-log_alpha))
 
        return -(first_term - second_term - self.k[0]).sum() / (self.input_size * self.out_size)
 
    def forward(self, input):
        """
        :param input: An float tensor with shape of [batch_size, input_size]
        :return: An float tensor with shape of [batch_size, out_size] and negative layer-kld estimation
        """
 
        log_alpha = self.clip(self.log_sigma2 - t.log(self.theta ** 2))
        kld = self.kld(log_alpha)
 
        if not self.training:
            mask = log_alpha &gt; self.threshold
            return t.addmm(self.bias, input, self.theta.masked_fill(mask, 0))
 
        mu = t.mm(input, self.theta)
        std = t.sqrt(t.mm(input ** 2, self.log_sigma2.exp()) + 1e-6)
 
        eps = Variable(t.randn(*mu.size()))
        if input.is_cuda:
            eps = eps.cuda()
 
        return std * eps + mu + self.bias, kld
 
    def max_alpha(self):
        log_alpha = self.log_sigma2 - self.theta ** 2
        return t.max(log_alpha.exp())
</code></pre>
<p>From scrolling through the paper and skimming the equations, it looks like they are both fine although I might also misread something.</p>

hughperkins: <p>Ah, awesome, thanks! Question: how to use them? Like:</p>
<ul>
<li>where do we put them in an rnn?</li>
<li>are the masks shared between eg different timesteps, or something like that?</li>
<li>will this work with eg an nn.LSTM? or do we need to use an nn.LSTMCell, and plug those together, with this dropout in between?
<ul>
<li>if so, is there an LSTM implementation that handles this for us?</li>
</ul>
</li>
</ul>

justusschock: <p>To be honest:</p>
<p>I don’t know if this works with an LSTM, but my first guess would be No, since both of them seem to be designed to be used with sequential networks:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch.nn.functional as F
 
from variational_dropout.variational_dropout import VariationalDropout
 
 
class VariationalDropoutModel(nn.Module):
    def __init__(self):
        super(VariationalDropoutModel, self).__init__()
 
        self.fc = nn.ModuleList([
            VariationalDropout(784, 500),
            VariationalDropout(500, 50),
            nn.Linear(50, 10)
        ])
 
    def forward(self, input, train=False):
        """
        :param input: An float tensor with shape of [batch_size, 784]
        :param train: An boolean value indicating whether forward propagation called when training is performed
        :return: An float tensor with shape of [batch_size, 10]
                 filled with logits of likelihood and kld estimation
        """
 
        result = input
 
        if train:
            kld = 0
 
            for i, layer in enumerate(self.fc):
                if i != len(self.fc) - 1:
                    result, kld = layer(result, train)
                    result = F.elu(result)
                    kld += kld
 
            return self.fc[-1](result), kld
 
        for i, layer in enumerate(self.fc):
            if i != len(self.fc) - 1:
                result = F.elu(layer(result, train))
 
        return self.fc[-1](result)
 
    def loss(self, **kwargs):
        if kwargs['train']:
            out, kld = self(kwargs['input'], kwargs['train'])
            return F.cross_entropy(out, kwargs['target'], size_average=kwargs['average']), kld
 
        out = self(kwargs['input'], kwargs['train'])
        return F.cross_entropy(out, kwargs['target'], size_average=kwargs['average'])
</code></pre>
<p>So I guess, you would have to plug those together yourself.</p>

hughperkins: <aside class="quote no-group" data-username="justusschock" data-post="6" data-topic="23030">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/justusschock/48/4637_2.png" class="avatar"> justusschock:</div>
<blockquote>
<p>So I guess, you would have to plug those together yourself.</p>
</blockquote>
</aside>
<p>Ya, that was sort of my conclusion too <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=12" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:" loading="lazy" width="20" height="20">  Unless there is some paper implementation somewhere, that is in PyTorch, and uses variational dropout?</p>

Shisho_Sama: <p>Hello guys<br>
Any new updates on this a year later?</p>

Iridium_Blue: <p>Hey all, new here, perhaps I can help.</p>
<p>Thanks to kreitkurita of Carnegie Mellon University, who sub-classed LSTM to use Variational Dropout : <a href="https://github.com/keitakurita/Better_LSTM_PyTorch/blob/master/better_lstm/model.py" rel="nofollow noopener">Better LSTM with Variational Dropout</a></p>
<p>Just use this as a drop-in replacement for LSTM. The implementation is an almost faithful implementation of the original paper <a href="https://arxiv.org/abs/1512.05287" rel="nofollow noopener">https://arxiv.org/abs/1512.05287</a> (see code comments for minor deviations.)</p>
<p>Some tips :</p>
<ul>
<li>
<p>0.25 is a good initial choice for dropouti, dropoutw and dropouto (input, weight, and output respectively.)</p>
</li>
<li>
<p>It is probably best to avoid using other dropout techniques alongside this one (embedding-, batch-, layer-, etc.)  At least at first.  And possibly always.   I need to look into this more.</p>
</li>
<li>
<p>In their original paper, Gal and Ghahramani note that <strong>weight_decay</strong> takes on new importance.  They suggest 0.001 as a default.    (This is set on the Optimizer.  If you’re using Adam, I suggest looking into AdamW instead.)</p>
</li>
</ul>