Propagating gradients through an "Item Selector" network

Question

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.

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?

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:

• [Paper by Bengio et al. on ST-Gumbel-Softmax]
• [PyTorch Code Example]
• [Blogpost by Fabian Fuchs]
• [Maddison et al.]