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, but in the backward pass replace the argmax with a softmax (with temperature).

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.

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

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