Token embedding
×learnedtf.gather(wte, X)
Look up a learned feature vector for each token ID.
- in
- [B batch, T positions]
- out
- [B batch, T positions, 768 d]
- parameters
- 38,597,376
Parameters
38,597,376 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Token embedding | 38,597,376 | |
| Per-instance total | 38,597,376 | |
The vocabulary projection reuses this table’s transpose; its parameters are counted here.
Derived from the configured table dimensions; checkpoint tensors have not been independently counted.
Codemodel.py 152–157 · modeling_gpt2.py 492–577
wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd], wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd], h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length)) self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) inputs_embeds = self.wte(input_ids) position_embeds = self.wpe(position_ids)
hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)Sources2
Direct evidence
The original TensorFlow graph creates both tables and adds their gathered rows.
openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 152–157
Read code excerpt
wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd],… lines 153–153 omitted …
wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd],… lines 155–156 omitted …
h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))Direct evidence
Transformers names the tables wte and wpe, matching the released checkpoint keys.
transformers · huggingface/transformers / src/transformers/models/gpt2/modeling_gpt2.py
Pinned revision 71c6f699ac9b ↗ · lines 492–577
Read code excerpt
self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)… lines 494–568 omitted …
inputs_embeds = self.wte(input_ids)… lines 570–575 omitted …
position_embeds = self.wpe(position_ids)
hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)