Minimal working example or tutorial showing how to use Pytorch's nn.TransformerDecoder for batch text generation in training and inference modes?

Question

I want to solve a sequence-to-sequence text generation task (e.g. question answering, language translation, etc.).

For the purposes of this question, you may assume that I already have the input part already handled. (I already have a tensor of dimensions batch_size x num_input_tokens x input_dim representing the input sequences. Also, all input sequences in my problem are of the same length, so no masking is required on the input side of things).

Now, I want to generate the output sequences using nn.TransformerDecoder. I'm aware of Pytorch's official tutorial SEQUENCE-TO-SEQUENCE MODELING WITH NN.TRANSFORMER AND TORCHTEXT. Unfortunately, the official tutorial doesn't meet my needs, for the following reasons:

• nn.TransformerDecoder is not used in the example.
• The example is about language modeling, not text generation. There is no forward loop that generates text word by word.

I've searched around the web and I've found a few things, but nothing like a simple and minimal working example that directly applies to my problem setting. Concretely, on the output side of things I need the following:

• I want to generate output sequences in batch. I've found codes on GitHub where people appear to be doing text generation, but they do it for a single sequence at a time, not a batch of multiple sequences.
• The output sequences may have different lengths.
• I want to train my model with the teacher-forcing strategy and batches of multiple sequences. Given that in training I know the lengths of the sequences in advance, you may assume that I already have my batches padded with zeroes. However, I still need to figure out how to implement the forward function of my model, with a generation loop that uses nn.TransformerDecoder. Basically, I need to figure out how to iterate word-wise over my batch of output sequences, masking out the future words in each step (so that the model doesn't cheat by trivially predicting the next words).
• Then, I need a similar forward function for inference mode. I need to figure out how to implement the generation loop to do basically the same as in training mode, except that instead of teacher-forcing I want to implement greedy search (i.e. use the tokens with highest predicted probability at iteration i as the next input for iteration i+1).

I already know how to do all this using LSTMs. Below you can see the forward function of a model that I implemented in the past to do exactly what I just said with an LSTM. The same forward function is used for both training and inference, depending on the value of the variable 'mode':

  def forward(
      self,
      image_local_features,
      question_vectors,
      answers=None,
      max_answer_length=None,
      mode='train',
  ):
    if mode == 'train':
      batch_size, max_answer_length = answers.shape
      assert answers is not None
    else:
      batch_size = image_local_features.size(0)
      assert max_answer_length is not None
    
    y = self.embedding_table(self.start_idx).expand(batch_size, -1)
    o = torch.zeros(batch_size, self.hidden_size).to(DEVICE)
    h = self.W_h(question_vectors)
    c = self.W_c(question_vectors)

    if mode == 'train':
      answer_embeddings = self.embedding_table(answers.permute(1,0))
      assert answer_embeddings.shape == (max_answer_length, batch_size, self.embed_size)

    output = []

    for t in range(max_answer_length):
      y_bar = torch.cat((y,o),1)
      assert y_bar.shape == (batch_size, self.embed_size + self.hidden_size)
      assert h.shape == (batch_size, self.hidden_size)
      assert c.shape == (batch_size, self.hidden_size)
      h, c = self.lstm_cell(y_bar, (h, c))
      e = (self.W_attn(image_local_features) * h.unsqueeze(1)).sum(-1)
      att = torch.softmax(e,-1)
      a = (image_local_features * att.unsqueeze(2)).sum(1)
      assert a.shape == (batch_size, self.image_local_feat_size)
      u = torch.cat((a,h),1)
      assert u.shape == (batch_size, self.hidden_size + self.image_local_feat_size)
      v = self.W_u(u)
      o = self.dropout(torch.tanh(v))
      assert o.shape == (batch_size, self.hidden_size)
      output.append(self.W_vocab(o))
      if mode == 'train':
        y = answer_embeddings[t] # teacher-forcing
      else:
        y = self.embedding_table(torch.argmax(output[t], 1)) # greedy search
      assert y.shape == (batch_size, self.embed_size)

    output = torch.stack(output, 1)
    assert output.shape == (batch_size, max_answer_length, self.vocab_size)
    return output

Another way to phrase my question would be: how can I reimplement what I did with LSTMs using nn.TransformerDecoder instead?

Any minimal working / hello world example that shows how to do batch training and batch inference with nn.TransformerDecoder for text generation will be very appreciated.

---

Note: alternatively, if there is a straightforward way of accomplishing the same with an out-of-the-box solution from hugginface, that would be awesome too.

Answer 1

After a Googling around, I think this tutorial may suit your needs.

However, it seems you have a misconception about the Transformer decoder: in training mode there is no iteration at all. While LSTM-based decoders are autoregressive by nature, Transformers are not. Instead, all predictions are generated at once based on the real target tokens (i.e. teacher forcing). To train a Transformer decoder to later be used autoregressively, we use the self-attention masks, to ensure that each prediction only depends on the previous tokens, despite having access to all tokens. You can have a look at the Annotated Transformer tutorial in its Training loop section to see how they do it.

Another difference between LSTMs and Transformers is positional encodings, which are used by Transformers to be able to know the position of each token.

Regarding inference time, the easiest approach is to implement greedy decoding (e.g. this), where at each timestep you simply take the most probable token. This decoding strategy, however, will probably give poor results (e.g. the typical token repetitions). A better option is beam search, where at each timestep you keep the most probable K partially decoded sequences, although it is more complex to implement and I have not found any implementation online meant for nn.TransformerDecoder; maybe you can have a look at OpenNMT's implementation.