Hello.
I’m training a transformer to replicate “attention is all you need” for german to english translation and I’ve found that training starts to diverge pretty quickly. I wanted to start with a baseline, so I was following the pytorch transformer tutorial. I.e. the architecture is 2 layers of an encoder, and then the decoder is just a linear layer and a softmax.
Here is an abridged version of the architecture:
class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

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

    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x: Tensor, shape [seq_len, batch_size, embedding_dim]
        """
        # Logic above assumes [seq_len, batch_size, embedding_dim]
        # Swap axes for batch_size first, then swap back. 
        x = torch.swapaxes(x, 1, 0)
        x = x + self.pe[:x.size(0)]
        x = torch.swapaxes(x, 0, 1)        
        return self.dropout(x)  
    
    
class BaselineTransformer(nn.Module): 
    def __init__(self, 
                 model_dim=512, 
                 nhead=8, 
                 num_encoder_layers=6,
                 num_decoder_layers=6, 
                 output_vocab_size=10000, 
                 padding_idx=0, 
                 dropout=0.0): 
        super(BaselineTransformer, self).__init__()
        self.model_dimension = model_dim
        self.nhead = nhead
        self.num_encoder_layers = num_encoder_layers
        self.num_decoder_layers = num_decoder_layers
        self.output_vocab_size = output_vocab_size
        
        pretrained_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-de-en")
        
        self.encoder_input_embeddings = pretrained_model.get_encoder().embed_tokens.to(device)
        
        self.pos_encoder = PositionalEncoding(model_dim, dropout=dropout)
        encoder_layers = TransformerEncoderLayer(d_model=model_dim, nhead=nhead, batch_first=True, dropout=dropout, norm_first=config['norm_first'])
        self.transformer_encoder = TransformerEncoder(encoder_layers, num_encoder_layers).to(device)#, encoder_norm).to(device)
    
        self.target_lm_head = nn.Linear(self.model_dimension, self.output_vocab_size, device=device)
        self.log_softmax = torch.nn.LogSoftmax(dim=2)   
        self.init_weights()
        
    def init_weights(self):
        initrange = 0.1
        self.encoder_input_embeddings.weight.data.uniform_(-initrange, initrange)
        self.target_lm_head.weight.data.uniform_(-initrange, initrange)
        self.target_lm_head.bias.data.uniform_(math.log(1 / self.output_vocab_size) - initrange, math.log(1 / self.output_vocab_size) + initrange)
        
    def forward(self, source, target, tokenizer, 
                src_key_padding_mask=None, tgt_key_padding_mask=None, tgt_mask=None, src_mask=None):
        source_embeddings_raw = self.encoder_input_embeddings(source) * math.sqrt(self.model_dimension)
        source_embeddings = self.pos_encoder(source_embeddings_raw)        
        memory = self.transformer_encoder(
            source_embeddings, 
            src_key_padding_mask=src_key_padding_mask,
            mask=src_mask
        ) 
        
        logits = self.target_lm_head(memory)
        res = self.log_softmax(logits)   
        return res

It looks like training loss reaches some low level (≈ 7.5 for NLL) and starts to diverge rapidly thereafter. It also appears as if there is some interaction between the bias of the layer norm and the bias of the attention operation.
Here is the loss: link
Image:

Screen Shot 2022-01-30 at 1.46.28 PM1232×648 30.5 KB

Here are some charts plotting the parameter values evolving over training:

Group 1 (45)1851×2751 519 KB

And here are some annotations of the chart that indicate an interaction between norm bias and attention bias (though not positive!)

Group 2 (47)1851×2751 748 KB

I’m also using 1/5 the learning rate of the “attention is all you need” paper, so a bit surprised about the training divergence.
Finally, right around when the divergence starts happening, the attention weights of the first training layer get nuked, and it seems like it just starts paying attention to the same tokens for every other token (note the rows sum to 1 in this heatmap):

Screen Shot 2022-01-30 at 1.58.52 PM1626×1606 228 KB


Screen Shot 2022-01-30 at 1.59.08 PM1626×1666 174 KB


Screen Shot 2022-01-30 at 1.59.21 PM1670×1614 183 KB

Any idea on what might be happening? This model doesn’t perform much better than an input-independent baseline (which reaches 8.5 NLL whereas this model hits around 7 before diverging).
Also, here’s the hyper parameters used for training in case that helps.

Screen Shot 2022-01-30 at 2.02.49 PM1514×560 76.7 KB