Answers.Body: <p>As soon as you discretize the selection, i.e. make a hard selection (argmax) instead of a soft selection (softmax), you have a biased gradient. This is because the things you didn't select are not included in the gradient calculation.</p>
  <p>However, it is possible and one way to create such a biased gradient is to apply both techniques. You can use the <strong>argmax method in the forward pass</strong> and the <strong>softmax (with temperature) in the backward pass</strong> .</p>
  <p>Here is what that could look like, independent of framework:</p>
  <pre><code>def differentiable_select(logits, temp=0.6):
      # turn logits into a categorical distribution (soft-selection)
      logits_dist = softmax(logits / temp)
      # get the maximum probability for the forward pass
      hard_selection_idx = logits_dist.argmax(-1)
      hard_selection_mask = one_hot(hard_selection_idx, depth=logits.shape[-1])
      # for the forward pass, return the hard_selection_mask, 
      # for the backward pass, return the logits_dist (soft-selection)
      return logits_dist + stop_gradient(hard_selection_mask - logits_dist)
  
  # generate the selection mask
  select_mask = differentiable_select(logits)
  # select from x (x-shape: [n_items, n_dims])
  x_select = x * selection_mask
  </code></pre>
  <p>The central part is the <code>stop_gradient</code> function. In PyTorch this corresponds to <code>Tensor.detatch()</code> and in tensorflow you'd call <code>tf.stop_gradient</code>.
  That way, in the forward pass, you have the hard selection:</p>
  <pre><code>logits_dist + hard_selection_mask - logits_dist == hard_selection_mask
  </code></pre>
  <p>and in the backward pass, you have the soft selection by using <code>logits_dist</code>. The temperature (<span class="math-container">$\tau$</span>) is a hyperparameter that trades training stability (if <span class="math-container">$\tau$</span> is close to <span class="math-container">$1$</span>) vs. similarity to the forward pass (if <span class="math-container">$\tau$</span> is close to <span class="math-container">$0$</span>). In my experiments, a temperature of <span class="math-container">$\tau = 0.6$</span> worked pretty well.</p>
  <p>This solution is derived from the Straight-Through Gumbel-Softmax (ST-Gumbel-Softmax) function. In the original implementation, the ST-Gumbel-Softmax function includes non-deterministic sampling from the categorical distribution.</p>
  <p>For further reading, you can look at the following sources:</p>
  <ul>
  <li><a href="https://arxiv.org/abs/1308.3432" rel="nofollow noreferrer">[Paper by Bengio et al. on ST-Gumbel-Softmax]</a></li>
  <li><a href="https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax" rel="nofollow noreferrer">[PyTorch Code Example]</a></li>
  <li><a href="https://fabianfuchsml.github.io/gumbel/" rel="nofollow noreferrer">[Blogpost by Fabian Fuchs]</a></li>
  <li><a href="https://arxiv.org/abs/1611.00712" rel="nofollow noreferrer">[Maddison et al.]</a></li>
  </ul>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-30T21:38:58.133Z
