nn.TransformerEncoder-based language model always generates the most common token

username: sampdoria
cooked: <p>I’m trying to do language modeling on a custom dataset using <code>nn.TransformerEncoder</code>. I’m using <a href="https://github.com/pytorch/examples/tree/master/word_language_model" rel="noopener nofollow ugc">https://github.com/pytorch/examples/tree/master/word_language_model</a> as a reference.</p>
<p>Previously, I used Google’s Trax and its <code>TransformerLM</code> model to train a transformer with this dataset, based on this example: <a href="https://github.com/jalammar/jalammar.github.io/blob/master/notebooks/Trax_TransformerLM_Intro.ipynb" rel="noopener nofollow ugc">https://github.com/jalammar/jalammar.github.io/blob/master/notebooks/Trax_TransformerLM_Intro.ipynb</a>. There, I reduced the Adam learning rate to <code>1e-04</code>, replaced the data with my actual dataset, adjusted the hyperparameters for my use-case, and managed to get very good results.</p>
<p>I’m now trying to replicate the same result in PyTorch, but without much luck.</p>
<p>Here are the things I modified from the <code>word_language_model</code> example above:</p>
<ul>
<li>Changed the batching and dataset load logic. Since my dataset consists of separate sentences, my input data is of shape <code>(num_examples, seq_len)</code>, and each element is a token index, with 0 reserved for padding. I generate inputs of shape <code>(max_seq_len, batch_size)</code>.</li>
<li>Calculated a padding mask that contains <code>True</code> wherever the original data has 0 (the padding token), passing it as <code>src_key_padding_mask</code> to <code>nn.TransformerEncoder</code>.</li>
<li>Added an Adam optimizer instead of the previously inline LR annealing on plateau.</li>
<li>Removed <code>F.log_softmax</code>, so now returning raw data from the last <code>nn.Linear</code> layer, and changed the loss to <code>nn.CrossEntropyLoss(ignore_index=0)</code>.</li>
<li>Added accuracy measurements by masking the final logits using the key padding mask, taking an argmax, and comparing to the actual output.</li>
</ul>
<p>My training plateaus within a few dozen batches to an accuracy of ~10% and a constant loss, and simply predicts the most common overall token, no matter the input (this gives an accuracy of ~10%, as that token is around 10% of the training data). I’ve tried using a learning rate scheduler, adding some warmup steps, played around with hyperparameters and the initial learning rate, tried changing the initialization of the embedding and final <code>nn.Linear</code> weights to <code>nn.init.xavier_uniform_</code>, but nothing helps.</p>
<p>For comparison, I tried randomly shuffling my data around to make it nonsense, and the model arrives at similar numbers. So I’m pretty sure it’s learning nothing.</p>
<p>Here’s an example of how my training data looks (right before it goes into the model):</p>
<pre><code class="lang-auto">data =
       [[  1,   1,   1,   1],
        [  2,   2,   2,   2],
        [  3, 107,  81,  81],
        [115,   4, 111,  46],
        [  5,   5,   5,   5],
        [ 80,  41,  63,  61],
        # some zeroes here, starting at different rows,
        # depending on the size of the example
        ...]]

target =
    [  2,   3, 115,   5,  80,  42,  ...
       0,   0,   0,   0,   0,  0,   ...
       2, 107,   4,   5,  41,  96,  ...
       0,   0,   0,   0,   0,  0,   ...]]
</code></pre>
<p>(Almost every example starts with <code>1 2</code>, so I imagined it’d be easy to fit on at least this, but not even that happens.)</p>
<p>Here’s the model:</p>
<pre><code class="lang-python"># PositionalEncoding definition from word_language_model omitted
# for brevity

class Transformer(nn.Module):
    def __init__(self, n_tokens, d_model, n_heads, d_ff, n_layers, dropout=0.1,
            max_len=4096, activation='relu'):
        super(Transformer, self).__init__()

        self.mask = None
        self.d_model = d_model

        self.pos_encoder = PositionalEncoding(d_model, dropout=dropout)
        self.embedding = nn.Embedding(n_tokens, d_model, padding_idx=0)
        enc_layer = nn.TransformerEncoderLayer(d_model, n_heads, d_ff, dropout)
        self.tf_encoder = nn.TransformerEncoder(enc_layer, n_layers)
        self.decoder = nn.Linear(d_model, n_tokens)

        self.init_weights()

    def _from_binary_mask(self, mask):
        return mask.float() \
            .masked_fill(mask == False, float('-inf')) \
            .masked_fill(mask == True, float(0.0))

    def _generate_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        return self._from_binary_mask(mask)

    def init_weights(self):
        initrange = 0.1
        nn.init.uniform_(self.embedding.weight, -initrange, initrange)
        nn.init.zeros_(self.decoder.weight)
        nn.init.uniform_(self.decoder.weight, -initrange, initrange)
        # Alternative I tried:
        # nn.init.xavier_uniform_(self.embedding.weight)
        # nn.init.xavier_uniform_(self.decoder.weight)
        # nn.init.normal_(self.decoder.bias, 1e-6)

    def forward(self, src, use_mask=True, src_key_padding_mask=None):
        if use_mask:
            device = src.device
            if self.mask is None or self.mask.size(0) != len(src):
                mask = self._generate_mask(len(src)).to(device)
                self.mask = mask
        else:
            self.mask = None

        # embed and add positional information
        src = self.embedding(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)

        output = self.tf_encoder(src, self.mask,
                src_key_padding_mask=src_key_padding_mask)
        output = self.decoder(output)
        return output
</code></pre>
<p>And here’s my training loop:</p>
<pre><code class="lang-python">model = model.Transformer(n_tokens, args.d_model, args.n_heads, args.d_ff,
        args.n_layers, args.dropout).to(device)

optimizer = torch.optim.Adam(model.parameters(), weight_decay=1e-3)

# init data, etc.

def train(train_data):
    model.train()
    bs = args.batch_size

    total_loss = 0.0
    total_correct = 0
    total_total = 0

    start_time = time.time()
    for batch_i in range(0, train_data.size(0) // bs):
        transposed_data, target = get_batch(train_data, batch_i) # (bs, seq_len), (bs*seq_len)
        data = transposed_data.transpose(0, 1).type(torch.LongTensor).to(device) # (seq_len, bs)
        target = target.type(torch.LongTensor).to(device) # (bs*seq_len)

        optimizer.zero_grad()
        key_padding_mask = calc_key_padding_mask(transposed_data, target) # (bs, seq_len), 1 if not padding
        output = model(data, src_key_padding_mask=torch.logical_not(key_padding_mask))
        output = output.view(-1, n_tokens) # (bs*seq_len, n_tokens)

        logits = F.log_softmax(output, dim=-1) # (bs*seq_len, n_tokens)
        masked_logits = apply_mask_to_logits(logits, key_padding_mask) # (bs*seq_len, n_tokens), sets element to -inf if padding
        confidences, predictions = torch.max(masked_logits.exp(), 1) # (bs*seq_len), confidence will be 0 if padding
        total_correct += (torch.logical_and(predictions == target, confidences &gt; 0.0)).float().sum()
        total_total += key_padding_mask.float().sum()

        e = loss(output, target)
        e.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.25)
        optimizer.step()

        total_loss += e.item()

        if batch_i % args.log_interval == 0 and batch_i &gt; 0:
            curr_loss = total_loss / args.log_interval / bs
            curr_accuracy = total_correct / total_total
            total_loss = 0.
            total_correct = 0
            total_total = 0

            # print the stats...
</code></pre>
<p>Please help - I’m at a loss (no pun intended). I feel like I’ve tried everything.</p>