Answers.Body: <p>Though I may be answering this question quite late, I would like to address your inquiry (because it's possible that other individuals may also have similar queries. :D)</p>
  <p>You can think of the last 3 lines of the function <code>predict</code> as implementing a <a href="https://en.wikipedia.org/wiki/Softmax_function" rel="nofollow noreferrer">softmax</a> layer.
  In this case, the function returns the log probability rather than the probability.</p>
  <p><a href="https://i.sstatic.net/VGlCt.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/VGlCt.png" alt="enter image description here" /></a></p>
  <p>To be specific, <code>logits</code> stores <span class="math-container">$z_i=\ln{e^{z_i}}$</span> and <code>logsumexp(...)</code> computes the <span class="math-container">$\ln\sum_{j=1}^{K}{e^{z_j}}$</span>. The returned value is thus <span class="math-container">$\ln{z_i} - \ln\sum_{j=1}^{K}{e^{z_j}} = \ln{ \frac{e^{z_i}}{\sum_{j=1}^{K}{e^{z_j}}}}$</span>.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-12-29T11:02:34.443Z
Answers.LastActivityDate: 2023-12-29T11:02:34.443Z
Answers.UserDisplayName: present42
Answers.Body: <p>It's apparently for numerical stability. From the Wikipedia page for <a href="https://en.wikipedia.org/wiki/LogSumExp" rel="nofollow noreferrer">LogSumExp</a>:</p>
  <blockquote>
  <p>A common purpose of using log-domain computations is to increase accuracy and avoid underflow and overflow problems when very small or very large numbers are represented directly (i.e. in a linear domain) using limited-precision floating point numbers.</p>
  </blockquote>
  <p>And <a href="https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning">this answer</a> from <code>stats.stackexchange.com</code>:</p>
  <blockquote>
  <p>This is a simple trick to improve the numerical stability. As you probably know, exponential function grows very fast, and so does the magnitude of any numerical errors. This trick is based on [...]</p>
  </blockquote>
  <p>That said, I'm not sure why it's required in your example because, unless I'm missing something, there is no exponential function which could introduce the aforementioned numerical instability.</p>
  <p>I guess it's possible that it's a mistake in the JAX docs, but it seems more likely that I just don't understand it. Hopefully someone else will be able to comment/answer to provide an explanation. I figured I'd post this half-answer since there are no others answers here.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-08-09T21:03:40.510Z
Answers.LastActivityDate: 2022-08-09T21:03:40.510Z
Answers.UserDisplayName: joe
Body: <p>I'm trying to understand <a href="https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html" rel="nofollow noreferrer">this tutorial for Jax</a>.</p>
  <p>Here's an excerpt. It's for a neural net that is designed to classify MNIST images:</p>
  <pre><code>from jax.scipy.special import logsumexp
  
  def relu(x):
    return jnp.maximum(0, x)
  
  def predict(params, image):
    # per-example predictions
    activations = image
    for w, b in params[:-1]:
      outputs = jnp.dot(w, activations) + b
      activations = relu(outputs)
    
    final_w, final_b = params[-1]
    logits = jnp.dot(final_w, activations) + final_b
    return logits - logsumexp(logits)
  </code></pre>
  <p>I don't understand why they would subtract a constant value from all the final predictions, given that the only thing that matters is their relative values.</p>
  <p>For clarity, here's the loss function:</p>
  <pre><code>def loss(params, images, targets):
    preds = batched_predict(params, images)
    return -jnp.mean(preds * targets)
  </code></pre>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2022-06-29T22:44:22.477Z
Comments.Text: Can you please provide more context? What task are they trying to solve? Which loss function? Which dataset? Ideally, questions here should not require code. Code can be provided as an additional source of info, but note that programming issues are off-topic. I don't think your question is a programming issue, but, like I just said, it lacks details.
Comments.UserDisplayName: nbro
Tags: <neural-networks><machine-learning><loss><mnist>
Title: Why do we subtract logsumexp from the outputs of this neural network?
UserDisplayName: Foobar
ThreadText: # Why do we subtract logsumexp from the outputs of this neural network?
  
  (Asked by: username_0 on 2022-06-25)
  
  I'm trying to understand this tutorial for Jax.
  Here's an excerpt. It's for a neural net that is designed to classify MNIST images:
  from jax.scipy.special import logsumexp
  
  def relu(x):
    return jnp.maximum(0, x)
  
  def predict(params, image):
    # per-example predictions
    activations = image
    for w, b in params[:-1]:
      outputs = jnp.dot(w, activations) + b
      activations = relu(outputs)
    
    final_w, final_b = params[-1]
    logits = jnp.dot(final_w, activations) + final_b
    return logits - logsumexp(logits)
  
  I don't understand why they would subtract a constant value from all the final predictions, given that the only thing that matters is their relative values.
  For clarity, here's the loss function:
  def loss(params, images, targets):
    preds = batched_predict(params, images)
    return -jnp.mean(preds * targets)
  
  **Tags:** neural-networks, machine-learning, loss, mnist
  > **Comments:**
  > **username_1**: Can you please provide more context? What task are they trying to solve? Which loss function? Which dataset? Ideally, questions here should not require code. Code can be provided as an additional source of info, but note that programming issues are off-topic. I don't think your question is a programming issue, but, like I just said, it lacks details.  
  
  ---
  
  ## 2 answers
  
  ### Answer by username_2 on 2023-12-29. Score: 2 | Accepted answer
  
  Though I may be answering this question quite late, I would like to address your inquiry (because it's possible that other individuals may also have similar queries. :D)
  
  You can think of the last 3 lines of the function `predict` as implementing a [softmax](https://en.wikipedia.org/wiki/Softmax_function) layer.
  In this case, the function returns the log probability rather than the probability.
  
  [![enter image description here](https://i.sstatic.net/VGlCt.png)](https://i.sstatic.net/VGlCt.png)
  
  To be specific, `logits` stores $z\_i=\ln{e^{z\_i}}$ and `logsumexp(...)` computes the $\ln\sum\_{j=1}^{K}{e^{z\_j}}$. The returned value is thus $\ln{z\_i} - \ln\sum\_{j=1}^{K}{e^{z\_j}} = \ln{ \frac{e^{z\_i}}{\sum\_{j=1}^{K}{e^{z\_j}}}}$.
  
  
  ---
  
  ### Answer by username_3 on 2022-08-09. Score: 1
  
  It's apparently for numerical stability. From the Wikipedia page for [LogSumExp](https://en.wikipedia.org/wiki/LogSumExp):
  
  > A common purpose of using log-domain computations is to increase accuracy and avoid underflow and overflow problems when very small or very large numbers are represented directly (i.e. in a linear domain) using limited-precision floating point numbers.
  
  And [this answer](https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning) from `stats.stackexchange.com`:
  
  > This is a simple trick to improve the numerical stability. As you probably know, exponential function grows very fast, and so does the magnitude of any numerical errors. This trick is based on [...]
  
  That said, I'm not sure why it's required in your example because, unless I'm missing something, there is no exponential function which could introduce the aforementioned numerical instability.
  
  I guess it's possible that it's a mistake in the JAX docs, but it seems more likely that I just don't understand it. Hopefully someone else will be able to comment/answer to provide an explanation. I figured I'd post this half-answer since there are no others answers here.
  
  
  ---
  
  