Implementing Truncated Normal Initializer

SivanK: <p>I would like to implement an initialization identical to tf.truncated_normal_initializer.<br>
I read a suggestion of doing: m.weight.data.normal_(1.0, 0.02).clamp_(min=0,max=2), but it is not good as values with more than two standard deviations from the mean should be discarded and re-drawn, and not clamped to a constant max/min value.<br>
Do you have a suggestion of how can I implement it?</p>

alexis-jacq: <p>I would suggest a monte-carlo sampling imitating the one used by <code>torch.normal</code>, but instead of using the normal pdf (let’s cal it “f”), you use this pdf:</p>
<p>trunc_f = max( 0 , f - f(2*std) )</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/1X/d6e532f6de9f24721f1a487915423cebe7641224.jpg" data-download-href="https://discuss.pytorch.org/uploads/default/d6e532f6de9f24721f1a487915423cebe7641224" title="2017-07-11 14.14.48.jpg"><img src="https://discuss.pytorch.org/uploads/default/optimized/1X/d6e532f6de9f24721f1a487915423cebe7641224_2_690x388.jpg" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/1X/d6e532f6de9f24721f1a487915423cebe7641224_2_690x388.jpg, https://discuss.pytorch.org/uploads/default/optimized/1X/d6e532f6de9f24721f1a487915423cebe7641224_2_1035x582.jpg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/1X/d6e532f6de9f24721f1a487915423cebe7641224_2_1380x776.jpg 2x" data-dominant-color="BEBBB6"><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">2017-07-11 14.14.48.jpg</span><span class="informations">3771×2121 1.05 MB</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>In fact, you have to first truncate a centered normal with std=1, do the sampling, then you multiply the values by std and add the mean.</p>

uapatira: <p><a class="mention" href="/u/sivank">@SivanK</a> - the reason <a class="mention" href="/u/alexis-jacq">@alexis-jacq</a> suggests you substract 2stdev (or however much you want to truncate) from your normal is because if you do a naive min/max, you’ll end up with two spikes on either end of your curve instead of a nice taper.</p>

alexis-jacq: <p>Wow, in fact, there are no monte-carlo method behind torch.normal (wich is a backend of the TH normal function from THRandom.c). It is based on the <a href="https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform">box-muller transform</a>.</p>

alexis-jacq: <p>The box-muller trick is to sample, uniformly between 0 and 1 two random variables, U1 and U2. Then, you can verify that</p>
<pre><code class="lang-auto">Z = sqrt(-2*log(U1)) * cos(2*pi*U2) = R*cos(Theta)
</code></pre>
<p>follows a normal distribution. Now, we want to create Z only living between -2 and 2. For this truncation, we still accept all the angles Theta, but we just want the ray R bigger than -2 and smaller than 2. In other words:</p>
<pre><code class="lang-auto">-2 &lt;= sqrt(-2*log(U1)) &lt;= 2
</code></pre>
<p>iif U1 is uniformly sampled between exp(-2) and 1 <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Then you have the python function:</p>
<pre><code class="lang-python"># sample u1:
size = m.weight.size()
u1 = torch.rand(size)*(1-np.exp(-2)) + np.exp(-2)
# sample u2:
u2 = torch.rand(size)
# sample the truncated gaussian ~TN(0,1,[-2,2]):
z = torch.sqrt(-2*log(u1)) * torch.cos(2*np.pi*u2)
m.weight.data = z
</code></pre>
<p>This led to a pull request for THRandom.c: <a href="https://github.com/torch/TH/pull/4">https://github.com/torch/TH/pull/4</a></p>

SivanK: <p><a class="mention" href="/u/uapatira">@uapatira</a>  - that’s why I asked this question in the first place <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
<a class="mention" href="/u/alexis-jacq">@alexis-jacq</a> - Thanks for the great answer!</p>

ruotianluo: <p><a class="mention" href="/u/alexis-jacq">@alexis-jacq</a>  This approximation is not close to the true distribution.</p>

alexis-jacq: <p>Well it seems ok with plots:</p>
<p>in blue, the normal distribution (sorted)<br>
in red, the truncated distribution (also sorted)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/1X/34b462e6df8ef346bb513fa7049cf8012c11dff1.png" data-download-href="https://discuss.pytorch.org/uploads/default/34b462e6df8ef346bb513fa7049cf8012c11dff1" title="plots.png"><img src="https://discuss.pytorch.org/uploads/default/optimized/1X/34b462e6df8ef346bb513fa7049cf8012c11dff1_2_666x500.png" width="666" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/1X/34b462e6df8ef346bb513fa7049cf8012c11dff1_2_666x500.png, https://discuss.pytorch.org/uploads/default/original/1X/34b462e6df8ef346bb513fa7049cf8012c11dff1.png 1.5x, https://discuss.pytorch.org/uploads/default/original/1X/34b462e6df8ef346bb513fa7049cf8012c11dff1.png 2x" data-dominant-color="FBFAFA"><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">plots.png</span><span class="informations">800×600 14.8 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>

