How to prevent very large values in final Linear layer?

username: Laya1
cooked: <p>I have a classification model with 3 hidden layers with ReLU as activation. The final layer of the model is a linear layer. <strong>At the time of prediction, the model produces extremely large values, like <code>1.418e4</code></strong>. I use a function to perform the <code>softmax</code> operation and obtain the class probabilities as follows:</p>
<pre><code class="lang-auto">def softmax(x):
    y = torch.zeros_like(x, dtype=torch.float64)
    for i in range(x.shape[0]):
        x_temp = x[i, :] - max(x[i, :]) # numerical stability
        z = torch.exp(x_temp)
        y[i, :] = z / torch.sum(z)
    return y

</code></pre>
<p>Since the range of values of the linear layer is very wide (over 4 orders of magnitude), the output of the  <strong><code>softmax</code> operation results in a one-hot encoded output instead of continuous-valued probabilities</strong>. Since this is the last layer of the model, I am not using <code>BatchNormalization</code> in the model.<br>
<strong>Is there any way to prevent the model from producing extremely large numbers as outputs?</strong></p>

username: KFrank
cooked: <p>Hi Laya!</p>
<aside class="quote no-group" data-username="Laya1" data-post="1" data-topic="147054" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/8c91f0/48.png" class="avatar"> Laya1:</div>
<blockquote>
<p>the  <strong><code>softmax</code> operation results in a one-hot encoded output instead of continuous-valued probabilities</strong>.</p>
</blockquote>
</aside>
<p>Based on the information you have provided, this is not necessarily<br>
a problem.</p>
<p>Having the predicted probabilities be, in effect, one-hot encoded just<br>
means that your model is making highly certain predictions.  Is there<br>
any reason that your model shouldn’t make highly certain predictions?</p>
<p>Depending on the character of your data and how you train, this could<br>
be the expected outcome of successful training.</p>
<p>(As an aside, why do you write your own <code>softmax()</code> function instead<br>
of using <code>torch.softmax()</code>?)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Laya1
cooked: <p>Thank you for the reply <a class="mention" href="/u/kfrank">@KFrank</a>. I think the model can make highly certain predictions, but not all the time. The data used for training is synthetically generated, and I know that there are some samples where even the ground truth has some uncertainty associated with it. What is troubling me is that the network is confident for all samples, even after training 20 of them with different initialisations and train-test splits.</p>
<p>I wrote my own <code>softmax()</code> because <code>torch.Softmax()</code> was producing values that were a bunch of ones and something like <code>2.712</code> for each example, which violates everything we know about probability distributions. When I wrote my own <code>softmax()</code>, I found out that the bunch of ones became zeros and the <code>2.712</code> became <code>1</code> and I ended up with one-hot encoded outputs. The classification was correct, but the output of <code>torch.Softmax()</code> was weird.</p>

username: KFrank
cooked: <p>Hi Laya!</p>
<aside class="quote no-group" data-username="Laya1" data-post="3" data-topic="147054" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/8c91f0/48.png" class="avatar"> Laya1:</div>
<blockquote>
<p>I wrote my own <code>softmax()</code> because <code>torch.Softmax()</code> was producing values that were a bunch of ones and something like <code>2.712</code> for each example, which violates everything we know about probability distributions.</p>
</blockquote>
</aside>
<p><code>torch.softmax()</code> works for me – and agrees with your <code>softmax()</code><br>
version:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt; def softmax(x):
...     y = torch.zeros_like(x, dtype=torch.float64)
...     for i in range(x.shape[0]):
...         x_temp = x[i, :] - max(x[i, :]) # numerical stability
...         z = torch.exp(x_temp)
...         y[i, :] = z / torch.sum(z)
...     return y
...
&gt;&gt;&gt; t = torch.randn (3, 5, dtype = torch.float64)
&gt;&gt;&gt; r1 = softmax (t)
&gt;&gt;&gt; r2 = torch.softmax (t, dim = 1)
&gt;&gt;&gt; torch.allclose (r1, r2)
True
</code></pre>
<p>Double-check that you don’t just have a bug somewhere,  but if you<br>
are seeing <code>torch.softmax()</code> return bogus values, could you tell us<br>
what version of pytorch you are using and post a complete, runnable<br>
script that reproduces the issue?</p>
<p>You don’t say how you are training your classifier or what loss function<br>
you are using, but if you’re using <code>CrossEntropyLoss</code>, please be sure<br>
that you are passing the output of your final <code>Linear</code> layer directly to<br>
<code>CrossEntropyLoss</code> <em>without</em> passing it through <code>softmax()</code> or anything<br>
similar.  Doing so would “weaken” the certainty of the model’s predictions<br>
so in order to be feeding moderately certain predictions to your loss<br>
function, you would actually be training your model to make very certain<br>
predictions.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Laya1
cooked: <p>Thank you for the reply <a class="mention" href="/u/kfrank">@KFrank</a>. I am using <code>torch.__version__</code> <code>'1.10.2'</code>. I am using an assumed density filtering version of a fully connected neural network inspired from this link: <a href="https://github.com/uzh-rpg/deep_uncertainty_estimation" rel="noopener nofollow ugc">deep_uncertainty_estimation</a>. Every layer takes in two arguments (mean, variance) and produces two outputs (mean, variance).</p>
<p>For the prupose of training of the model, the difference from a regular neural network is the <code>ReLU</code> layer here: <a href="https://github.com/uzh-rpg/deep_uncertainty_estimation/blob/master/contrib/adf.py#L73" rel="noopener nofollow ugc">relu</a> and <code>BatchNormalization</code> layer here: <a href="https://github.com/uzh-rpg/deep_uncertainty_estimation/blob/master/contrib/adf.py#L211" rel="noopener nofollow ugc">bn</a> (I am using a variant for 1D BN). Here is my network architecture:</p>
<pre><code class="lang-auto">class FCModel(torch.nn.Module):
    def __init__(self, input_dim, output_dim, layer_sizes, dropout_prob=0.2):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        layers = []
        layers.append(adf.Linear(in_features=self.input_dim, out_features=layer_sizes[0]))
        layers.append(adf.BatchNorm1D(num_features=layer_sizes[0], track_running_stats=False))
        layers.append(adf.ReLU())
        for i in range(1, len(layer_sizes)):
            layers.append(adf.Linear(in_features=layer_sizes[i - 1], out_features=layer_sizes[i]))
            layers.append(adf.BatchNorm1D(num_features=layer_sizes[i], track_running_stats=False))
            layers.append(adf.ReLU())
            layers.append(adf.Dropout1D(dropout_prob))
        layers.append(adf.Linear(in_features=layer_sizes[-1], out_features=self.output_dim))
        self.layers = adf.Sequential(*layers)

    def forward(self, x_mean, x_var):
        x_mean, x_var = self.layers(x_mean, x_var)
        return x_mean, x_var

</code></pre>
<p>I am using <code>CrossEntropyLoss()</code> and there is no <code>torch.nn.Softmax()</code> layer in the model. I use <code>softmax()</code> only at the time of evaluation. I actually made the mistake of using <code>torch.nn.Softmax()</code> along with <code>CrossEntropyLoss()</code> and then removed the layer later on.</p>
<p>In a different version of the model, I had a custom <code>Softmax()</code> layer which was activated only when <code>self.training</code> was <code>False</code> so that the unnormalized outputs of the last <code>Linear</code> layer are fed to <code>CrossEntropyLoss()</code> and the network produces class probabilities instead of logits at the time of evaluation. That is when I found weird values output by the <code>torch.nn.Softmax()</code> layer.</p>
<p>Hope this helps.</p>