Как заставить модель переобучиться

Недавно я изучал модель Transformer, используя реализацию на Pytorch, но моя модель не сходилась. Я задумался, была ли проблема в моем коде или в чем-то еще. Поэтому я подумал, что если я “упрощу” задачу для обучения, то ее будет легче обучать, и создал предложение и повторил его несколько раз (чтобы заставить модель переобучиться), но моя модель все равно не сходилась. Так что мой вопрос заключается в том, есть ли какая-то ошибка в моей реализации/цикле обучения или переобучение не так тривиально?

Модель:

class PositionalEncoding(nn.Module):
   def __init__(self, d_model, dropout=0.1, max_len=5000):
       super(PositionalEncoding, self).__init__()
       self.dropout = nn.Dropout(p=dropout)

       pe = torch.zeros(max_len, d_model)
       position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
       div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
       pe[:, 0::2] = torch.sin(position * div_term)
       pe[:, 1::2] = torch.cos(position * div_term)
       pe = pe.unsqueeze(0).transpose(0, 1)
       self.register_buffer('pe', pe)

   def forward(self, x):
       x = x + self.pe[:x.size(0), :]
       return self.dropout(x)

class TransformerModel(nn.Module):
   def __init__(self, d_model, max_seq_len, enc_vocab_size, dec_vocab_size):
       super(TransformerModel, self).__init__()
       self.transformer = nn.Transformer(d_model=d_model, nhead=8, num_encoder_layers=4, num_decoder_layers=8, dim_feedforward=2048, batch_first=True)
       self.d_model = d_model
       self.enc_vocab_size = enc_vocab_size
       self.dec_vocab_size = dec_vocab_size

       self.encoder_emb = nn.Embedding(enc_vocab_size, d_model)
       self.decoder_emb = nn.Embedding(dec_vocab_size, d_model)

       self.pe = PositionalEncoding(d_model)
       self.ol = nn.Linear(d_model, dec_vocab_size)

   def forward(self, enc_input, enc_mask, dec_input, dec_mask, causal_mask):
       enc_input = self.encoder_emb(enc_input)
       enc_input = self.pe(enc_input)

       dec_input = self.decoder_emb(dec_input)
       dec_input = self.pe(dec_input)

       out = self.transformer(enc_input, dec_input, tgt_mask=causal_mask, src_key_padding_mask=enc_mask, tgt_key_padding_mask=dec_mask)

       return self.ol(out)

Цикл обучения:

en_vocab_size = len(tokenizer_en.get_vocab())
pt_vocab_size = len(tokenizer_pt.get_vocab())
seq_len = 40

model = TransformerModel(d_model=256, max_seq_len=seq_len, enc_vocab_size=en_vocab_size, dec_vocab_size=pt_vocab_size)
model.to(device)
causal_mask = torch.triu(torch.ones((seq_len, seq_len)), diagonal=1) * -10**9
causal_mask = causal_mask.to(device)

opt = torch.optim.Adam(model.parameters(), lr=10e-3)
loss_fn = nn.CrossEntropyLoss()

model.train()

def get_input_and_mask(sequence_batch, tokenizer, device):
   sequence_batch = tokenizer.encode_batch(sequence_batch)
   mask  = torch.tensor([e.attention_mask for e in sequence_batch], dtype=bool, device=device)
   sequence_batch = torch.tensor([e.ids for e in sequence_batch], device=device)
   return sequence_batch, mask

num_batches = len(train_data)
print(f"num_batches = {num_batches}")

for epoch in range(100):
   print(f" EPOCH {epoch} ".center(80, '-'))
   for iteration, (enc, dec) in enumerate(train_data):
       enc_input, enc_mask = get_input_and_mask(enc, tokenizer_en, device)
       dec_input, dec_mask = get_input_and_mask(dec, tokenizer_pt, device)

       y = model(enc_input, enc_mask, dec_input, dec_mask, causal_mask)

       opt.zero_grad()
       loss = loss_fn(y.view(-1, pt_vocab_size), dec_input.view(-1))
       loss.backward()

Набор данных:

Input> "Porque semeiam vento, colherão tempestade; não haverá seara, pois o talo não produzirá cereal; o pouco que der, os estrangeiros o devorarão."
Target> "For they have sown the wind, and they shall reap the whirlwind: it hath no stalk: the bud shall yield no meal: if so be it yield, the strangers shall swallow it up."

Я повторил это предложение 20 тысяч раз.