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