tom: <p>Well, the method confines the random variables to a circle when you want a square [Edit: The square would not work either as it has the wrong marginals, too. One needs the full “other dimension” for each coordinate]. It probably works just fine for the initialisation heuristic, but it is far from general purpose.<br>
(Interestingly, a similar transformation to polar coordinates is the easiest way I know to compute the normalisation constant in the normal transform. If you could apply it to squares you would get a closed form for the normal CDF.)</p>
<p>Best regards</p>
<p>Thomas</p>

alexis-jacq: <p>I don’t understand why do you want a square… Don’t we want to confine into a circle of ray 2? And anyway, here my example is in 1d, so the circle and the square are the same, it’s the segment [-2, 2]. The box-muller method for truncated normal is already known and used, since this paper: <a href="http://oa.upm.es/22659/1/INVE_MEM_2012_153251.pdf">http://oa.upm.es/22659/1/INVE_MEM_2012_153251.pdf</a></p>

ruotianluo: <p><a class="mention" href="/u/alexis-jacq">@alexis-jacq</a> I haven’t looked at the paper. But from my simulation, the sampled points are not following the true truncated normal distribution,<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/7/76955fb9a09822741378a37ba1bed392206f8004.png" width="396" height="280"></p>
<p>The red line is the pdf given by scipy.stats.truncnorm</p>

tom: <p>This seems to be the picture I had in mind when I said that it was with cutting off too much by restricting to the circle…</p>
<p>If you can do with an approximation, <code>torch.fmod(torch.randn(size),2)</code> might be easiest.<br>
It moves more of the cut-off mass to the center than to the boundaries, but at least the good samples  stay where they belong…</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/1/17ca93f305f657a7e41889475750c5cfab94f276.png" width="375" height="252"></p>

alexis-jacq: <p>Wow, and your approximation takes one line!</p>
<p>So it looks like you take all the element over 2 and put them close to the mean, like in a snake game, and that’s why you have still too many people close to the mean. And it seems that the worst case comes with variance = 1. I would worth to write down the mathematics, to see how far we are from a Gaussian with both methods</p>

tom: <p>Hi Alexis,</p>
<p>Well, I think one of the things is that your sketch seems different from the <a href="https://en.wikipedia.org/wiki/Truncated_normal_distribution" rel="nofollow noopener">usual definition of truncated Gaussian</a>.<br>
So what the .fmod in sampling does in terms of pdf is to sum the pdf over all x+sign(x)<code>*</code>cutoff<code>*</code>i | i=0,1,2,… . If the shape of pdf in [cutoff<code>*</code>i, cutoff<code>*</code>(i+1)] were proportional to [0,cutoff], this would be exact. As it is falls of much more quickly than proportional (a 3 std event conditional on the magnitude being at least 2 std is much less likely than a 1 std event). This is why summation puts too much mass in the center. With a smaller cutoff, the pdf in [-cutoff,cutoff] becomes flatter (closer to uniform) and thus the error is more pronounced. And this would further increase if you cut of at less than 1 std, but at some point you might just use a uniform distribution anyway. Even at 1 you could use a cutoff of 2/3 (or so) and then multiply the sample by 3/2 to approximate the truncated normal better (i.e. <code>torch.fmod(torch.randn(size),2/3)*3/2</code>).</p>
<p>In the plain Box-Muller method, you are essentially using two transforms: first that for U_1 uniform on [0,1] R = -sqrt(2 ln(U_1)) has the distribution of the radii in a 2d standard normal, χ(2). Thus that X,Y = sin/cos(U_2) * R will be 2d standard normal. Then you are using that the marginals of the 2d standard normal are 1d standard normals.</p>
<p>If you now restrict R to R &lt;= c, and take the pdf of X close to c, you have only very little of the total mass along the vertical (Y-coordinate) lies in the disc R&lt;=c. This is why the “restricted BM” distribution has the density going to 0 as X approaches c.</p>
<p>But then I would be surprised if the initialisation heuristic broke down on with a simple approximation.</p>
<p>Best regards</p>
<p>Thomas</p>

