Transformer model not overfitting, but good performance on datasets with patterns

Crossword: <p>I have written a transformer model in Pytorch to try to predict crossword answers from clues. I am training it on a dataset of about 200,000 answer-clue pairs. I have tried Adam and SGD optimizers, and am using cross-entropy loss as the cost function.</p>
<p>The model trains well on an artificial dataset with about 100 examples and a simple pattern: the source array is randomly selected from all integers between the minimum and the maximum of the vocabulary tokens, and each element y of the target array is either equal to the corresponding element x, or equal to x+20, or equal to x-20, with some padding and arbitrarily chosen numbers for EOS and SOS tokens. The loss quickly converges to a low value (0.008) from its intial high value, and would probably keep decreasing until approaching zero if the learning rate were lowered.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/a/ba0bea5238b905b56e5207fc58ef5acb4ec81299.png" data-download-href="https://discuss.pytorch.org/uploads/default/ba0bea5238b905b56e5207fc58ef5acb4ec81299" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba0bea5238b905b56e5207fc58ef5acb4ec81299_2_690x379.png" alt="image" data-base62-sha1="qxQkX5Ak7vTwxPVf9Hntub9bE0x" width="690" height="379" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba0bea5238b905b56e5207fc58ef5acb4ec81299_2_690x379.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba0bea5238b905b56e5207fc58ef5acb4ec81299_2_1035x568.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba0bea5238b905b56e5207fc58ef5acb4ec81299_2_1380x758.png 2x" data-dominant-color="FAFBFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">3406×1873 221 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, it does not train well on the actual dataset. The loss decreases by at most 90%, and then seems to get stuck around a value where the model predicts primarily the first (SOS) token. Its performance on a very small subset of the dataset (say, five items), where we would expect it to overfit, is similarly poor, with the loss function not going below 0.6 even after thousands of epochs:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e190c3363074a0e0e8a7d216c5de936ed437acb4.png" data-download-href="https://discuss.pytorch.org/uploads/default/e190c3363074a0e0e8a7d216c5de936ed437acb4" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e190c3363074a0e0e8a7d216c5de936ed437acb4_2_690x356.png" alt="image" data-base62-sha1="wbrA9IPUxEQTI9rqq4eFBsM1cdC" width="690" height="356" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e190c3363074a0e0e8a7d216c5de936ed437acb4_2_690x356.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e190c3363074a0e0e8a7d216c5de936ed437acb4_2_1035x534.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e190c3363074a0e0e8a7d216c5de936ed437acb4_2_1380x712.png 2x" data-dominant-color="F7F9FB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">3571×1846 231 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In conjunction with this, the model tends to disproportionately predict a single value, such as “the,” for all inputs.</p>
<p>Does anyone have any ideas for fixing this?</p>