Selecting a label smoothing factor for seq2seq NMT with a massive imbalanced vocabulary

Question

I'm training a seq2seq RNN with a vocabulary of 8192 words. This means that the typical categorical cross entropy label smoothing factor suggested in papers like 'Attention is all you need' of $0.1$ would result in true labels with a value around $0.9$ but false labels with a value around $1\cdot10^{-4}$. I hadn't initially consider this an issue at all, but after running training online for around 6000 epochs (with each epoch training on between 1 and 13 samples) the model would almost always only predict tokens such as i, yes, the, that, do, not, know, . when using Beam Search with a beam width of 8 and $\alpha,\beta$ penalties of 0.2.

To me this suggests two potential causes: firstly, the model is underfitting due either too much regularisation (which I doubt due to an LSTM dropout of 0.1 and MHSA dropout of 0.1 with residual and dense skip connections) or due to the model having a low capacity (which I also don't think is true as the MHSA in the encoder and the LSTMs in the encoder and decoder both have a latent dimension of 2048); or secondly, the massive imbalance in the frequencies of different tokens is causing the model to preferentially chose those particular tokens.

If the second case is true would increasing the label smoothing to a higher factor such as $0.9$ help reduce the bias towards the frequent tokens? Or should I employ a 'weighted' label smoothing technique which takes into account the class frequencies?

Alternatively, would a different loss function help mitigate this issue? Binary cross entropy comes to mind, but I don't think that would be a particularly good fit either.

Answer 1

Turns out I was partially on the right track. Label smoothing will help reduce class bias at the cost of introducing some underfitting which can be desirable in complex models which have the capacity to well-and-truly overfit, but not for a simple model like mine.

What I needed was a loss function which was robust to class imbalance with an in-built weight mechanism. The answer here is Focal Loss introduced by this paper with an example keras implementation.

I will however note that for such a massive vocabulary simply using an inversely proportional weighting such as $\textit{weight}_i={1}/{\textit{ClassFreq}_i}$ is not advisable due to the heavy tailed nature (Zipf distribution) of natural language, resulting in the opposite issue where words like i or and are never predicted. Instead, it makes sense to use a much 'softer' weighting distribution using something like the reciprocal of log frequencies normalised by the maximum value which seems to work well with focal loss.