Gradients Exploding with Custom Gradient and 2 layer MLP

Question

I am training a 2 layer MLP on an off-policy learning-to-rank task, where the input is a list of documents against a query with a feature vector for each query-document pair, i.e. if there are M documents against a query, I have M feature vectors. I use a 2 layer MLP, with ReLU activations to map the feature vectors to a query-document ranking score which is used to rank the documents later. The objective is similar to RL objective, where the goal is to maximize the average ranking metric given the documents and labels, wherein average is taken over multiple (stochastic) sample of rankings, and computing the ranking metric for each ranking, and finally averaging it.

To train the model, I am using a custom-gradient implementation (as sketched in this paper: https://harrieo.github.io/files/2021-sigir-plrank.pdf), where the method gives me a weight for each document, and the gradient is finally defined as:

$\delta R(q) = \sum_{d} w_d \delta m(d) $, where $R$ is the reward for the query, $w$ is the weight returned by the custom gradient method, and $\delta m(d)$ is the gradient of the model weights from the autograd.

The training for this method is highly unstable, and the loss quickly shoots over 1e+6 value, which returns mostly garbage metric. Some observations:

1. I have scaled each feature column in the input to [0,1] range.
2. The training is stable as long as the sum of the weight $\sum_{d} \lambda_d \leq 1$, as soon as the sum of weight shoots over 1, the training becomes unstable.
3. The document-query score $m(d)$ is unbounded because of ReLU, and if I switch the activations to Tanh/sigmoid, the training stabilizes, but the model underfits.
4. Tried different network initialization methods, but all fail.

Any suggestions to stabilize the training?

Answer 1

Try the following:

• Gradient clipping, either by local or global norm. A typical clipping value is 1, but expect slower learning.
• A unit-norm contraint (or similar) on the weights.
• Normalize the rewards or even better the returns to have a suitably (small) scale.
• Try using a different activation like ReLU6 (like ReLU but bounded to a maximum of six), swish, or GELU.
• If you use activations from the ReLU family pair them with He-initialization of weights.
• Introduce batch normalization in your network.
• Check for numerical instabilities in the custom gradient computation, if so: try a more numerically stable but equivalent implementation, or add some small epsilon.