alexis-jacq: <p>Ok thanks, now I understand better. So, close to x= 0, all samples should be accepted, since cos(U_2) is small, but with my approach, I still multiply by a restricted R and I make it even smaller. That explains the plot of <a class="mention" href="/u/ruotianluo">@ruotianluo</a> where I pass over expected density close to the mean.</p>
<p>It was too beautiful to work…</p>
<p>Once again, I should be more careful when emotions mix with mathematics…</p>

ruotianluo: <p>Coming back to this thread after one year.</p>
<p>One simple approximated implementation for truncated_normal to 2*std.</p>
<pre><code class="lang-auto">def truncated_normal_(tensor, mean=0, std=1):
    size = tensor.shape
    tmp = tensor.new_empty(size + (4,)).normal_()
    valid = (tmp &lt; 2) &amp; (tmp &gt; -2)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
</code></pre>

arijit17: <p>Hi! Does it implement same thing as <a href="https://www.tensorflow.org/api_docs/python/tf/truncated_normal" rel="nofollow noopener">https://www.tensorflow.org/api_docs/python/tf/truncated_normal</a> ?</p>

yunjey: <p>Here is a simpler way to sample values from truncated normal distribution.</p>
<pre><code class="lang-auto">from scipy.stats import truncnorm
import torch


def truncated_normal(size, threshold=1):
    values = truncnorm.rvs(-threshold, threshold, size=size)
    return values

# usage example
x= truncnorm([10, 20], threshold=1)   # sample 10x20 sized tensor
x = torch.from_numpy(x).cuda()
</code></pre>

Omry_Yadan: <p>Thanks <a class="mention" href="/u/ruotianluo">@ruotianluo</a>.<br>
I confirmed that your function is equivalent to scipy truncnorm.</p>
<p>if anyone wants reproduce:</p>
<pre><code class="lang-auto">import torch
from scipy.stats import truncnorm
import matplotlib.pyplot as plt

def truncated_normal_(tensor, mean=0, std=1):
    size = tensor.shape
    tmp = tensor.new_empty(size + (4,)).normal_()
    valid = (tmp &lt; 2) &amp; (tmp &gt; -2)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
    return tensor



fig, ax = plt.subplots(1, 1)


def test_truncnorm():
    a, b = -2, 2
    size = 1000000
    r = truncnorm.rvs(a, b, size=size)
    ax.hist(r, density=True, histtype='stepfilled', alpha=0.2, bins=50)

    tensor = torch.zeros(size)
    utils.truncated_normal_(tensor)
    r = tensor.numpy()

    ax.hist(r, density=True, histtype='stepfilled', alpha=0.2, bins=50)
    ax.legend(loc='best', frameon=False)
    plt.show()


if __name__ == '__main__':
    test_truncnorm()

</code></pre>

heiner: <p>The Tensorflow implementation of <code>truncated_normal</code> can be found <a href="https://github.com/tensorflow/tensorflow/blob/b1a6b315a63bb29b4593bfb98095da4397d8cd5a/tensorflow/compiler/tf2xla/lib/random.cc#L42" rel="nofollow noopener">on GitHub</a> and references <a href="https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf" rel="nofollow noopener">https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf</a> (see page 24) there.</p>
<p>Analogous Numpy/PyTorch code would look like</p>
<pre><code class="lang-auto">import numpy as np
import torch


def parameterized_truncated_normal(uniform, mu, sigma, a, b):
    normal = torch.distributions.normal.Normal(0, 1)

    alpha = (a - mu) / sigma
    beta = (b - mu) / sigma

    alpha_normal_cdf = normal.cdf(alpha)
    p = alpha_normal_cdf + (normal.cdf(beta) - alpha_normal_cdf) * uniform

    p = p.numpy()
    one = np.array(1, dtype=p.dtype)
    epsilon = np.array(np.finfo(p.dtype).eps, dtype=p.dtype)
    v = np.clip(2 * p - 1, -one + epsilon, one - epsilon)
    x = mu + sigma * np.sqrt(2) * torch.erfinv(torch.from_numpy(v))
    x = torch.clamp(x, a, b)

    return x


def truncated_normal(uniform):
    return parameterized_truncated_normal(uniform, mu=0.0, sigma=1.0, a=-2, b=2)


def sample_truncated_normal(shape=()):
    return truncted_normal(torch.from_numpy(np.random.uniform(0, 1, shape)))
</code></pre>
<p>However in the case of <code>mu=0.0, sigma=1.0, a=-2, b=2</code>, I suspect repeated normal sampling is also fine: 5% of entries need to be resampled each iteration. E.g. since 0.05**5 * 1000000 &lt; 1 one would need roughly 5 iterations for 1M entries.</p>