Make nn.Transformer work for Text Generation

username: AfonsoSalgadoSousa
cooked: <p>Hi. I am trying to make a Transformer work for paraphrase generation but the generations are not useful (the same everytime, full of BOS tokens or “?” tokens).<br>
I followed <a href="https://pytorch.org/tutorials/beginner/translation_transformer.html" rel="noopener nofollow ugc">this tutorial</a> for reference. My implementation is embedded into a framework which requires an Encoder and a Decoder:</p>
<p>The encoder is like this:</p>
<pre><code class="lang-auto">class TransformerEncoder(nn.Module):
    def __init__(
        self,
        vocab_size,
        pad_token_id=None,
        embedding_size=256,
        num_heads=8,
        num_layers=3,
        ffnn_size=512,
        dropout=0.1,
    ):
        super(TransformerEncoder, self).__init__()
        self.vocab_size = vocab_size
        self.pad_token_id = pad_token_id

        self.embedding_size = embedding_size
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.ffnn_size = ffnn_size

        self.embed_tokens = TokenEmbedding(vocab_size, embedding_size)
        self.embed_positions = PositionalEmbedding(embedding_size, dropout=dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            embedding_size,
            num_heads,
            ffnn_size,
            dropout,
        )
        encoder_norm = nn.LayerNorm(embedding_size)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers, encoder_norm)

    def forward(
        self,
        input_ids,
    ):

        # seq_len = input_ids.shape[1]
        # device = next(self.parameters()).device

        embedded_tokens = self.embed_positions(self.embed_tokens(input_ids))
        # B x T x C -&gt; T x B x C
        embedded_tokens = embedded_tokens.transpose(0, 1)

        memory = self.encoder(embedded_tokens)

        return (memory,)
</code></pre>
<p>The decoder is like this:</p>
<pre><code class="lang-auto">class TransformerDecoder(nn.Module):
    def __init__(
        self,
        vocab_size,
        pad_token_id=None,
        embedding_size=256,
        num_heads=8,
        num_layers=3,
        ffnn_size=512,
        dropout=0.1,
    ):

        super(TransformerDecoder, self).__init__()
        self.vocab_size = vocab_size
        self.pad_token_id = pad_token_id

        self.embedding_size = embedding_size
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.ffnn_size = ffnn_size

        self.dropout_module = nn.Dropout(p=dropout)

        self.embed_tokens = TokenEmbedding(vocab_size, embedding_size)
        self.embed_positions = PositionalEmbedding(embedding_size, dropout=dropout)

        decoder_layer = nn.TransformerDecoderLayer(
            embedding_size, num_heads, ffnn_size, dropout
        )
        decoder_norm = nn.LayerNorm(embedding_size)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers, decoder_norm)
        self.fc_out = nn.Linear(embedding_size, vocab_size)

    def forward(
        self,
        input_ids,
        encoder_out,
    ):
        seq_len = input_ids.shape[1]

        device = next(self.parameters()).device
        mask = generate_square_subsequent_mask(seq_len).to(device)

        embedded_tokens = self.embed_positions(self.embed_tokens(input_ids))

        # B x T x C -&gt; T x B x C
        embedded_tokens = embedded_tokens.transpose(0, 1)

        output = self.decoder(embedded_tokens, encoder_out[0], tgt_mask=mask)

        # T x B x C -&gt; B x T x C
        output = output.transpose(1, 0)

        return (self.fc_out(output),)
</code></pre>
<p>TokenEmbedding and PositionalEmbedding are as in the tutorial.<br>
The main model just invokes encoder and decoder like:</p>
<pre><code class="lang-auto">        encoder_outputs = self.encoder(input_ids=input_ids, **kwargs)

        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            encoder_out=encoder_outputs,
            **kwargs,
        )
</code></pre>
<p>The loss is going down, but the generations are real bad. Following is an example of the generations:<br>
<strong>Source:</strong> &lt; s &gt; Can I jailbreak iOS 10 ? &lt; /s &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt;<br>
<strong>Preds:</strong> &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt; &lt; s &gt;<br>
<strong>Target:</strong> &lt; s &gt; Can you jailbreak iOS 10 ? &lt; /s &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt; &lt; pad &gt;</p>
<p>As you can see, the predictions in this case are only BOS tokens. The output of the decoder on each decoder step is always almost the same for every iteration. The model does not seem to be learning. I have tried learning rates from 0.1 to 1e-4.</p>
<p>Do you have an intuition on what might be wrong? Sorry for the question not being self-contained. Thanks in advance for any help you can provide.</p>