Answers.Text: Thank you for your in-depth response @Chillston! I had looked into Gumbel Softmax before posting the question - as you can see I attempted both a temperature-based approach as well as a "hard" softmax. However I found in practice that it did not converge to a solution.  [Here's](https://gist.github.com/JustASquid/aece279ac63cf399744cc0f3f28a21e1) a gist of a Keras implementation of the problem, where the goal is to select the option with the value '5'. It's also confusing to me that all the literature focuses on stochastic networks and Gumbel noise, but this problem is totally deterministic.
Answers.UserDisplayName: Daniel
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-30T23:24:19.803Z
Answers.Text: Ah, right - I must have missed that you wrote that. Anyways, in your case, I could imagine that providing numbers from 0 to 10 as inputs may cause the inability to converge, you could try providing numbers 0.0, 0.1, .., 1.0 and train on selecting 0.5. 
  Beyond that the method should work like you implemented it.
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-30T23:27:29.030Z
Answers.Text: Regarding the sampling from the Gumbel distribution: My current understanding is that this method has an exploration/exploitation problem (as in Reinforcement Learning), so sampling instead of selecting the argmax should improve exploration of the solution space. In your case this shouldn't be an issue though, as you only select once and not cascading multiple consecutive selections. But I'm interested in that as well, so it would be cool if you'd let me know if scaling the dataset values to [0, 1] works.
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-31T01:43:00.670Z
Answers.Text: Scaling the inputs doesn't have an effect (I had tried this before, but left them unscaled for simplicity in the example). In both cases, If you check the network outputs, the logits trend towards being equal (presumably because the backwards pass softmax loss kind of cancels out numbers above and below 0.5)
Answers.UserDisplayName: Daniel
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-31T08:57:01.177Z
Answers.Text: Okay, another idea: Your model does not use permutation symmetry, because you flatten the input directly after passing it to the model. Thus the model has to learn all possible combinations of inputs, when instead you could apply the model to the options separately. You can try leaving out the flatten (L.31) and setting the output neurons of the `Dense` layer in L.35 to 1. Maybe this helps?
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-01T00:53:42.100Z
Answers.Text: I see what you're going for, and I do something similar in my "full" application with 1d convolution layers to apply the same kernel to every item individually. But At the end of the day it *is* order-dependent because the output needs to match up to the input in order for the selection to work
Answers.UserDisplayName: Daniel
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-02T09:52:28.010Z
Answers.Text: I don't fully understand these points. When you use a 1D-Convolution, do you set the kernel-size to 1? Otherwise you'd still have the sequential bias in your model. And from what I know about your project so far, I'd disagree with you on the order-dependency. The model should not be concerned with order if the only important thing is that input and output are ordered similarly - here is an example: If you use an MLP and apply it separately to each item in the sequence, you can still ensure that each output $y_i$ belongs to the corresponding input $x_i$, right? Or maybe I'm missing something
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-17T21:36:09.580Z
Answers.Text: Hey @Chillston - just to follow up on this. It turns out the secret is actually all in the Gumbel noise. If you apply Gumbel noise to the example I posted, it works flawlessly. And it translated to my real application as well! A really interesting result that I don't fully understand...
Answers.UserDisplayName: Daniel
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-18T11:26:41.203Z
Answers.Text: Hey thanks so much for sharing! That's really interesting! You are only selecting once, right? So it's probably not the case that the noise boosts exploration of the solution space.. Are you planning on investigating this further?
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-19T22:38:57.883Z
Answers.Text: what do you mean by only selecting once? It only selects once (per item in the batch, per training step) if that's what you mean?
Answers.UserDisplayName: Daniel
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-11-20T20:08:25.057Z
Answers.Text: Exactly, in a current project of mine, we have cascading selections where the pool of options depends on a prior selection and in that scenario it makes sense to me that the randomness is a benefit. However, it is not intuitive to me at all why it helps in your case :D If you ever find out, I would love to hear what it is!
Answers.UserDisplayName: Chillston
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-10-30T10:15:12.293Z
Answers.LastActivityDate: 2022-10-30T10:26:26.760Z
Answers.LastEditDate: 2022-10-30T10:26:26.760Z
Answers.UserDisplayName: Chillston
Body: <p>Consider the following problem:</p>
  <p>There are <span class="math-container">$N$</span> items and <span class="math-container">$S$</span> slots. Each item is a vector of length <span class="math-container">$D$</span>.
  The goal is to train a neural network to select one item per slot in order to minimize the loss function <span class="math-container">$L$</span>, given any arbitrary list of items. This loss function could also be thought of as a reward function; it's not comparing labels, it just outputs the 'value' of some slots/items configuration, and it's differentiable. (In my particular application, this function is another MLP which can be trained separately; it's an actor-critic setup. But this problem applies to any sort of differentiable loss/reward function).</p>
  <p>Assumptions:</p>
  <ul>
  <li>The same item can be used multiple times in different slots.</li>
  <li>We have an infinite number of training examples available; they don't need to be labelled so they can be generated automatically.</li>
  <li><span class="math-container">$S \times N$</span> is large, such that we can't practically enumerate the value of <span class="math-container">$L$</span> for all possible configurations of items in slots for a particular example. (If we could, then this would turn into a simple supervised learning case).</li>
  </ul>
  <p>Here's an example where <span class="math-container">$N = 8$</span>, <span class="math-container">$S = 4$</span>, <span class="math-container">$D = 1$</span>. The neural network is an MLP which outputs an <span class="math-container">$S \times N$</span> matrix of logits. These are then converted to a binary matrix which is multiplied with the original list of items to get the selected item in every slot, on which the loss can be calculated.
  <a href="https://i.sstatic.net/Gasr6.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/Gasr6.png" alt="enter image description here" /></a></p>
  <p>The problem is what should happen between the neural net output and the binary 'selection' matrix (The ??? bubble in the diagram).</p>
  <ul>
  <li>It's natural to do an argmax/one-hot operation on the logits (i.e. push the max logit to 1 and the rest to zero). But this operation isn't differentiable, so the MLP can't be trained.</li>
  <li>You could do a softmax instead, which has good gradients, but then the problem is broken because the selected items are weighted combinations of items in the original list, which isn't allowed.</li>
  <li>You can also do softmax with &quot;temperature&quot;, e.g. <code>softmax(logits * 10e5)</code>, to get a more argmax-like behavior in a differentiable function. But in practice, the gradients are low quality and the network doesn't train properly.</li>
  <li>You can try argmax for the forward pass, but softmax for the backwards pass - this kind of works, but again in practice the neural network does not converge to a good solution.</li>
  </ul>
  <p>Is there any kind of neural network architecture or training regimen that can solve or sidestep this problem? Or is this application fundamentally unsuited for back-propagation?</p>
  
Tags: <deep-learning><deep-rl><architecture>
Title: Propagating gradients through an "Item Selector" network
UserDisplayName: Daniel
ThreadText: # Propagating gradients through an "Item Selector" network
  
  (Asked by: username_0 on 2022-10-30)
  
  Consider the following problem:
  
  There are $N$ items and $S$ slots. Each item is a vector of length $D$.
  The goal is to train a neural network to select one item per slot in order to minimize the loss function $L$, given any arbitrary list of items. This loss function could also be thought of as a reward function; it's not comparing labels, it just outputs the 'value' of some slots/items configuration, and it's differentiable. (In my particular application, this function is another MLP which can be trained separately; it's an actor-critic setup. But this problem applies to any sort of differentiable loss/reward function).
  
  Assumptions:
  
  * The same item can be used multiple times in different slots.
  * We have an infinite number of training examples available; they don't need to be labelled so they can be generated automatically.
  * $S \times N$ is large, such that we can't practically enumerate the value of $L$ for all possible configurations of items in slots for a particular example. (If we could, then this would turn into a simple supervised learning case).
  
  Here's an example where $N = 8$, $S = 4$, $D = 1$. The neural network is an MLP which outputs an $S \times N$ matrix of logits. These are then converted to a binary matrix which is multiplied with the original list of items to get the selected item in every slot, on which the loss can be calculated.
  [![enter image description here](https://i.sstatic.net/Gasr6.png)](https://i.sstatic.net/Gasr6.png)
  
  The problem is what should happen between the neural net output and the binary 'selection' matrix (The ??? bubble in the diagram).
  
  * It's natural to do an argmax/one-hot operation on the logits (i.e. push the max logit to 1 and the rest to zero). But this operation isn't differentiable, so the MLP can't be trained.
  * You could do a softmax instead, which has good gradients, but then the problem is broken because the selected items are weighted combinations of items in the original list, which isn't allowed.
  * You can also do softmax with "temperature", e.g. `softmax(logits * 10e5)`, to get a more argmax-like behavior in a differentiable function. But in practice, the gradients are low quality and the network doesn't train properly.
  * You can try argmax for the forward pass, but softmax for the backwards pass - this kind of works, but again in practice the neural network does not converge to a good solution.
  
  Is there any kind of neural network architecture or training regimen that can solve or sidestep this problem? Or is this application fundamentally unsuited for back-propagation?
  
  **Tags:** deep-learning, deep-rl, architecture
  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2022-10-30. Score: 2
  
  As soon as you discretize the selection, i.e. make a hard selection (argmax) instead of a soft selection (softmax), you have a biased gradient. This is because the things you didn't select are not included in the gradient calculation.
  
  However, it is possible and one way to create such a biased gradient is to apply both techniques. You can use the **argmax method in the forward pass** and the **softmax (with temperature) in the backward pass** .
  
  Here is what that could look like, independent of framework:
  
  ```
  def differentiable_select(logits, temp=0.6):
      # turn logits into a categorical distribution (soft-selection)
      logits_dist = softmax(logits / temp)
      # get the maximum probability for the forward pass
      hard_selection_idx = logits_dist.argmax(-1)
      hard_selection_mask = one_hot(hard_selection_idx, depth=logits.shape[-1])
      # for the forward pass, return the hard_selection_mask, 
      # for the backward pass, return the logits_dist (soft-selection)
      return logits_dist + stop_gradient(hard_selection_mask - logits_dist)
  
  # generate the selection mask
  select_mask = differentiable_select(logits)
  # select from x (x-shape: [n_items, n_dims])
  x_select = x * selection_mask
  
  ```
  
  The central part is the `stop_gradient` function. In PyTorch this corresponds to `Tensor.detatch()` and in tensorflow you'd call `tf.stop_gradient`.
  That way, in the forward pass, you have the hard selection:
  
  ```
  logits_dist + hard_selection_mask - logits_dist == hard_selection_mask
  
  ```
  
  and in the backward pass, you have the soft selection by using `logits_dist`. The temperature ($\tau$) is a hyperparameter that trades training stability (if $\tau$ is close to $1$) vs. similarity to the forward pass (if $\tau$ is close to $0$). In my experiments, a temperature of $\tau = 0.6$ worked pretty well.
  
  This solution is derived from the Straight-Through Gumbel-Softmax (ST-Gumbel-Softmax) function. In the original implementation, the ST-Gumbel-Softmax function includes non-deterministic sampling from the categorical distribution.
  
  For further reading, you can look at the following sources:
  
  * [[Paper by Bengio et al. on ST-Gumbel-Softmax]](https://arxiv.org/abs/1308.3432)
  * [[PyTorch Code Example]](https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax)
  * [[Blogpost by Fabian Fuchs]](https://fabianfuchsml.github.io/gumbel/)
  * [[Maddison et al.]](https://arxiv.org/abs/1611.00712)
  
  > **Comments:**
  > **username_0**: Thank you for your in-depth response @username_1! I had looked into Gumbel Softmax before posting the question - as you can see I attempted both a temperature-based approach as well as a "hard" softmax. However I found in practice that it did not converge to a solution.  [Here's](https://gist.github.com/JustASquid/aece279ac63cf399744cc0f3f28a21e1) a gist of a Keras implementation of the problem, where the goal is to select the option with the value '5'. It's also confusing to me that all the literature focuses on stochastic networks and Gumbel noise, but this problem is totally deterministic.  
  > **username_1**: Ah, right - I must have missed that you wrote that. Anyways, in your case, I could imagine that providing numbers from 0 to 10 as inputs may cause the inability to converge, you could try providing numbers 0.0, 0.1, .., 1.0 and train on selecting 0.5.  Beyond that the method should work like you implemented it.  
  > **username_1**: Regarding the sampling from the Gumbel distribution: My current understanding is that this method has an exploration/exploitation problem (as in Reinforcement Learning), so sampling instead of selecting the argmax should improve exploration of the solution space. In your case this shouldn't be an issue though, as you only select once and not cascading multiple consecutive selections. But I'm interested in that as well, so it would be cool if you'd let me know if scaling the dataset values to [0, 1] works.  
  > **username_0**: Scaling the inputs doesn't have an effect (I had tried this before, but left them unscaled for simplicity in the example). In both cases, If you check the network outputs, the logits trend towards being equal (presumably because the backwards pass softmax loss kind of cancels out numbers above and below 0.5)  
  > **username_1**: Okay, another idea: Your model does not use permutation symmetry, because you flatten the input directly after passing it to the model. Thus the model has to learn all possible combinations of inputs, when instead you could apply the model to the options separately. You can try leaving out the flatten (L.31) and setting the output neurons of the `Dense` layer in L.35 to 1. Maybe this helps?  
  > **username_0**: I see what you're going for, and I do something similar in my "full" application with 1d convolution layers to apply the same kernel to every item individually. But At the end of the day it *is* order-dependent because the output needs to match up to the input in order for the selection to work  
  > **username_1**: I don't fully understand these points. When you use a 1D-Convolution, do you set the kernel-size to 1? Otherwise you'd still have the sequential bias in your model. And from what I know about your project so far, I'd disagree with you on the order-dependency. The model should not be concerned with order if the only important thing is that input and output are ordered similarly - here is an example: If you use an MLP and apply it separately to each item in the sequence, you can still ensure that each output $y_i$ belongs to the corresponding input $x_i$, right? Or maybe I'm missing something  
  > **username_0**: Hey @username_1 - just to follow up on this. It turns out the secret is actually all in the Gumbel noise. If you apply Gumbel noise to the example I posted, it works flawlessly. And it translated to my real application as well! A really interesting result that I don't fully understand...  
  > **username_1**: Hey thanks so much for sharing! That's really interesting! You are only selecting once, right? So it's probably not the case that the noise boosts exploration of the solution space.. Are you planning on investigating this further?  
  > **username_0**: what do you mean by only selecting once? It only selects once (per item in the batch, per training step) if that's what you mean?  
  > **username_1**: Exactly, in a current project of mine, we have cascading selections where the pool of options depends on a prior selection and in that scenario it makes sense to me that the randomness is a benefit. However, it is not intuitive to me at all why it helps in your case :D If you ever find out, I would love to hear what it is!  
  
  ---
  
  