Here's a basic GPT2 implementation:

class GPT(nn.Module):
    def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers):
        self.seq_len = seq_len
        self.wte = nn.Embedding(vocab_size, model_dim)
        self.wpe = nn.Embedding(seq_len, model_dim)
        self.h = nn.Sequential(
            *[DecodeLayer(model_dim, n_heads) for _ in range(n_layers)]
        self.ln_f = nn.LayerNorm(model_dim)

    def forward(self, x):
        embeds = self.wte(x) + self.wpe(torch.arange(0, self.seq_len).to(x.device)) # [seq_len, model_dim]
        embeds = self.h(embeds)
        embeds = self.ln_f(embeds)
        # (seq_len, model_dim)  x (model_dim, vocab_size) => seq_len vocab_size
        return embeds @ self.wte.weight.T

The positional embeddings are added to the input with:

self.wpe(torch.arange(0, self.seq_len).to(x.device))

I wonder, if it would be possible to just do: self.wpe = nn.Parameter(seq_len, model_dim), and change the line in the forward pass to:

embeds = self.wte(x) + self.wpe
  • $\begingroup$ Is this an implementation question or a theory question? i.e., are you asking whether the two implementations do the same thing or are you asking whether a slightly different way of doing PE would work effectively. $\endgroup$ Commented Jan 25 at 23:32
  • $\begingroup$ Good question. There's 2 questions hidden in here: Implementation question -- the 2 things should do the same thing as far as I can tell. Since, torch.arange(0, self.seq_len) will select the entire embedding matrix. However, from some testing, they do not do the same thing. Theory question -- why not just add a seq_len, model_dim size matrix to the output of self.wte(x), will it not have the same effect as using an embedding layer? $\endgroup$
    – Foobar
    Commented Jan 26 at 5:11

1 Answer 1


Depends... for ViT yes, they are called learned positional encodings, and the reason is that you always have a fixed number of patches (at least theoretically)

However, for NLP, you cannot do that, because at inference, you might have strings longer or shorter than the reference, and if you have trained only seq_len positional encoding, and the inference string is longer than that, than you won't be able to apply a positional encoding to the final part of it


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .