reverse word embeddings
keras
python
natural language processing
machine learning

reverse word embeddings in keras - python

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Reversing word embeddings usually means mapping vectors back to likely tokens, which is not an exact inverse operation. Embedding spaces are many-to-one in practice, and nearby vectors can correspond to semantically related words rather than one perfect answer. In Keras workflows, the practical solution is nearest-neighbor lookup against the embedding matrix.

Core Sections

Understand Why Exact Reversal Is Hard

An embedding layer maps integer token IDs to dense vectors. That mapping is direct from ID to vector, but the reverse direction from arbitrary vector to token ID needs a similarity search. If the vector came from a transformed hidden state, it may not align exactly with one embedding row.

python
1import numpy as np
2
3# Example embedding table, vocab_size x embedding_dim
4embedding_matrix = np.array([
5    [0.1, 0.2, 0.3],  # token 0
6    [0.9, 0.1, 0.2],  # token 1
7    [0.0, 0.8, 0.2],  # token 2
8], dtype=np.float32)
9
10query = np.array([0.88, 0.12, 0.18], dtype=np.float32)

The task is finding which row is closest to query.

Use Cosine Similarity for Token Recovery

Cosine similarity is commonly used because embeddings are often compared by direction rather than magnitude.

python
1import numpy as np
2
3def top_k_by_cosine(query_vec, matrix, k=3):
4    q = query_vec / (np.linalg.norm(query_vec) + 1e-8)
5    m = matrix / (np.linalg.norm(matrix, axis=1, keepdims=True) + 1e-8)
6    scores = m @ q
7    top_idx = np.argsort(scores)[::-1][:k]
8    return top_idx, scores[top_idx]
9
10idx, scores = top_k_by_cosine(query, embedding_matrix, k=2)
11print(idx, scores)

In real models, map idx back to tokens using your tokenizer index-to-word table.

Keras Example with Real Embedding Layer

When using tf.keras.layers.Embedding, extract weights after training and run nearest-neighbor lookup.

python
1import tensorflow as tf
2import numpy as np
3
4vocab_size = 1000
5embedding_dim = 16
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim),
9    tf.keras.layers.GlobalAveragePooling1D(),
10    tf.keras.layers.Dense(1, activation="sigmoid"),
11])
12
13# After training, fetch embedding table
14embedding_table = model.layers[0].get_weights()[0]  # shape: [vocab_size, embedding_dim]
15
16query_vec = embedding_table[42] + np.random.normal(0, 0.01, embedding_dim)
17indices, sims = top_k_by_cosine(query_vec, embedding_table, k=5)
18print("closest token ids:", indices)

This is the standard pattern for approximate reverse lookup.

Improve Quality with Approximate Nearest Neighbors

For large vocabularies, brute-force similarity can be slow. Libraries for approximate nearest neighbors reduce latency while keeping high recall. This is useful in retrieval tasks, embedding diagnostics, and interactive tooling.

Even with fast indexes, you should keep a validation routine that compares approximate results with exact top-k results on random samples.

Handle Out-of-Vocabulary and Subword Tokenization

If your tokenizer uses subwords, nearest token IDs may correspond to pieces instead of complete words. That is expected. Build post-processing that reconstructs human-readable text where needed.

Also track normalization strategy. If training used normalized embeddings but lookup uses raw vectors, ranking quality can drop.

Evaluate Reverse Lookup Quality

If reverse lookup is part of a production pipeline, measure quality rather than trusting visual spot checks. Build a benchmark set where you know expected nearest tokens and compute top-k recall. For contextual vectors, compare retrieval quality across different model layers to see where embeddings remain token-aligned.

python
1import numpy as np
2
3def recall_at_k(true_ids, pred_topk):
4    hits = 0
5    for t, preds in zip(true_ids, pred_topk):
6        if t in preds:
7            hits += 1
8    return hits / len(true_ids)
9
10true_ids = [1, 2, 0]
11pred_topk = [[1, 3, 5], [4, 2, 9], [7, 0, 6]]
12print("recall@3:", recall_at_k(true_ids, pred_topk))

Quality metrics make retrieval changes measurable when you update tokenizers or retrain embeddings.

Common Pitfalls

  • Assuming embedding reversal is exact rather than nearest-neighbor approximation.
  • Using Euclidean distance blindly when cosine similarity better matches model semantics.
  • Forgetting to normalize vectors consistently before similarity search.
  • Ignoring tokenizer details and misinterpreting subword token IDs as full words.
  • Running brute-force search on very large vocabularies without latency planning.

Summary

  • Reverse embedding lookup is a similarity search problem, not true inversion.
  • Extract embedding weights and retrieve top-k nearest rows.
  • Use tokenizer mappings to convert IDs back to readable tokens.
  • Normalize vectors consistently to maintain quality.
  • Consider approximate nearest-neighbor indexing for large vocabularies.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.