Answers.Body: <p>Mathematically it does not matter at all. The results will be the same. However there is a strong reason to prefer it being in the loss function: numeric stability.</p>
  <p>Because the loss function knows more information (ie what the loss is), it can compute the softmax and loss together using some tricks to improve numeric stability.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T15:33:14.623Z
Answers.Text: Thanks for the explanation. What are these tricks you mention here?
Answers.UserDisplayName: starbeamrainbowlabs
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T16:05:39.540Z
Answers.Text: Here is a blog post explaining the basics. If the above answer was helpful please consider accepting it.
  https://jaykmody.com/blog/stable-softmax/
Answers.UserDisplayName: chessprogrammer
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T16:37:19.003Z
Answers.Text: Many thanks for that link, but I'm afraid I still don't understand. That log post mentioned log(softmax()) for some reason? It did not seem to cover my question
Answers.UserDisplayName: starbeamrainbowlabs
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T17:00:44.083Z
Answers.Text: It did cover your question. Log softmax is more numerically stable. You can convert back to regular softmax after. Perhaps you do not have a sufficient mathematical background to understand the answer. But that link directly addresses your question. The trick is to use log softmax during the intermediate computations and then convert it back after. If this still does not make sense to you, make sure you have a good understanding of what numerical stability means for floating point arithmetic, and then the log softmax trick should be fairly easy to understand.
Answers.UserDisplayName: chessprogrammer
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T17:17:23.687Z
Answers.Text: Right, I understand that, but it still doesn't explain why you wouldn't do the (log) softmax in advance in the last layer of the model.
Answers.UserDisplayName: starbeamrainbowlabs
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-22T17:53:44.647Z
Answers.Text: You can. But then you cannot use cross entropy anymore since cross entropy is not expecting a log softmax. You would need to use a log loss instead. That is a fine alternative.
Answers.UserDisplayName: chessprogrammer
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-17T22:12:21.090Z
Answers.LastActivityDate: 2023-03-17T22:12:21.090Z
Answers.UserDisplayName: chessprogrammer
Body: <p>Suppose I have an image segmentation model with an output of <code>[ 128, 128, 2 ]</code>, segmenting an input image into 2 parts.</p>
  <p>Commonly, loss functions have the sigmoid or softmax needed to produce a probability distribution in the loss function itself, for example:</p>
  <pre><code>tf.nn.sigmoid_cross_entropy_with_logits
  tf.nn.softmax_cross_entropy_with_logits
  </code></pre>
  <p>Alternatively, one could include the <code>softmax</code> or <code>sigmoid</code> call in the last layer fo a model:</p>
  <pre><code>model = tf.keras.Sequential([
      # ...
      tf.keras.layers.Conv2D(2, kernel_size=1, padding=&quot;same&quot;, activation=&quot;softmax&quot;)
  ])
  </code></pre>
  <p>My question is this: Why could you keep the <code>softmax</code> / <code>sigmoid</code> in the loss function instead of in the last layer of the model?</p>
  <p>Wouldn't putting it in the last layer of the model ensure that the output of the model is constrained, enabling the model to more effectively learn the boundaries of the given input as smaller gradients would be needed for the model to change its mind on a given pixel?</p>
  
Tags: <training><image-segmentation><softmax>
Title: Should softmax be in the model or in the loss function?
UserDisplayName: starbeamrainbowlabs
ThreadText: # Should softmax be in the model or in the loss function?
  
  (Asked by: username_0 on 2023-03-17)
  
  Suppose I have an image segmentation model with an output of `[ 128, 128, 2 ]`, segmenting an input image into 2 parts.
  
  Commonly, loss functions have the sigmoid or softmax needed to produce a probability distribution in the loss function itself, for example:
  
  ```
  tf.nn.sigmoid_cross_entropy_with_logits
  tf.nn.softmax_cross_entropy_with_logits
  
  ```
  
  Alternatively, one could include the `softmax` or `sigmoid` call in the last layer fo a model:
  
  ```
  model = tf.keras.Sequential([
      # ...
      tf.keras.layers.Conv2D(2, kernel_size=1, padding="same", activation="softmax")
  ])
  
  ```
  
  My question is this: Why could you keep the `softmax` / `sigmoid` in the loss function instead of in the last layer of the model?
  
  Wouldn't putting it in the last layer of the model ensure that the output of the model is constrained, enabling the model to more effectively learn the boundaries of the given input as smaller gradients would be needed for the model to change its mind on a given pixel?
  
  **Tags:** training, image-segmentation, softmax
  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2023-03-17. Score: 3
  
  Mathematically it does not matter at all. The results will be the same. However there is a strong reason to prefer it being in the loss function: numeric stability.
  
  Because the loss function knows more information (ie what the loss is), it can compute the softmax and loss together using some tricks to improve numeric stability.
  
  > **Comments:**
  > **username_0**: Thanks for the explanation. What are these tricks you mention here?  
  > **username_1**: Here is a blog post explaining the basics. If the above answer was helpful please consider accepting it. https://jaykmody.com/blog/stable-softmax/  
  > **username_0**: Many thanks for that link, but I'm afraid I still don't understand. That log post mentioned log(softmax()) for some reason? It did not seem to cover my question  
  > **username_1**: It did cover your question. Log softmax is more numerically stable. You can convert back to regular softmax after. Perhaps you do not have a sufficient mathematical background to understand the answer. But that link directly addresses your question. The trick is to use log softmax during the intermediate computations and then convert it back after. If this still does not make sense to you, make sure you have a good understanding of what numerical stability means for floating point arithmetic, and then the log softmax trick should be fairly easy to understand.  
  > **username_0**: Right, I understand that, but it still doesn't explain why you wouldn't do the (log) softmax in advance in the last layer of the model.  
  > **username_1**: You can. But then you cannot use cross entropy anymore since cross entropy is not expecting a log softmax. You would need to use a log loss instead. That is a fine alternative.  
  
  ---
  
  