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).


  • 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

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?


1 Answer 1


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:

  • $\begingroup$ 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 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. $\endgroup$
    – Daniel
    Oct 30, 2022 at 21:38
  • $\begingroup$ 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. $\endgroup$
    – Chillston
    Oct 30, 2022 at 23:24
  • $\begingroup$ 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. $\endgroup$
    – Chillston
    Oct 30, 2022 at 23:27
  • $\begingroup$ 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) $\endgroup$
    – Daniel
    Oct 31, 2022 at 1:43
  • 1
    $\begingroup$ 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... $\endgroup$
    – Daniel
    Nov 17, 2022 at 21:36

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .