Is an embedding layer doing any computation?
I got stuck on the constructor before reaching the actual model. Was I using an existing embedding layer, or implementing one?
The intuition
An embedding is a trainable lookup table. With V tokens and width D, its weight has shape (V, D). Indexing it with token IDs of shape (B, T) produces (B, T, D). It is also equivalent to multiplying one-hot vectors by that table, without constructing the one-hot vectors.
E[token_ids] # (B, T) -> (B, T, D)The useful distinction is ownership: a custom nn.Module registers its own Parameter; nn.Embedding already does that work. A constructor error says nothing about the quality of the learned representation.
Related reading · PyTorch: Embedding