Transformer Decoder always predict same token with CLIP embedding

pabloruizp: <p>I’m trying to train a Transformer Decoder to generate textual captions from a CLIP embedding in an autoregressive way. I followed this <a href="https://pytorch.org/tutorials/beginner/translation_transformer.html" rel="noopener nofollow ugc">official Pytorch Tutorial</a> as a base for my project. However, when I train the model the loss decreases, but at inference it always predicts the same token for all the positions in the sequence. I have further debugged my model and during training time it also does the same. It seems that the Positional Encoding or the <code>tgt_mask</code> are not working as expected.</p>
<p>My Transformer Decoder:</p>
<pre><code class="lang-auto">class TextDecoder_TRANSFORMER(nn.Module):
    def __init__(self,latent_dim=512, ff_size=1024, num_layers=4, num_heads=4, dropout=0.1, activation="gelu", **kargs):
        super().__init__()
        
        self.latent_dim = latent_dim
        self.ff_size = ff_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dropout = dropout
        self.activation = activation
                
        self.sequence_pos_encoder = PositionalEncoding(self.latent_dim, self.dropout)
        seqTransDecoderLayer = nn.TransformerDecoderLayer(d_model=self.latent_dim,
                                                          nhead=self.num_heads,
                                                          dim_feedforward=self.ff_size,
                                                          dropout=self.dropout,
                                                          activation=activation)
        self.seqTransDecoder = nn.TransformerDecoder(seqTransDecoderLayer,
                                                     num_layers=self.num_layers)
        
        self.finallayer = nn.Linear(self.latent_dim, 49408)
        
    def decode(self, tgt, memory, tgt_mask):
        return self.seqTransDecoder(self.sequence_pos_encoder(tgt), memory, tgt_mask)
    
    def forward(self, batch):

        # z is the CLIP embedding
        z, mask, padding_mask = batch["z"], batch["text_mask"], batch["text_padding_mask"]
        
        # Text Captions
        captions = batch["clip_text_embedding"]
        captions = self.sequence_pos_encoder(captions)

        output = self.seqTransDecoder(tgt=captions, 
                                      tgt_mask=mask,
                                      tgt_key_padding_mask=padding_mask,
                                      memory=z)
        
        output = self.finallayer(output)
        batch["output_caption"] = output
        return batch
</code></pre>
<p>The Positional Encoding Used:</p>
<pre><code class="lang-auto">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.shape[0], :]
        return self.dropout(x)
</code></pre>
<p>And the code used to encode the textual information feed to the model</p>
<pre><code class="lang-auto">    # Extracted from: https://pytorch.org/tutorials/beginner/translation_transformer.html
    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones((sz, sz), device=self.device)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask
    
    def create_mask(self, tokens):
        tk_seq_len = tokens.shape[0]
        tk_mask = self.generate_square_subsequent_mask(tk_seq_len)
        tk_padding_mask = (tokens == 0).transpose(0, 1)
        return tk_mask, tk_padding_mask

    def encode_text(self, batch):
        captions = batch["clip_text"]

        captions_tensor = []
        for caption in captions:
            caption = clip.tokenize(caption).to(self.device)
            captions_tensor.append(caption)
        
        captions_tensor = torch.cat(captions_tensor, dim=0).to(self.device)
        captions_tensor = captions_tensor.permute(1,0)
        batch["text_tokens"] = captions_tensor

        mask, padding_mask = self.create_mask(captions_tensor[:-1,:])
        batch["text_mask"] = mask
        batch["text_padding_mask"] = padding_mask

        caption_embedding = self.clip_model.token_embedding(captions_tensor[:-1,:]).type(self.clip_model.dtype)  
        caption_embedding.to(self.device)
        batch["clip_text_embedding"] = caption_embedding

        return batch
    

    # function to generate output sequence using greedy algorithm
    def greedy_decode(self,text, memory, max_len):
        # The tokenizer is the SimpleTokenizer from the CLIP repository
        START_IDX = self.tokenizer.encoder['&lt;|startoftext|&gt;']
        EOS_IDX = self.tokenizer.encoder['&lt;|endoftext|&gt;']

        ys = torch.ones(1, 1).fill_(START_IDX).type(torch.long).to(self.device)
        for i in range(max_len-1):
            memory = memory.to(self.device)
            tgt_mask = (self.generate_square_subsequent_mask(ys.size(0))
                        .type(torch.bool)).to(self.device)

            token_embs = self.clip_model.token_embedding(ys)
            
            out = self.textDecoder.decode(token_embs, memory, tgt_mask)
            out = out.transpose(0, 1)
            prob = self.textDecoder.finallayer(out[0,:,:])
            _, next_word = torch.max(prob, dim=1)
            next_word = next_word[-1].item()
            
            ys = torch.cat([ys, torch.ones(1, 1).type(torch.long).fill_(next_word).to(self.device)], dim=0)

            if next_word == EOS_IDX:
                break
        
        return ys
</code></pre>
<p>I’ve been days trying to figure out what is wrong, but I haven’t been able to find it. I would really appreciate your help and wisdom.</p>