Token Embedding Lookup

easyFREEEmbeddings & Positional EncodingLLM 1: Foundations

Implement a token embedding lookup table.

Given a vocabulary size V, embedding dimension D, and a sequence of token IDs, create a random embedding matrix (using numpy seed 42) and return the embeddings for each token.

Function: token_embedding_lookup(V, D, token_ids)

V: vocab size (int)
D: embedding dimension (int)
token_ids: list of token IDs

Output: A numpy array of shape (len(token_ids), D) with the embedding row for each input token, in input order. The tests print each row with values rounded to 4 decimal places.

Example

Input

5 3
0 2 4

Output

[0.4967 -0.1383 0.6477]
[1.5792 0.7674 -0.4695]
[0.2420 -1.9133 -1.7249]

Reasoning

Inputs: $V = 5$ vocabulary entries, each embedded in $D = 3$ dimensions, and token_ids holds the 3 IDs $0, 2, 4$ , each in $[0, V-1]$ .
The embedding table $E$ is a $5 \times 3$ matrix of standard-normal numbers generated with seed 42, identical on every call. Row $i$ is the vector for token $i$ ; for example $E_1 = [1.5230, -0.2342, -0.2341]$ and $E_3 = [0.5426, -0.4634, -0.4657]$ .
Each ID $t_k$ is replaced by its row $E_{t_k}$ : token 0 gives $E_0 = [0.4967, -0.1383, 0.6477]$ , token 2 gives $E_2 = [1.5792, 0.7674, -0.4695]$ , token 4 gives $E_4 = [0.2420, -1.9133, -1.7249]$ .
Output: these rows stacked in input order, a $3 \times 3$ matrix: one row per token, $D$ columns, shown to 4 decimal places.
Notice that rows 1 and 3 never appear: the row count follows the sequence length, not $V$ ; a repeated ID would yield the same row twice. The random draws fill the table row by row, so with $D = 2$ instead, $E_1$ would hold different numbers.

Constraints

Use numpy with seed 42: np.random.seed(42); embeddings = np.random.randn(V, D)
1 <= V <= 100, 1 <= D <= 10
Token IDs are in range [0, V-1]
Round output to 4 decimal places

Learn this in

Text RepresentationNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium

Concept primer

Background Knowledge

The problem revolves around token embedding lookup, a fundamental concept in natural language processing (NLP) and machine learning. In NLP, tokens refer to the basic units of text, such as words or characters. Embeddings are a way to represent these tokens as dense vectors in a high-dimensional space, capturing their semantic meaning. The embedding dimension (D) determines the size of this vector space, while the vocabulary size (V) is the number of unique tokens in the language model.

Token embeddings are typically learned during the training process of a language model, but in this problem, we are tasked with creating a random embedding matrix. This matrix will have a size of $V \times D$ , where each row represents the embedding vector for a particular token. The embedding lookup process involves retrieving the corresponding embedding vector for a given token ID from this matrix. This is a crucial step in many NLP tasks, as it allows the model to understand the context and relationships between different tokens.

The use of numpy with a fixed seed (42) ensures that the random embedding matrix generated will be reproducible, which is important for debugging and testing purposes. The output will be the embedding matrix for the input tokens, with each row representing the embedding vector for a token, and the values rounded to 4 decimal places.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Create a random embedding matrix with the specified size ( $V \times D$ ) using numpy.
Parse the input token IDs and retrieve the corresponding embedding vectors from the matrix.
Round the embedding values to 4 decimal places and output the resulting matrix.

This problem can be seen as a simple application of array indexing and matrix operations, which are fundamental concepts in linear algebra and machine learning.

Step-by-Step Strategy

To implement the solution, follow these steps:

Import the necessary libraries, including numpy, and set the seed for reproducibility.
Read the input values for vocabulary size (V) and embedding dimension (D).
Create a random embedding matrix with size $V \times D$ using numpy.
Read the input token IDs and parse them into a list or array.
Use array indexing to retrieve the corresponding embedding vectors for each token ID from the matrix.
Round the embedding values to 4 decimal places.
Output the resulting embedding matrix.

Common Pitfalls

Some common pitfalls to watch out for include:

Forgetting to set the numpy seed, resulting in non-reproducible results.
Incorrectly parsing the input token IDs or vocabulary size.
Using the wrong data type or precision for the embedding matrix.
Failing to round the embedding values to the correct number of decimal places.

Time & Space Complexity

The time complexity of this problem is expected to be $O(V \times D + n)$ , where $n$ is the number of input token IDs. This is because we need to create the random embedding matrix (which takes $O(V \times D)$ time) and then retrieve the corresponding embedding vectors for each token ID (which takes $O(n)$ time). The space complexity is $O(V \times D)$ , as we need to store the entire embedding matrix in memory.

solution.py

Python 3.13

Test results