Introduction

An RNN is usually trained to predict logits over a vocabulary, then a softmax converts those logits into next-token probabilities. But you can train the model to output an embedding vector instead, including a Word2Vec-style target. That design can work well when you care more about semantic similarity than exact token classification, but it changes both the loss function and the way you decode predictions.

What Changes When You Predict Embeddings

In the usual language-model setup, the final layer has one unit per vocabulary item. The target is a token id, and the loss is commonly sparse categorical cross-entropy.

If you predict embeddings, the final layer instead has embedding_dim outputs. The target is a dense vector such as the precomputed embedding for the expected next word. Now the problem looks like regression or metric learning rather than plain classification.

That can be attractive when the vocabulary is large. A dense vector head can be much smaller than a huge softmax head, and near-miss predictions may still land close to the correct word in embedding space.

Building the Model

The basic architecture is straightforward: token ids go into an embedding layer, the sequence flows through an RNN variant such as LSTM or GRU, and a dense layer projects the hidden state into the target embedding space.

1import tensorflow as tf
2
3vocab_size = 1000
4input_embedding_dim = 64
5target_embedding_dim = 100
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Embedding(vocab_size, input_embedding_dim),
9    tf.keras.layers.GRU(128),
10    tf.keras.layers.Dense(target_embedding_dim)
11])
12
13model.compile(
14    optimizer="adam",
15    loss=tf.keras.losses.MeanSquaredError()
16)

This model does not output a word index. It outputs a vector of length 100.

Preparing Target Embeddings

The crucial part is building the training targets. If your next-token label is an integer id, you need to replace that id with the corresponding pretrained embedding vector.

1import numpy as np
2import tensorflow as tf
3
4embedding_matrix = np.random.randn(vocab_size, target_embedding_dim).astype("float32")
5
6x_train = np.array([
7    [1, 5, 8],
8    [2, 7, 9],
9    [4, 3, 6],
10], dtype="int32")
11
12y_token_ids = np.array([10, 11, 12], dtype="int32")
13y_train = embedding_matrix[y_token_ids]
14
15model.fit(x_train, y_train, epochs=3, verbose=0)

That is the main conceptual shift. The labels are no longer token ids. They are embedding vectors.

Choosing a Loss Function

Mean squared error is the simplest starting point, but it is not the only choice. Cosine-based losses can be useful when vector direction matters more than raw magnitude.

1model.compile(
2    optimizer="adam",
3    loss=tf.keras.losses.CosineSimilarity(axis=-1)
4)

In Keras, cosine similarity loss is minimized, so more similar vectors lead to lower loss values. In practice, many teams normalize both target and predicted embeddings before comparing them, especially if retrieval later uses cosine similarity.

Decoding Back to Words

Once the model predicts an embedding, you still need a way to turn that vector into a token. The common strategy is nearest-neighbor lookup against the embedding table.

1pred = model.predict(x_train[:1], verbose=0)[0]
2normalized_table = embedding_matrix / np.linalg.norm(embedding_matrix, axis=1, keepdims=True)
3normalized_pred = pred / np.linalg.norm(pred)
4nearest_id = np.argmax(normalized_table @ normalized_pred)
5print(nearest_id)

This is the tradeoff. You may reduce the output dimension during training, but inference now needs a similarity search step if the final result should be a word.

When This Approach Helps

Predicting embeddings is useful when semantic closeness matters. If the correct next word is car and the model predicts a vector closest to vehicle, that may be acceptable for some ranking, retrieval, or representation-learning tasks.

It is less natural when you need calibrated probabilities over exact tokens, beam search over a language model, or direct use of cross-entropy metrics. In those settings, logits are usually the cleaner tool.

Common Pitfalls

The biggest mistake is forgetting that an embedding target changes the problem definition. If you ultimately need token probabilities, a vector-regression head may make downstream decoding harder.

Another mistake is assuming any Word2Vec space will work well as a target. If the embedding space does not align with the prediction task, the model may learn vectors that look numerically close but decode poorly.

A third issue is skipping normalization during similarity search. Magnitude differences can distort nearest-neighbor results when the retrieval metric and training loss are not aligned.

Summary

• Yes, an RNN can be trained to output Word2Vec-style embeddings instead of logits
• The final layer size becomes the embedding dimension rather than the vocabulary size
• Target labels must be embedding vectors, not token ids
• 'MeanSquaredError and cosine-style losses are common starting points'
• If you need words back, you usually decode with nearest-neighbor search in the embedding space